From b5c2c5d6e0017b7a250ad8b12c05c2f26dd9ce90 Mon Sep 17 00:00:00 2001 From: Alex Kasko Date: Sat, 4 Apr 2026 01:06:22 +0100 Subject: [PATCH] Pool: expose health check query option This PR adds `pg_pool_health_check_query` option that is run to check the connection when it is taken from the pool or returned to the pool. The option is only effective for the Postgres databases attached after the option is set. To disable the health check query set this option to `NULL`: ```sql SET pg_pool_health_check_query = NULL ``` Testing: option name check is added, behavior test coverage is pending on the pool introspection support. --- src/include/postgres_connection.hpp | 4 +- .../storage/postgres_connection_pool.hpp | 13 +++- src/postgres_connection.cpp | 11 ++-- src/postgres_extension.cpp | 3 + src/storage/postgres_catalog.cpp | 5 -- src/storage/postgres_connection_pool.cpp | 59 +++++++++++++++---- test/sql/storage/attach_connection_pool.test | 7 +++ test/sql/storage/attach_multi_join.test | 2 +- 8 files changed, 78 insertions(+), 26 deletions(-) diff --git a/src/include/postgres_connection.hpp b/src/include/postgres_connection.hpp index b846e7433..8d5cc7b72 100644 --- a/src/include/postgres_connection.hpp +++ b/src/include/postgres_connection.hpp @@ -73,8 +73,8 @@ class PostgresConnection { bool IsOpen(); void Close(); - bool PingServer(); - void Reset(); + bool PingServer(const std::string &health_check_query); + void Reset(const std::string &health_check_query); shared_ptr GetConnection() { return connection; diff --git a/src/include/storage/postgres_connection_pool.hpp b/src/include/storage/postgres_connection_pool.hpp index 3c3277642..2ab44893d 100644 --- a/src/include/storage/postgres_connection_pool.hpp +++ b/src/include/storage/postgres_connection_pool.hpp @@ -8,6 +8,8 @@ #pragma once +#include + #include "duckdb/common/common.hpp" #include "duckdb/common/mutex.hpp" #include "duckdb/common/optional_ptr.hpp" @@ -24,7 +26,7 @@ using PostgresPoolConnection = dbconnector::pool::PooledConnection { public: - PostgresConnectionPool(PostgresCatalog &postgres_catalog, idx_t maximum_connections = DefaultPoolSize()); + PostgresConnectionPool(PostgresCatalog &postgres_catalog); public: bool TryGetConnection(PostgresPoolConnection &connection); @@ -32,7 +34,11 @@ class PostgresConnectionPool : public dbconnector::pool::ConnectionPool CreateNewConnection() override; @@ -42,7 +48,8 @@ class PostgresConnectionPool : public dbconnector::pool::ConnectionPoolSetMaxConnections(UBigIntValue::Get(connection_limit)); - } auto connection = connection_pool->GetConnection(); this->version = connection.GetConnection().GetPostgresVersion(context); diff --git a/src/storage/postgres_connection_pool.cpp b/src/storage/postgres_connection_pool.cpp index 9ba89bea3..c82d7da26 100644 --- a/src/storage/postgres_connection_pool.cpp +++ b/src/storage/postgres_connection_pool.cpp @@ -6,9 +6,22 @@ namespace duckdb { -PostgresConnectionPool::PostgresConnectionPool(PostgresCatalog &postgres_catalog, idx_t maximum_connections_p) - : dbconnector::pool::ConnectionPool(CreateConfig(maximum_connections_p)), - postgres_catalog(postgres_catalog) { +static dbconnector::pool::ConnectionPoolConfig CreateConfig(PostgresCatalog &postgres_catalog); + +static std::string GetHealthCheckQueryFromConfig(PostgresCatalog &postgres_catalog) { + Value val; + if (postgres_catalog.GetDatabase().TryGetCurrentSetting("pg_pool_health_check_query", val)) { + if (val.IsNull()) { + return std::string(); + } + return StringValue::Get(val); + } + return PostgresConnectionPool::DefaultHealthCheckQuery(); +} + +PostgresConnectionPool::PostgresConnectionPool(PostgresCatalog &postgres_catalog) + : dbconnector::pool::ConnectionPool(CreateConfig(postgres_catalog)), + postgres_catalog(postgres_catalog), health_check_query(GetHealthCheckQueryFromConfig(postgres_catalog)) { } PostgresPoolConnection PostgresConnectionPool::ForceGetConnection() { @@ -39,23 +52,47 @@ bool PostgresConnectionPool::CheckConnectionHealthy(PostgresConnection &conn) { if (!conn.IsOpen()) { return false; } - return conn.PingServer(); + std::string query = GetHealthCheckQuery(); + return conn.PingServer(query); } void PostgresConnectionPool::ResetConnection(PostgresConnection &conn) { - conn.Reset(); + std::string query = GetHealthCheckQuery(); + conn.Reset(query); } -dbconnector::pool::ConnectionPoolConfig PostgresConnectionPool::CreateConfig(idx_t max_connections) { - dbconnector::pool::ConnectionPoolConfig config; - config.max_connections = max_connections; - return config; +std::string PostgresConnectionPool::GetHealthCheckQuery() { + std::lock_guard guard(config_mutex); + return std::string(health_check_query.data(), health_check_query.length()); } -idx_t PostgresConnectionPool::DefaultPoolSize() noexcept { +void PostgresConnectionPool::SetHealthCheckQuery(const std::string &query) { + std::lock_guard guard(config_mutex); + this->health_check_query = std::string(query.data(), query.length()); +} + +idx_t PostgresConnectionPool::DefaultPoolSize() { idx_t detected = static_cast(std::thread::hardware_concurrency()); idx_t default_num = static_cast(8); - return detected < default_num ? detected : default_num; + return detected > default_num ? detected : default_num; +} + +std::string PostgresConnectionPool::DefaultHealthCheckQuery() { + return "SELECT 1"; +} + +static dbconnector::pool::ConnectionPoolConfig CreateConfig(PostgresCatalog &postgres_catalog) { + DatabaseInstance &db = postgres_catalog.GetDatabase(); + + Value connection_limit; + uint64_t max_connections = PostgresConnectionPool::DefaultPoolSize(); + if (db.TryGetCurrentSetting("pg_connection_limit", connection_limit) && !connection_limit.IsNull()) { + max_connections = UBigIntValue::Get(connection_limit); + } + + dbconnector::pool::ConnectionPoolConfig config; + config.max_connections = max_connections; + return config; } } // namespace duckdb diff --git a/test/sql/storage/attach_connection_pool.test b/test/sql/storage/attach_connection_pool.test index 494a25cc6..639fde915 100644 --- a/test/sql/storage/attach_connection_pool.test +++ b/test/sql/storage/attach_connection_pool.test @@ -59,3 +59,10 @@ query I SELECT COUNT(*) FROM connection_pool ---- 1000000 + +# todo: pool introspection is required to check the effect +statement ok +SET pg_pool_health_check_query = 'SELECT FAIL' + +statement ok +RESET pg_pool_health_check_query diff --git a/test/sql/storage/attach_multi_join.test b/test/sql/storage/attach_multi_join.test index 5e88708df..4226de960 100644 --- a/test/sql/storage/attach_multi_join.test +++ b/test/sql/storage/attach_multi_join.test @@ -29,7 +29,7 @@ select * from test1 a left join test2 b on a.id = b.id left join test3 c on a.id ---- statement ok -SET pg_connection_limit=1 +SET pg_connection_limit=2 query III select * from test1 a left join test2 b on a.id = b.id left join test3 c on a.id = c.id;