From c0dff6ca64126a29ef1e395dbe8cc2176f465b68 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Fri, 1 Jan 2021 15:02:28 +0530
Subject: [PATCH v5 3/3] 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             | 35 ++++++++++++---
 .../postgres_fdw/expected/postgres_fdw.out    | 44 ++++++++++++++++++-
 contrib/postgres_fdw/option.c                 |  9 +++-
 contrib/postgres_fdw/sql/postgres_fdw.sql     | 26 +++++++++++
 doc/src/sgml/postgres-fdw.sgml                | 42 ++++++++++++++++++
 5 files changed, 149 insertions(+), 7 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 323242c1d6..acb97d40cf 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -62,6 +62,8 @@ typedef struct ConnCacheEntry
 	Oid			serverid;
 	uint32		server_hashvalue;	/* hash value of foreign server OID */
 	uint32		mapping_hashvalue;	/* hash value of user mapping OID */
+	/* Keep or discard this connection at the end of xact */
+	bool            keep_connection;
 } ConnCacheEntry;
 
 /*
@@ -124,6 +126,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)
@@ -261,6 +265,15 @@ GetConnection(UserMapping *user, bool will_prep_stmt)
 		begin_remote_xact(entry);
 	}
 
+	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;
 
@@ -285,6 +298,7 @@ make_new_connection(ConnCacheEntry *entry, UserMapping *user)
 	entry->changing_xact_state = false;
 	entry->invalidated = false;
 	entry->serverid = server->serverid;
+	entry->keep_connection = true;
 	entry->server_hashvalue =
 		GetSysCacheHashValue1(FOREIGNSERVEROID,
 							  ObjectIdGetDatum(server->serverid));
@@ -963,15 +977,18 @@ pgfdw_xact_callback(XactEvent event, void *arg)
 
 		/*
 		 * If the connection isn't in a good idle state or it is marked as
-		 * invalid or keep_connections GUC is false and this connection is used
-		 * in current xact, then discard it to recover. Next GetConnection will
-		 * open a new connection.
+		 * invalid or this connection is used in current xact and
+		 * keep_connections GUC is false or GUC is true but the server level
+		 * option is false,  then discard it to recover. Next GetConnection will
+		 * open a new connection. Note that keep_connections GUC overrides the
+		 * server level keep_connection option.
 		 */
 		if (PQstatus(entry->conn) != CONNECTION_OK ||
 			PQtransactionStatus(entry->conn) != PQTRANS_IDLE ||
 			entry->changing_xact_state ||
 			entry->invalidated ||
-			(used_in_current_xact && !keep_connections))
+			(used_in_current_xact &&
+			(!keep_connections || !entry->keep_connection)))
 		{
 			elog(DEBUG3, "discarding connection %p", entry->conn);
 			disconnect_pg_server(entry);
@@ -1132,7 +1149,15 @@ pgfdw_inval_callback(Datum arg, int cacheid, uint32 hashvalue)
 	if (!ConnectionHash)
 		return;
 
-	/* ConnectionHash must exist already, if we're registered */
+	/*
+	 * Since we can discard ConnectionHash with postgres_fdw_disconnect, we may
+	 * have a NULL ConnectionHash. So return in that case. We do not need to
+	 * invalidate the cache entries as the cache itself would have discarded
+	 * with postgres_fdw_disconnect.
+	 */
+	if (!ConnectionHash)
+		return;
+
 	hash_seq_init(&scan, ConnectionHash);
 	while ((entry = (ConnCacheEntry *) hash_seq_search(&scan)))
 	{
diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index e08d9f38be..100a342c8b 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -8933,7 +8933,7 @@ DO $d$
     END;
 $d$;
 ERROR:  invalid option "password"
-HINT:  Valid options in this context are: service, passfile, channel_binding, connect_timeout, dbname, host, hostaddr, port, options, application_name, keepalives, keepalives_idle, keepalives_interval, keepalives_count, tcp_user_timeout, sslmode, sslcompression, sslcert, sslkey, sslrootcert, sslcrl, requirepeer, ssl_min_protocol_version, ssl_max_protocol_version, gssencmode, krbsrvname, gsslib, target_session_attrs, use_remote_estimate, fdw_startup_cost, fdw_tuple_cost, extensions, updatable, fetch_size
+HINT:  Valid options in this context are: service, passfile, channel_binding, connect_timeout, dbname, host, hostaddr, port, options, application_name, keepalives, keepalives_idle, keepalives_interval, keepalives_count, tcp_user_timeout, sslmode, sslcompression, sslcert, sslkey, sslrootcert, sslcrl, requirepeer, ssl_min_protocol_version, ssl_max_protocol_version, gssencmode, krbsrvname, gsslib, target_session_attrs, use_remote_estimate, fdw_startup_cost, fdw_tuple_cost, extensions, updatable, fetch_size, keep_connection
 CONTEXT:  SQL statement "ALTER SERVER loopback_nopw OPTIONS (ADD password 'dummypw')"
 PL/pgSQL function inline_code_block line 3 at EXECUTE
 -- If we add a password for our user mapping instead, we should get a different
@@ -9498,5 +9498,47 @@ SELECT * FROM postgres_fdw_disconnect();	/* TRUE */
  t
 (1 row)
 
+-- ===================================================================
+-- Test foreign server level option keep_connection
+-- ===================================================================
+-- By default, the connections associated with foreign server are cached i.e.
+-- keep_connection option is on. Set it to off.
+ALTER SERVER loopback OPTIONS (keep_connection 'off');
+-- loopback server connection is closed by the local session at the end of xact
+-- as the keep_connection was set to off.
+SELECT 1 FROM ft1 LIMIT 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+-- Should output nothing.
+SELECT * FROM fdw_unnest(postgres_fdw_get_connections()) ORDER BY 1;
+ fdw_unnest 
+------------
+(0 rows)
+
+ALTER SERVER loopback OPTIONS (SET keep_connection 'on');
+-- loopback server connection should get cached.
+SELECT 1 FROM ft1 LIMIT 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+-- Should output loopback.
+SELECT * FROM fdw_unnest(postgres_fdw_get_connections()) ORDER BY 1;
+    fdw_unnest    
+------------------
+ (loopback, true)
+(1 row)
+
+-- Discard loopback connection.
+SELECT * FROM postgres_fdw_disconnect();	/* TRUE */
+ postgres_fdw_disconnect 
+-------------------------
+ t
+(1 row)
+
 -- Clean up.
 DROP FUNCTION fdw_unnest;
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}
 	};
 
diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index b37aa569d6..44e076b4f4 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2905,5 +2905,31 @@ SELECT * FROM fdw_unnest(postgres_fdw_get_connections()) ORDER BY 1;
 -- Discard loopback connection.
 SELECT * FROM postgres_fdw_disconnect();	/* TRUE */
 
+-- ===================================================================
+-- Test foreign server level option keep_connection
+-- ===================================================================
+
+-- By default, the connections associated with foreign server are cached i.e.
+-- keep_connection option is on. Set it to off.
+ALTER SERVER loopback OPTIONS (keep_connection 'off');
+
+-- loopback server connection is closed by the local session at the end of xact
+-- as the keep_connection was set to off.
+SELECT 1 FROM ft1 LIMIT 1;
+
+-- Should output nothing.
+SELECT * FROM fdw_unnest(postgres_fdw_get_connections()) ORDER BY 1;
+
+ALTER SERVER loopback OPTIONS (SET keep_connection 'on');
+
+-- loopback server connection should get cached.
+SELECT 1 FROM ft1 LIMIT 1;
+
+-- Should output loopback.
+SELECT * FROM fdw_unnest(postgres_fdw_get_connections()) ORDER BY 1;
+
+-- Discard loopback connection.
+SELECT * FROM postgres_fdw_disconnect();	/* TRUE */
+
 -- Clean up.
 DROP FUNCTION fdw_unnest;
\ No newline at end of file
diff --git a/doc/src/sgml/postgres-fdw.sgml b/doc/src/sgml/postgres-fdw.sgml
index 15def846f5..a62460704f 100644
--- a/doc/src/sgml/postgres-fdw.sgml
+++ b/doc/src/sgml/postgres-fdw.sgml
@@ -478,6 +478,35 @@ OPTIONS (ADD password_required 'false');
 
   </sect3>
 
+  <sect3>
+    <title>Connection Management Options</title>
+
+    <para>
+     By default the foreign server connections made with
+     <filename>postgres_fdw</filename> are kept in local session for re-use.
+     This may be overridden using the following
+     <xref linkend="sql-createserver"/> option:
+    </para>
+ 
+    <variablelist>
+ 
+     <varlistentry>
+      <term><literal>keep_connection</literal></term>
+      <listitem>
+       <para>
+        This option controls whether <filename>postgres_fdw</filename> keeps the
+        server connection that is made with a specific foreign server. It can be
+        specified for a foreign server. Default is <literal>on</literal>. When
+        set to <literal>off</literal>, the associated foreign server connection
+        is discarded at the end of the transaction.
+      </para>
+ 
+      </listitem>
+     </varlistentry>
+ 
+    </variablelist>
+   </sect3>
+
 <sect2>
   <title>Functions</title>
 
@@ -542,6 +571,14 @@ OPTIONS (ADD password_required 'false');
       connection is discarded at the end of transaction in which it is used.
      </para>
 
+      <para>
+       <varname>postgres_fdw.keep_connections</varname> configuration parameter
+       overrides the server level <literal>keep_connection</literal> option.
+       Which means that when the configuration parameter is set to
+       <literal>on</literal>, irrespective of the server option
+       <literal>keep_connection</literal> setting, the connections are discarded.
+      </para>
+
      <para>
       Note that setting <varname>postgres_fdw.keep_connections</varname> to
       <literal>off</literal> does not discard any previously made and still open
@@ -588,6 +625,11 @@ OPTIONS (ADD password_required 'false');
    servers. Each connection is discarded at the end of transaction in which it
    is used.
   
+   A server level option, <literal>keep_connection</literal> that is used with
+   <xref linkend="sql-createserver"/>. Default being <literal>on</literal>,
+   when set to <literal>off</literal> the local session doesn't keep remote
+   connection associated with the foreign server. The connection is
+   discarded at the end of the transaction.
   </para>
  </sect2>
 
-- 
2.25.1

