From 8c875725c04ec3b86a041d7ffac8d7d3091f3a29 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Mon, 23 Nov 2020 11:53:33 +0530
Subject: [PATCH v1] postgres_fdw add keep_connections GUC to not cache
 connections

This patch adds a new GUC postgres_fdw.keep_connections, default
being on, when set to off no remote connections are cached by the
local session.
---
 contrib/postgres_fdw/connection.c   | 14 ++++++++++++--
 contrib/postgres_fdw/postgres_fdw.c | 16 ++++++++++++++++
 contrib/postgres_fdw/postgres_fdw.h |  3 +++
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 4f79897cc3..eae5113b60 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -60,6 +60,8 @@ typedef struct ConnCacheEntry
 	bool		invalidated;	/* true if reconnect is pending */
 	uint32		server_hashvalue;	/* hash value of foreign server OID */
 	uint32		mapping_hashvalue;	/* hash value of user mapping OID */
+	/* Is this connection used in current xact?*/
+	bool		used_in_current_xact;
 } ConnCacheEntry;
 
 /*
@@ -261,6 +263,8 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 		begin_remote_xact(entry);
 	}
 
+	entry->used_in_current_xact = true;
+
 	/* Remember if caller will prepare statements */
 	entry->have_prep_stmt |= will_prep_stmt;
 
@@ -284,6 +288,7 @@ make_new_connection(ConnCacheEntry *entry, UserMapping *user)
 	entry->have_error = false;
 	entry->changing_xact_state = false;
 	entry->invalidated = false;
+	entry->used_in_current_xact = false;
 	entry->server_hashvalue =
 		GetSysCacheHashValue1(FOREIGNSERVEROID,
 							  ObjectIdGetDatum(server->serverid));
@@ -954,13 +959,18 @@ pgfdw_xact_callback(XactEvent event, void *arg)
 		 * If the connection isn't in a good idle state, discard it to
 		 * recover. Next GetConnection will open a new connection.
 		 */
-		if (PQstatus(entry->conn) != CONNECTION_OK ||
+		if ((PQstatus(entry->conn) != CONNECTION_OK ||
 			PQtransactionStatus(entry->conn) != PQTRANS_IDLE ||
-			entry->changing_xact_state)
+			entry->changing_xact_state) ||
+			(entry->used_in_current_xact &&
+			!keep_connections))
 		{
 			elog(DEBUG3, "discarding connection %p", entry->conn);
 			disconnect_pg_server(entry);
 		}
+
+		if (entry->used_in_current_xact)
+			entry->used_in_current_xact = false;
 	}
 
 	/*
diff --git a/contrib/postgres_fdw/postgres_fdw.c b/contrib/postgres_fdw/postgres_fdw.c
index 9c5aaacc51..05c8817224 100644
--- a/contrib/postgres_fdw/postgres_fdw.c
+++ b/contrib/postgres_fdw/postgres_fdw.c
@@ -301,6 +301,8 @@ typedef struct
 	List	   *already_used;	/* expressions already dealt with */
 } ec_member_foreign_arg;
 
+bool keep_connections = true;
+
 /*
  * SQL functions
  */
@@ -505,6 +507,20 @@ static void merge_fdw_options(PgFdwRelationInfo *fpinfo,
 							  const PgFdwRelationInfo *fpinfo_o,
 							  const PgFdwRelationInfo *fpinfo_i);
 
+void
+_PG_init(void)
+{
+	DefineCustomBoolVariable("postgres_fdw.keep_connections",
+							 "Enables postgres_fdw connection caching.",
+							 "When off postgres_fdw will close connections at the end of transaction.",
+							 &keep_connections,
+							 true,
+							 PGC_USERSET,
+							 0,
+							 NULL,
+							 NULL,
+							 NULL);
+}
 
 /*
  * Foreign-data wrapper handler function: return a struct with pointers
diff --git a/contrib/postgres_fdw/postgres_fdw.h b/contrib/postgres_fdw/postgres_fdw.h
index eef410db39..7f1bdb96d6 100644
--- a/contrib/postgres_fdw/postgres_fdw.h
+++ b/contrib/postgres_fdw/postgres_fdw.h
@@ -124,9 +124,12 @@ typedef struct PgFdwRelationInfo
 	int			relation_index;
 } PgFdwRelationInfo;
 
+extern bool keep_connections;
+
 /* in postgres_fdw.c */
 extern int	set_transmission_modes(void);
 extern void reset_transmission_modes(int nestlevel);
+extern void _PG_init(void);
 
 /* in connection.c */
 extern PGconn *GetConnection(UserMapping *user, bool will_prep_stmt);
-- 
2.25.1