diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 0bc407844..f8138d09d 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -17,6 +17,7 @@ add_library( postgres_filter_pushdown.cpp postgres_hstore.cpp postgres_parameters.cpp + postgres_pin_connection.cpp postgres_query.cpp postgres_scanner.cpp postgres_secrets.cpp diff --git a/src/include/postgres_scanner.hpp b/src/include/postgres_scanner.hpp index c992c319d..46a4b2803 100644 --- a/src/include/postgres_scanner.hpp +++ b/src/include/postgres_scanner.hpp @@ -65,6 +65,8 @@ struct PostgresBindData : public FunctionData { string limit_clause; + uint64_t pinned_connection_id = 0; + public: void SetTablePages(idx_t approx_num_pages); @@ -102,4 +104,14 @@ class PostgresExecuteFunction : public TableFunction { PostgresExecuteFunction(); }; +class PostgresPinConnectionFunction : public ScalarFunction { +public: + PostgresPinConnectionFunction(); +}; + +class PostgresClosePinnedConnectionFunction : public ScalarFunction { +public: + PostgresClosePinnedConnectionFunction(); +}; + } // namespace duckdb diff --git a/src/include/postgres_utils.hpp b/src/include/postgres_utils.hpp index a3b14a8da..8fc9f0ad2 100644 --- a/src/include/postgres_utils.hpp +++ b/src/include/postgres_utils.hpp @@ -17,6 +17,7 @@ #include "duckdb/main/secret/secret_manager.hpp" namespace duckdb { +class PostgresCatalog; class PostgresSchemaEntry; class PostgresTransaction; @@ -105,6 +106,7 @@ class PostgresUtils { static PostgresQueryMode ExtractQueryMode(ClientContext &ctx); static PostgresQueryMode ExtractQueryMode(Value &value); + static PostgresCatalog &GetCatalogByName(ClientContext &context, const std::string &name); private: static mutex libpq_init_lock; diff --git a/src/postgres_extension.cpp b/src/postgres_extension.cpp index aac0a14fd..d911a4826 100644 --- a/src/postgres_extension.cpp +++ b/src/postgres_extension.cpp @@ -143,6 +143,12 @@ static void LoadInternal(ExtensionLoader &loader) { PostgresConfigurePoolFunction configure_pool_function; loader.RegisterFunction(configure_pool_function); + PostgresPinConnectionFunction pin_connection_function; + loader.RegisterFunction(pin_connection_function); + + PostgresClosePinnedConnectionFunction postgres_close_pinned_connection_function; + loader.RegisterFunction(postgres_close_pinned_connection_function); + RegisterHstoreFunctions(loader); // Register the new type diff --git a/src/postgres_pin_connection.cpp b/src/postgres_pin_connection.cpp new file mode 100644 index 000000000..527f3bf62 --- /dev/null +++ b/src/postgres_pin_connection.cpp @@ -0,0 +1,64 @@ +#include "postgres_scanner.hpp" + +#include "storage/postgres_catalog.hpp" + +namespace duckdb { + +static void PostgresPinConnection(DataChunk &args, ExpressionState &state, Vector &result) { + UnifiedVectorFormat catalog_name_data; + args.data[0].ToUnifiedFormat(catalog_name_data); + if (!catalog_name_data.validity.RowIsValid(0)) { + throw InvalidInputException("Specified attached database name must be not null"); + } + + auto catalog_name_strt = UnifiedVectorFormat::GetData(catalog_name_data)[0]; + string catalog_name(catalog_name_strt.GetData(), catalog_name_strt.GetSize()); + + auto &catalog = PostgresUtils::GetCatalogByName(state.GetContext(), catalog_name); + auto &pool = catalog.GetConnectionPool(); + + auto conn = pool.ForceAcquire(); + uint64_t conn_id = pool.PinConnection(std::move(conn)); + + auto result_data = FlatVector::GetDataMutable(result); + result_data[0] = conn_id; +} + +PostgresPinConnectionFunction::PostgresPinConnectionFunction() + : ScalarFunction("postgres_pin_connection", {LogicalType::VARCHAR}, LogicalType::UBIGINT, PostgresPinConnection) { + SetNullHandling(FunctionNullHandling::SPECIAL_HANDLING); + SetFallible(); + SetVolatile(); +} + +static void PostgresClosePinnedConnection(DataChunk &args, ExpressionState &state, Vector &result) { + UnifiedVectorFormat catalog_name_data; + args.data[0].ToUnifiedFormat(catalog_name_data); + if (!catalog_name_data.validity.RowIsValid(0)) { + throw InvalidInputException("Specified catalog name must be not null"); + } + auto catalog_name_strt = UnifiedVectorFormat::GetData(catalog_name_data)[0]; + string catalog_name(catalog_name_strt.GetData(), catalog_name_strt.GetSize()); + + UnifiedVectorFormat conn_id_data; + args.data[1].ToUnifiedFormat(conn_id_data); + if (!conn_id_data.validity.RowIsValid(0)) { + throw InvalidInputException("Specified connection ID must be not null"); + } + auto conn_id = UnifiedVectorFormat::GetData(conn_id_data)[0]; + + auto &catalog = PostgresUtils::GetCatalogByName(state.GetContext(), catalog_name); + { auto conn = catalog.GetConnectionPool().UnpinConnection(conn_id); } + + auto &result_validity = FlatVector::ValidityMutable(result); + result_validity.SetInvalid(0); +} + +PostgresClosePinnedConnectionFunction::PostgresClosePinnedConnectionFunction() + : ScalarFunction("postgres_close_pinned_connection", {LogicalType::VARCHAR, LogicalType::UBIGINT}, + LogicalType::BOOLEAN, PostgresClosePinnedConnection) { + SetNullHandling(FunctionNullHandling::SPECIAL_HANDLING); + SetFallible(); + SetVolatile(); +} +} // namespace duckdb diff --git a/src/postgres_query.cpp b/src/postgres_query.cpp index 735de2250..a1b531405 100644 --- a/src/postgres_query.cpp +++ b/src/postgres_query.cpp @@ -1,6 +1,8 @@ #include "duckdb.hpp" -#include "duckdb/main/client_context.hpp" +#include "dbconnector/defer.hpp" + +#include "duckdb/main/client_context.hpp" #include "duckdb/parser/parsed_data/create_table_function_info.hpp" #include "postgres_parameters.hpp" #include "postgres_scanner.hpp" @@ -24,14 +26,15 @@ static bool ExtractFlag(TableFunctionBindInput &input, const string &name, bool static unique_ptr BindDML(ClientContext &context, TableFunctionBindInput &input, vector &return_types, vector &names, - PostgresCatalog &pg_catalog, PostgresConnection &con, std::string sql, - bool use_transaction) { + PostgresCatalog &catalog, PostgresConnection &con, std::string sql, + bool use_transaction, uint64_t pinned_connection_id) { // 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, pg_catalog.GetName()); + auto result = make_uniq(context, catalog.GetName()); result->command_only = true; + result->pinned_connection_id = pinned_connection_id; 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 @@ -46,7 +49,7 @@ static unique_ptr BindDML(ClientContext &context, TableFunctionBin result->read_only = false; result->sql = std::move(sql); result->use_transaction = use_transaction; - PostgresScanFunction::PrepareBind(pg_catalog.GetPostgresVersion(), context, *result, 0, pg_catalog); + PostgresScanFunction::PrepareBind(catalog.GetPostgresVersion(), context, *result, 0, catalog); return std::move(result); } @@ -80,36 +83,16 @@ static PostgresQueryMode ResolveQueryMode(ClientContext &context, TableFunctionB return PostgresQueryMode::AUTO; } -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"); - } - - // 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(); - auto &transaction = Transaction::Get(context, catalog).Cast(); - auto sql = input.inputs[1].GetValue(); - // strip any trailing semicolons +static string StripTrailingSemicolon(string sql) { StringUtil::RTrim(sql); while (!sql.empty() && sql.back() == ';') { sql = sql.substr(0, sql.size() - 1); StringUtil::RTrim(sql); } + return sql; +} - bool use_transaction = ExtractFlag(input, "use_transaction", true); - - vector param_values; +static vector ExtractParams(TableFunctionBindInput &input) { auto params_it = input.named_parameters.find("params"); if (params_it != input.named_parameters.end()) { Value &struct_val = params_it->second; @@ -119,16 +102,73 @@ static unique_ptr PGQueryBind(ClientContext &context, TableFunctio if (struct_val.type().id() != LogicalTypeId::STRUCT && struct_val.type().id() != LogicalTypeId::TUPLE) { throw BinderException("Query parameters must be specified in a STRUCT"); } - param_values = StructValue::GetChildren(struct_val); + return StructValue::GetChildren(struct_val); + } + return vector(); +} + +static uint64_t ExtractPinnedConnId(TableFunctionBindInput &input) { + uint64_t pinned_connection_id = 0; + auto conn_it = input.named_parameters.find("connection"); + if (conn_it != input.named_parameters.end()) { + Value &conn_val = conn_it->second; + if (conn_val.IsNull()) { + throw BinderException("Specified connection must be not null"); + } + pinned_connection_id = UBigIntValue::Get(conn_val); } + return pinned_connection_id; +} + +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"); + } + + auto db_name = input.inputs[0].GetValue(); + auto &catalog = PostgresUtils::GetCatalogByName(context, db_name); + auto input_sql = input.inputs[1].GetValue(); + auto sql = StripTrailingSemicolon(std::move(input_sql)); + bool use_transaction = ExtractFlag(input, "use_transaction", true); + vector param_values = ExtractParams(input); + uint64_t pinned_connection_id = ExtractPinnedConnId(input); + + optional_ptr con_ptr = nullptr; + PostgresPoolConnection pinned_connection; + + if (pinned_connection_id > 0) { + pinned_connection = catalog.GetConnectionPool().UnpinConnection(pinned_connection_id); + PostgresConnection &conn = pinned_connection.GetConnection(); + con_ptr = &conn; + } else { + auto &transaction = Transaction::Get(context, catalog).Cast(); + if (use_transaction) { + con_ptr = &transaction.GetConnection(); + } else { + con_ptr = &transaction.GetConnectionWithoutTransaction(); + } + } + + auto deferred_pin = dbconnector::Defer([&pinned_connection, &catalog] { + if (!pinned_connection) { + return; + } + try { + catalog.GetConnectionPool().PinConnection(std::move(pinned_connection)); + } catch (...) { + // suppress + } + }); - auto &con = use_transaction ? transaction.GetConnection() : transaction.GetConnectionWithoutTransaction(); + PostgresConnection &con = *con_ptr; 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); + return BindDML(context, input, return_types, names, catalog, con, std::move(sql), use_transaction, + pinned_connection_id); } auto conn = con.GetConn(); @@ -150,9 +190,11 @@ static unique_ptr PGQueryBind(ClientContext &context, TableFunctio } int nfields = PQnfields(describe_prepared); if (nfields <= 0) { - return BindDML(context, input, return_types, names, pg_catalog, con, std::move(sql), use_transaction); + return BindDML(context, input, return_types, names, catalog, con, std::move(sql), use_transaction, + pinned_connection_id); } - auto result = make_uniq(context, pg_catalog.GetName()); + auto result = make_uniq(context, catalog.GetName()); + result->pinned_connection_id = pinned_connection_id; auto type_config = PostgresTypeConfig::FromContext(context); for (idx_t c = 0; c < nfields; c++) { PostgresType postgres_type; @@ -188,7 +230,7 @@ static unique_ptr PGQueryBind(ClientContext &context, TableFunctio result->query_mode = ResolveQueryMode(context, input, param_values); result->params = PostgresParameters(std::move(param_types), std::move(param_values)); result->use_transaction = use_transaction; - PostgresScanFunction::PrepareBind(pg_catalog.GetPostgresVersion(), context, *result, 0, pg_catalog); + PostgresScanFunction::PrepareBind(catalog.GetPostgresVersion(), context, *result, 0, catalog); return std::move(result); } @@ -203,6 +245,7 @@ static FunctionSignature PostgresQuerySignature(bool suppress_dml_output) { } options.Add("prepare", LogicalType::BOOLEAN); options.Add("query_mode", LogicalType::VARCHAR); + options.Add("connection", LogicalType::UBIGINT); }); return signature; } diff --git a/src/postgres_scanner.cpp b/src/postgres_scanner.cpp index 20f5b0159..3a98b832c 100644 --- a/src/postgres_scanner.cpp +++ b/src/postgres_scanner.cpp @@ -3,6 +3,8 @@ #include +#include "dbconnector/attached.hpp" + #include "duckdb/main/extension/extension_loader.hpp" #include "duckdb/common/shared_ptr.hpp" #include "duckdb/common/helper.hpp" @@ -45,6 +47,8 @@ struct PostgresGlobalState : public GlobalTableFunctionState { explicit PostgresGlobalState(idx_t max_threads) : page_idx(0), batch_idx(0), max_threads(max_threads) { } + ~PostgresGlobalState(); + mutable mutex lock; idx_t page_idx; idx_t batch_idx; @@ -57,6 +61,8 @@ struct PostgresGlobalState : public GlobalTableFunctionState { PostgresConnection &GetConnection(); void SetConnection(PostgresConnection connection); void SetConnection(shared_ptr connection); + void SetPinnedConnection(PostgresPoolConnection pinned_connection, + dbconnector::attached::AttachedCatalog attached_catalog); bool TryOpenNewConnection(ClientContext &context, PostgresLocalState &lstate, const PostgresBindData &bind_data); idx_t MaxThreads() const override { @@ -65,6 +71,8 @@ struct PostgresGlobalState : public GlobalTableFunctionState { private: PostgresConnection connection; + PostgresPoolConnection pinned_connection; + dbconnector::attached::AttachedCatalog attached_catalog; }; static void PostgresGetSnapshot(ClientContext &context, const PostgresBindData &bind_data, @@ -302,18 +310,34 @@ static unique_ptr PostgresInitGlobalState(ClientContex attached_catalog = PostgresCatalog::Lookup(context, bind_data.catalog_name); pg_catalog = &attached_catalog.Get(); + optional_ptr con_ptr = nullptr; + PostgresPoolConnection pinned_connection; + + if (bind_data.pinned_connection_id > 0) { + pinned_connection = pg_catalog->GetConnectionPool().UnpinConnection(bind_data.pinned_connection_id); + PostgresConnection &conn = pinned_connection.GetConnection(); + con_ptr = &conn; + } else { + auto &transaction = Transaction::Get(context, *pg_catalog).Cast(); + if (bind_data.use_transaction) { + con_ptr = &transaction.GetConnection(); + } else { + con_ptr = &transaction.GetConnectionWithoutTransaction(); + } + } + PostgresConnection &con = *con_ptr; + + result->SetConnection(con.GetConnection()); + if (pinned_connection) { + result->SetPinnedConnection(std::move(pinned_connection), std::move(attached_catalog)); + } + if (bind_data.command_only) { // A command (no result columns). Execute it here — at execution time, so EXPLAIN does not // fire it — and hand back a rowcount result through the materialized-collection // fast path (PostgresScan returns gstate.collection directly, and GetLocalState short-circuits // when a collection is present, so no postgres scan connection is opened). - auto &transaction = Transaction::Get(context, *pg_catalog).Cast(); - unique_ptr pg_res; - if (bind_data.use_transaction) { - pg_res = transaction.Query(bind_data.sql); - } else { - pg_res = transaction.QueryWithoutTransaction(bind_data.sql); - } + unique_ptr pg_res = con.Query(context, bind_data.sql); vector rowcount_types {LogicalType::BIGINT}; auto materialized = make_uniq(Allocator::Get(context), rowcount_types); DataChunk rowcount_chunk; @@ -329,9 +353,6 @@ static unique_ptr PostgresInitGlobalState(ClientContex return std::move(result); } - auto &transaction = Transaction::Get(context, *pg_catalog).Cast(); - auto &con = bind_data.use_transaction ? transaction.GetConnection() : transaction.GetConnectionWithoutTransaction(); - result->SetConnection(con.GetConnection()); if (bind_data.requires_materialization) { // if requires_materialization is enabled we scan and materialize the table in its entirety up-front vector types; @@ -452,6 +473,27 @@ static unique_ptr GetLocalState(ClientContext &context, return std::move(local_state); } +void PostgresGlobalState::SetPinnedConnection(PostgresPoolConnection pinned_connection, + dbconnector::attached::AttachedCatalog attached_catalog) { + if (this->pinned_connection) { + throw InternalException("Invalid global state: pinned connection alread initialized"); + } + this->pinned_connection = std::move(pinned_connection); + this->attached_catalog = std::move(attached_catalog); +} + +PostgresGlobalState::~PostgresGlobalState() { + connection.Close(); + if (pinned_connection) { + try { + PostgresCatalog &catalog = attached_catalog.Get(); + catalog.GetConnectionPool().PinConnection(std::move(pinned_connection)); + } catch (...) { + // suppress + } + } +} + static unique_ptr PostgresInitLocalState(ExecutionContext &context, TableFunctionInitInput &input, GlobalTableFunctionState *global_state) { diff --git a/src/postgres_utils.cpp b/src/postgres_utils.cpp index 8dcbc296d..cd6e48c83 100644 --- a/src/postgres_utils.cpp +++ b/src/postgres_utils.cpp @@ -3,6 +3,7 @@ #include "dbconnector/query/query_writer.hpp" +#include "storage/postgres_catalog.hpp" #include "storage/postgres_schema_entry.hpp" #include "storage/postgres_transaction.hpp" #include "postgres_scanner.hpp" @@ -787,4 +788,17 @@ PostgresQueryMode PostgresUtils::ExtractQueryMode(Value &value) { } } +PostgresCatalog &PostgresUtils::GetCatalogByName(ClientContext &context, const std::string &name) { + auto &db_manager = DatabaseManager::Get(context); + auto db = db_manager.GetDatabase(context, Identifier(name)); + if (!db) { + throw BinderException("Failed to find attached database \"%s\"", name); + } + auto &catalog = db->GetCatalog(); + if (catalog.GetCatalogType() != "postgres") { + throw BinderException("Attached database \"%s\" does not refer to a Postgres database", name); + } + return catalog.Cast(); +} + } // namespace duckdb diff --git a/test/sql/scanner/query_pin_connection.test b/test/sql/scanner/query_pin_connection.test new file mode 100644 index 000000000..196417567 --- /dev/null +++ b/test/sql/scanner/query_pin_connection.test @@ -0,0 +1,96 @@ +# name: test/sql/scanner/query_pin_connection.test +# description: Test connection pinning with postgres_query +# group: [scanner] + +require postgres_scanner + +require-env POSTGRES_TEST_DATABASE_AVAILABLE + +statement ok +ATTACH 'dbname=postgresscanner' AS s (TYPE POSTGRES) + +statement error +SELECT postgres_pin_connection() +---- +No function matches the given name and argument types + +statement error +SELECT postgres_pin_connection(NULL) +---- +Specified attached database name must be not null + +statement error +SELECT postgres_pin_connection('fail') +---- +Failed to find attached database + +statement ok +SET VARIABLE conn = postgres_pin_connection('s') + +statement error +SELECT postgres_close_pinned_connection('fail', 0) +---- +Failed to find attached database + +statement error +SELECT postgres_close_pinned_connection('s', 0) +---- +not found in the specified attached database + +statement ok +SELECT postgres_close_pinned_connection('s', getvariable('conn')) + +statement error +SELECT postgres_close_pinned_connection('s', getvariable('conn')) +---- +not found in the specified attached database + +statement ok +RESET VARIABLE conn + +statement ok +SET VARIABLE conn = postgres_pin_connection('s') + +statement ok +CALL postgres_execute('s', 'CREATE TEMPORARY TABLE temp1(col1 INT)', connection=getvariable('conn')) + +statement ok +CALL postgres_execute('s', 'INSERT INTO temp1 values(42)', connection=getvariable('conn')) + +query I +FROM postgres_query('s', 'SELECT * FROM temp1', connection=getvariable('conn')) +---- +42 + +statement error +FROM postgres_query('s', 'SELECT 42', connection=getvariable('conn')) +UNION ALL +FROM postgres_query('s', 'SELECT 43', connection=getvariable('conn')) +---- +not found in the specified attached database + +statement ok +SET VARIABLE conn2 = postgres_pin_connection('s') + +query I +FROM postgres_query('s', 'SELECT 42', connection=getvariable('conn')) +UNION ALL +FROM postgres_query('s', 'SELECT 43', connection=getvariable('conn2')) +---- +42 +43 + +statement ok +SELECT postgres_close_pinned_connection('s', getvariable('conn2')) + +statement ok +SELECT postgres_close_pinned_connection('s', getvariable('conn')) + +statement ok +RESET VARIABLE conn1 + +statement ok +RESET VARIABLE conn + +statement ok +DETACH s \ No newline at end of file