diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 63d27c0b2..14c0b1268 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -14,7 +14,6 @@ add_library( postgres_connection.cpp postgres_copy_from.cpp postgres_copy_to.cpp - postgres_execute.cpp postgres_extension.cpp postgres_filter_pushdown.cpp postgres_hstore.cpp diff --git a/src/postgres_execute.cpp b/src/postgres_execute.cpp deleted file mode 100644 index 68bb19f51..000000000 --- a/src/postgres_execute.cpp +++ /dev/null @@ -1,78 +0,0 @@ -#include "duckdb.hpp" -#include "duckdb/main/client_context.hpp" - -#include "duckdb/parser/parsed_data/create_table_function_info.hpp" -#include "postgres_scanner.hpp" -#include "duckdb/main/database_manager.hpp" -#include "duckdb/main/attached_database.hpp" -#include "storage/postgres_catalog.hpp" -#include "storage/postgres_transaction.hpp" - -namespace duckdb { - -struct PGExecuteBindData : public TableFunctionData { - explicit PGExecuteBindData(PostgresCatalog &pg_catalog, string query_p, bool use_transaction) - : pg_catalog(pg_catalog), query(std::move(query_p)), use_transaction(use_transaction) { - } - - bool finished = false; - PostgresCatalog &pg_catalog; - string query; - bool use_transaction = true; -}; - -static duckdb::unique_ptr PGExecuteBind(ClientContext &context, TableFunctionBindInput &input, - vector &return_types, vector &names) { - return_types.emplace_back(LogicalType::BIGINT); - names.emplace_back("rowcount"); - - // look up the database to query - auto db_name = input.inputs[0].GetValue(); - auto &db_manager = DatabaseManager::Get(context); - auto db = db_manager.GetDatabase(context, Identifier(db_name)); - if (!db) { - throw BinderException("Failed to find attached database \"%s\" referenced in postgres_query", db_name); - } - auto &catalog = db->GetCatalog(); - if (catalog.GetCatalogType() != "postgres") { - throw BinderException("Attached database \"%s\" does not refer to a Postgres database", db_name); - } - auto &pg_catalog = catalog.Cast(); - - bool use_transaction = true; - for (auto &kv : input.named_parameters) { - if (kv.first == "use_transaction") { - use_transaction = BooleanValue::Get(kv.second); - } - } - - return make_uniq(pg_catalog, input.inputs[1].GetValue(), use_transaction); -} - -static void PGExecuteFunction(ClientContext &context, TableFunctionInput &data_p, DataChunk &output) { - auto &data = data_p.bind_data->CastNoConst(); - if (data.finished) { - output.SetChildCardinality(0); - return; - } - auto &transaction = Transaction::Get(context, data.pg_catalog).Cast(); - unique_ptr res; - if (data.use_transaction) { - res = transaction.Query(data.query); - } else { - res = transaction.QueryWithoutTransaction(data.query); - } - - int64_t *vec_data = FlatVector::GetDataMutable(output.data[0]); - vec_data[0] = res->AffectedRows(); - output.SetChildCardinality(1); - data.finished = true; -} - -PostgresExecuteFunction::PostgresExecuteFunction() - : TableFunction("postgres_execute", {LogicalType::VARCHAR, LogicalType::VARCHAR}, PGExecuteFunction, - PGExecuteBind) { - named_parameters["use_transaction"] = LogicalType::BOOLEAN; -} - -} // namespace duckdb diff --git a/src/postgres_query.cpp b/src/postgres_query.cpp index a7036238e..ee7f210e3 100644 --- a/src/postgres_query.cpp +++ b/src/postgres_query.cpp @@ -148,4 +148,16 @@ PostgresQueryFunction::PostgresQueryFunction() projection_pushdown = true; global_initialization = TableFunctionInitialization::INITIALIZE_ON_SCHEDULE; } + +PostgresExecuteFunction::PostgresExecuteFunction() + : TableFunction("postgres_execute", {LogicalType::VARCHAR, LogicalType::VARCHAR}, nullptr, PGQueryBind) { + named_parameters["use_transaction"] = LogicalType::BOOLEAN; + named_parameters["params"] = LogicalType::ANY; + PostgresScanFunction scan_function; + init_global = scan_function.init_global; + init_local = scan_function.init_local; + function = scan_function.function; + projection_pushdown = true; + global_initialization = TableFunctionInitialization::INITIALIZE_ON_SCHEDULE; +} } // namespace duckdb diff --git a/src/postgres_scanner.cpp b/src/postgres_scanner.cpp index 12caeee32..7414cccbe 100644 --- a/src/postgres_scanner.cpp +++ b/src/postgres_scanner.cpp @@ -534,7 +534,10 @@ static void PostgresScan(ClientContext &context, TableFunctionInput &data, DataC auto &gstate = data.global_state->Cast(); if (gstate.collection) { - gstate.collection->Scan(gstate.scan_state, output); + bool scan_happened = gstate.collection->Scan(gstate.scan_state, output); + if (!scan_happened) { + output.SetChildCardinality(0); + } return; } auto &local_state = data.local_state->Cast(); diff --git a/test/sql/storage/attach_rollback_throws.test b/test/sql/storage/attach_rollback_throws.test index 2c820f0d4..d81d6ae38 100644 --- a/test/sql/storage/attach_rollback_throws.test +++ b/test/sql/storage/attach_rollback_throws.test @@ -23,7 +23,7 @@ s 1 statement error CALL postgres_execute('s', 'SELECT pg_terminate_backend(pg_backend_pid())') ---- -pg_terminate_backend +terminating connection due to administrator command # connection must have been discarded, as it cannot pass the health check on returning it to the pool query II diff --git a/test/sql/storage/postgres_execute_transaction.test b/test/sql/storage/postgres_execute_transaction.test index 973e72ae4..ed8b1c7e3 100644 --- a/test/sql/storage/postgres_execute_transaction.test +++ b/test/sql/storage/postgres_execute_transaction.test @@ -39,7 +39,7 @@ CALL postgres_query('s', 'SELECT 42') 42 statement ok -CALL postgres_execute('s', 'INSERT INTO postgres_execute_attempt VALUES (42); INSERT INTO postgres_execute_attempt VALUES (84)') +CALL postgres_execute('s', 'INSERT INTO postgres_execute_attempt VALUES (42), (84)') statement ok COMMIT diff --git a/test/sql/storage/postgres_execute_use_transaction.test b/test/sql/storage/postgres_execute_use_transaction.test index 5257dd3a6..2c083720b 100644 --- a/test/sql/storage/postgres_execute_use_transaction.test +++ b/test/sql/storage/postgres_execute_use_transaction.test @@ -12,20 +12,17 @@ ATTACH 'dbname=postgresscanner' AS s (TYPE POSTGRES) statement error CALL postgres_execute('s', 'VACUUM') ---- -Invalid Error: Failed to execute query "BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ; -VACUUM": ERROR: VACUUM cannot run inside a transaction block +Invalid Error: Failed to execute query "VACUUM": ERROR: VACUUM cannot run inside a transaction block statement error CALL postgres_execute('s', 'VACUUM', use_transaction=true) ---- -Invalid Error: Failed to execute query "BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ; -VACUUM": ERROR: VACUUM cannot run inside a transaction block +Invalid Error: Failed to execute query "VACUUM": ERROR: VACUUM cannot run inside a transaction block statement error CALL postgres_execute('s', 'VACUUM', use_transaction=true) ---- -Invalid Error: Failed to execute query "BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ; -VACUUM": ERROR: VACUUM cannot run inside a transaction block +Invalid Error: Failed to execute query "VACUUM": ERROR: VACUUM cannot run inside a transaction block statement ok CALL postgres_execute('s', 'VACUUM', use_transaction=false)