From 3715f35c93bcd3def357dd59f3d6bff77bc0ef44 Mon Sep 17 00:00:00 2001 From: Alex Kasko Date: Wed, 2 Sep 2026 22:47:59 +0100 Subject: [PATCH] Allow non-preparable queries in postgres_query After #552 the `postgres_query()` function is also used to run DDL and DML queries. Though the queries are still being prepared on a remote server before execution (this is necessary to get the resulting columns during binding) - the same "general" code path is used for both `SELECT`s and DML queries. For majority of the queries this does not cause problems, as Postgres (comparing to other DBs) has minimal limitations to queries that cannot be prepared. Though two notable groups that cannot be prepared contain utility queries (like `VACUUM`) and concatenated multi-query strings. This PR adds support for a `prepare=FALSE` named parameter to `postgres_query()` that makes such calls to use the "fast-path", when the input string is run in Postgres at once without preparing it. This non-prepared mode does not allow to specify query parameters and does not return a result set - intended to be used only with utility or DML queries. Testing: existing `postgres_query` test is updated with multi-query strings coverage. --- src/postgres_query.cpp | 64 +++++++++++++++++----------- test/sql/scanner/postgres_query.test | 50 ++++++++++++++++++++++ 2 files changed, 90 insertions(+), 24 deletions(-) diff --git a/src/postgres_query.cpp b/src/postgres_query.cpp index 7069be166..77a2cfac6 100644 --- a/src/postgres_query.cpp +++ b/src/postgres_query.cpp @@ -22,10 +22,37 @@ static bool ExtractFlag(TableFunctionBindInput &input, const string &name, bool return default_val; } -static unique_ptr PGQueryBind(ClientContext &context, TableFunctionBindInput &input, - vector &return_types, vector &names) { +static unique_ptr BindDML(ClientContext &context, TableFunctionBindInput &input, + vector &return_types, vector &names, + PostgresCatalog &pg_catalog, PostgresConnection &con, std::string sql, + bool use_transaction) { + // The statement returns no result columns: it's a command (DDL, or DML without RETURNING). + // Instead of failing, run it as a command and return a single-row Success result. We reuse + // the prepare/describe just done — no extra round-trip — and defer execution to + // InitGlobalState (execution time, not bind, so EXPLAIN does not run it). auto result = make_uniq(context); + result->command_only = true; + if (ExtractFlag(input, "suppress_dml_output", false)) { + // This invocation wraps a command with no result set (DDL, or DML without RETURNING). Tell the + // binder via the return-type modifier so that when this is routed through CONNECT the outer + // statement is reported as NOTHING and displays like a native command (no spurious result table). + input.table_function.call_return_type = StatementReturnType::NOTHING; + } + return_types.emplace_back(LogicalType::BIGINT); + names.emplace_back(Identifier("rowcount")); + result->SetCatalog(pg_catalog); + result->dsn = con.GetDSN(); + result->types = return_types; + result->names.emplace_back(names[0].GetIdentifierName()); + result->read_only = false; + result->sql = std::move(sql); + result->use_transaction = use_transaction; + PostgresScanFunction::PrepareBind(pg_catalog.GetPostgresVersion(), context, *result, 0); + return std::move(result); +} +static unique_ptr PGQueryBind(ClientContext &context, TableFunctionBindInput &input, + vector &return_types, vector &names) { if (input.inputs[0].IsNull() || input.inputs[1].IsNull()) { throw BinderException("Parameters to postgres_query cannot be NULL"); } @@ -68,6 +95,13 @@ static unique_ptr PGQueryBind(ClientContext &context, TableFunctio auto &con = use_transaction ? transaction.GetConnection() : transaction.GetConnectionWithoutTransaction(); + if (!ExtractFlag(input, "prepare", true)) { + if (param_values.size() > 0) { + throw BinderException("query parameters cannot be used with 'prepare=FALSE'"); + } + return BindDML(context, input, return_types, names, pg_catalog, con, std::move(sql), use_transaction); + } + auto conn = con.GetConn(); // prepare execution of the query to figure out the result types and names auto prepared = PQprepare(conn, "", sql.c_str(), 0, nullptr); @@ -87,29 +121,9 @@ static unique_ptr PGQueryBind(ClientContext &context, TableFunctio } int nfields = PQnfields(describe_prepared); if (nfields <= 0) { - // The statement returns no result columns: it's a command (DDL, or DML without RETURNING). - // Instead of failing, run it as a command and return a single-row Success result. We reuse - // the prepare/describe just done — no extra round-trip — and defer execution to - // InitGlobalState (execution time, not bind, so EXPLAIN does not run it). - result->command_only = true; - if (ExtractFlag(input, "suppress_dml_output", false)) { - // This invocation wraps a command with no result set (DDL, or DML without RETURNING). Tell the - // binder via the return-type modifier so that when this is routed through CONNECT the outer - // statement is reported as NOTHING and displays like a native command (no spurious result table). - input.table_function.call_return_type = StatementReturnType::NOTHING; - } - return_types.emplace_back(LogicalType::BIGINT); - names.emplace_back(Identifier("rowcount")); - result->SetCatalog(pg_catalog); - result->dsn = con.GetDSN(); - result->types = return_types; - result->names.emplace_back(names[0].GetIdentifierName()); - result->read_only = false; - result->sql = std::move(sql); - result->use_transaction = use_transaction; - PostgresScanFunction::PrepareBind(pg_catalog.GetPostgresVersion(), context, *result, 0); - return std::move(result); + return BindDML(context, input, return_types, names, pg_catalog, con, std::move(sql), use_transaction); } + auto result = make_uniq(context); auto type_config = PostgresTypeConfig::FromContext(context); for (idx_t c = 0; c < nfields; c++) { PostgresType postgres_type; @@ -154,6 +168,7 @@ PostgresQueryFunction::PostgresQueryFunction() named_parameters["use_transaction"] = LogicalType::BOOLEAN; named_parameters["params"] = LogicalType::ANY; named_parameters["suppress_dml_output"] = LogicalType::BOOLEAN; + named_parameters["prepare"] = LogicalType::BOOLEAN; PostgresScanFunction scan_function; init_global = scan_function.init_global; init_local = scan_function.init_local; @@ -166,6 +181,7 @@ PostgresExecuteFunction::PostgresExecuteFunction() : TableFunction("postgres_execute", {LogicalType::VARCHAR, LogicalType::VARCHAR}, nullptr, PGQueryBind) { named_parameters["use_transaction"] = LogicalType::BOOLEAN; named_parameters["params"] = LogicalType::ANY; + named_parameters["prepare"] = LogicalType::BOOLEAN; PostgresScanFunction scan_function; init_global = scan_function.init_global; init_local = scan_function.init_local; diff --git a/test/sql/scanner/postgres_query.test b/test/sql/scanner/postgres_query.test index 68fb8fa75..e2aca8041 100644 --- a/test/sql/scanner/postgres_query.test +++ b/test/sql/scanner/postgres_query.test @@ -202,5 +202,55 @@ CALL postgres_query('s1', 'UPDATE dml_output_test SET col1 = 43', suppress_dml_o statement ok CALL postgres_query('s1', 'DROP TABLE dml_output_test') +# non-preparable statement + +statement error +CALL postgres_query('s1', 'PREPARE p1 AS SELECT 42; EXECUTE p1;') +---- +Failed to prepare + +statement error +CALL postgres_query('s1', 'PREPARE p1 AS SELECT 42; EXECUTE p1;', prepare=FALSE, params=row(42)) +---- +query parameters cannot be used with 'prepare=FALSE' + +# use the same connection + +statement ok +BEGIN TRANSACTION + +statement ok +CALL postgres_query('s1', 'PREPARE p1 AS SELECT 42; EXECUTE p1;', prepare=FALSE) + +query I +FROM postgres_query('s1', 'EXECUTE p1', prepare=FALSE) +---- +1 + +query II +SELECT column_name, column_type FROM ( + DESCRIBE FROM postgres_query('s1', 'EXECUTE p1', prepare=FALSE) +) +---- +rowcount BIGINT + +query II +SELECT column_name, column_type FROM ( + DESCRIBE FROM postgres_execute('s1', 'EXECUTE p1', prepare=FALSE) +) +---- +rowcount BIGINT + +statement ok +CALL postgres_query('s1', 'DEALLOCATE p1', prepare=FALSE) + +statement error +FROM postgres_query('s1', 'EXECUTE p1', prepare=FALSE) +---- +does not exist + +statement ok +ROLLBACK + statement ok DETACH s1