Merge pull request #170 from Mytherin/connectionpoolsetting

Add setting pg_experimental_connection_cache that is disabled by default
duckdb · Jan 18, 2024 · 2a32e01 · 2a32e01
2 parents e3e4f20 + 9e9ebae
commit 2a32e01
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 0 deletions.
diff --git a/src/include/storage/postgres_connection_pool.hpp b/src/include/storage/postgres_connection_pool.hpp
@@ -49,6 +49,8 @@ class PostgresConnectionPool {
 	void ReturnConnection(PostgresConnection connection);
 	void SetMaximumConnections(idx_t new_max);
 
+	static void PostgresSetConnectionCache(ClientContext &context, SetScope scope, Value &parameter);
+
 private:
 	PostgresCatalog &postgres_catalog;
 	mutex connection_lock;

diff --git a/src/postgres_extension.cpp b/src/postgres_extension.cpp
@@ -64,6 +64,9 @@ static void LoadInternal(DatabaseInstance &db) {
 	config.AddExtensionOption("pg_array_as_varchar",
 	                          "Read Postgres arrays as varchar - enables reading mixed dimensional arrays",
 	                          LogicalType::BOOLEAN, Value::BOOLEAN(false));
+	config.AddExtensionOption("pg_experimental_connection_cache",
+							  "Whether or not to use the connection cache (currently experimental)", LogicalType::BOOLEAN,
+							  Value::BOOLEAN(false), PostgresConnectionPool::PostgresSetConnectionCache);
 	config.AddExtensionOption("pg_experimental_filter_pushdown",
 	                          "Whether or not to use filter pushdown (currently experimental)", LogicalType::BOOLEAN,
 	                          Value::BOOLEAN(false));

diff --git a/src/storage/postgres_connection_pool.cpp b/src/storage/postgres_connection_pool.cpp
@@ -2,6 +2,7 @@
 #include "storage/postgres_catalog.hpp"
 
 namespace duckdb {
+static bool pg_use_connection_cache = false;
 
 PostgresPoolConnection::PostgresPoolConnection() : pool(nullptr) {
 }
@@ -61,6 +62,13 @@ bool PostgresConnectionPool::TryGetConnection(PostgresPoolConnection &connection
 	return true;
 }
 
+void PostgresConnectionPool::PostgresSetConnectionCache(ClientContext &context, SetScope scope, Value &parameter) {
+	if (parameter.IsNull()) {
+		throw BinderException("Cannot be set to NULL");
+	}
+	pg_use_connection_cache = BooleanValue::Get(parameter);
+}
+
 PostgresPoolConnection PostgresConnectionPool::GetConnection() {
 	PostgresPoolConnection result;
 	if (!TryGetConnection(result)) {
@@ -82,6 +90,9 @@ void PostgresConnectionPool::ReturnConnection(PostgresConnection connection) {
 		// immediately
 		return;
 	}
+	if (!pg_use_connection_cache) {
+		return;
+	}
 	// check if the underlying connection is still usable
 	auto pg_con = connection.GetConn();
 	if (PQstatus(connection.GetConn()) != CONNECTION_OK) {

diff --git a/test/sql/storage/attach_connection_pool.test b/test/sql/storage/attach_connection_pool.test
@@ -6,6 +6,9 @@ require postgres_scanner
 
 require-env POSTGRES_TEST_DATABASE_AVAILABLE
 
+statement ok
+SET pg_experimental_connection_cache=true
+
 statement ok
 SET pg_connection_limit=1000