Allow running postgres_query on a specific conn

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

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

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;

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

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<string_t>(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<uint64_t>(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<string_t>(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<uint64_t>(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

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<FunctionData> BindDML(ClientContext &context, TableFunctionBindInput &input,
                                         vector<LogicalType> &return_types, vector<Identifier> &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<PostgresBindData>(context, pg_catalog.GetName());
+	auto result = make_uniq<PostgresBindData>(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<FunctionData> 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<FunctionData> PGQueryBind(ClientContext &context, TableFunctionBindInput &input,
-                                            vector<LogicalType> &return_types, vector<Identifier> &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<string>();
-	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<PostgresCatalog>();
-	auto &transaction = Transaction::Get(context, catalog).Cast<PostgresTransaction>();
-	auto sql = input.inputs[1].GetValue<string>();
-	// 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<Value> param_values;
+static vector<Value> 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<FunctionData> 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<Value>();
+}
+
+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<FunctionData> PGQueryBind(ClientContext &context, TableFunctionBindInput &input,
+                                            vector<LogicalType> &return_types, vector<Identifier> &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<string>();
+	auto &catalog = PostgresUtils::GetCatalogByName(context, db_name);
+	auto input_sql = input.inputs[1].GetValue<string>();
+	auto sql = StripTrailingSemicolon(std::move(input_sql));
+	bool use_transaction = ExtractFlag(input, "use_transaction", true);
+	vector<Value> param_values = ExtractParams(input);
+	uint64_t pinned_connection_id = ExtractPinnedConnId(input);
+
+	optional_ptr<PostgresConnection> 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<PostgresTransaction>();
+		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<FunctionData> 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<PostgresBindData>(context, pg_catalog.GetName());
+	auto result = make_uniq<PostgresBindData>(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<FunctionData> 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;
 }