From 1d3304f680a1358f1d1f7a5151bb4b6e4446a486 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Mon, 23 Nov 2020 11:57:04 +0530
Subject: [PATCH v1] postgres_fdw server level option, keep_connection to not
 cache connection

This patch adds a new server level option, keep_connection,
default being on, when set to off, the local session doesn't cache
the connections associated with the foreign server.
---
 contrib/postgres_fdw/connection.c | 16 +++++++++++++++-
 contrib/postgres_fdw/option.c     |  9 ++++++++-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index eae5113b60..dfea4b3da5 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -62,6 +62,8 @@ typedef struct ConnCacheEntry
 	uint32		mapping_hashvalue;	/* hash value of user mapping OID */
 	/* Is this connection used in current xact?*/
 	bool		used_in_current_xact;
+	/* Keep or discard this connection at the end of xact */
+	bool		keep_connection;
 } ConnCacheEntry;
 
 /*
@@ -123,6 +125,8 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 	ConnCacheEntry *entry;
 	ConnCacheKey key;
 	MemoryContext ccxt = CurrentMemoryContext;
+	ListCell   *lc;
+	ForeignServer *server;
 
 	/* First time through, initialize connection cache hashtable */
 	if (ConnectionHash == NULL)
@@ -265,6 +269,15 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 
 	entry->used_in_current_xact = true;
 
+	server = GetForeignServer(user->serverid);
+	foreach(lc, server->options)
+	{
+		DefElem    *def = (DefElem *) lfirst(lc);
+
+		if (strcmp(def->defname, "keep_connection") == 0)
+			entry->keep_connection = defGetBoolean(def);
+	}
+
 	/* Remember if caller will prepare statements */
 	entry->have_prep_stmt |= will_prep_stmt;
 
@@ -289,6 +302,7 @@ make_new_connection(ConnCacheEntry *entry, UserMapping *user)
 	entry->changing_xact_state = false;
 	entry->invalidated = false;
 	entry->used_in_current_xact = false;
+	entry->keep_connection = true;
 	entry->server_hashvalue =
 		GetSysCacheHashValue1(FOREIGNSERVEROID,
 							  ObjectIdGetDatum(server->serverid));
@@ -963,7 +977,7 @@ pgfdw_xact_callback(XactEvent event, void *arg)
 			PQtransactionStatus(entry->conn) != PQTRANS_IDLE ||
 			entry->changing_xact_state) ||
 			(entry->used_in_current_xact &&
-			!keep_connections))
+			(!keep_connections || !entry->keep_connection)))
 		{
 			elog(DEBUG3, "discarding connection %p", entry->conn);
 			disconnect_pg_server(entry);
diff --git a/contrib/postgres_fdw/option.c b/contrib/postgres_fdw/option.c
index 1a03e02263..0fe2eff878 100644
--- a/contrib/postgres_fdw/option.c
+++ b/contrib/postgres_fdw/option.c
@@ -107,7 +107,8 @@ postgres_fdw_validator(PG_FUNCTION_ARGS)
 		 * Validate option value, when we can do so without any context.
 		 */
 		if (strcmp(def->defname, "use_remote_estimate") == 0 ||
-			strcmp(def->defname, "updatable") == 0)
+			strcmp(def->defname, "updatable") == 0 ||
+			strcmp(def->defname, "keep_connection") == 0)
 		{
 			/* these accept only boolean values */
 			(void) defGetBoolean(def);
@@ -213,6 +214,12 @@ InitPgFdwOptions(void)
 		{"sslcert", UserMappingRelationId, true},
 		{"sslkey", UserMappingRelationId, true},
 
+		/*
+		 * If true, cache the connection associated with this server, otherwise
+		 * remove it at the end of the xact. Default is true.
+		 */
+		{"keep_connection", ForeignServerRelationId, false},
+
 		{NULL, InvalidOid, false}
 	};
 
-- 
2.25.1