From 3407b6a48699350ccf6008832ec1143ee1263618 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Thu, 21 Jan 2021 07:10:17 +0530
Subject: [PATCH v15] 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             | 25 +++++++++--
 .../postgres_fdw/expected/postgres_fdw.out    | 44 ++++++++++++++++++-
 contrib/postgres_fdw/option.c                 |  9 +++-
 contrib/postgres_fdw/sql/postgres_fdw.sql     | 19 ++++++++
 doc/src/sgml/postgres-fdw.sgml                | 43 ++++++++++++++++++
 5 files changed, 134 insertions(+), 6 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 4e65a4e47a..da054d38c2 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -62,6 +62,8 @@ typedef struct ConnCacheEntry
 	Oid			serverid;		/* foreign server OID used to get server name */
 	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));
@@ -956,15 +970,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);
diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index ffd0af4931..7a0053b908 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -8939,7 +8939,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, batch_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, batch_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
@@ -9454,3 +9454,45 @@ SELECT * FROM postgres_fdw_get_connections() ORDER BY 1;
  loopback    | 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 postgres_fdw_get_connections() ORDER BY 1;
+ server_name | valid 
+-------------+-------
+(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 postgres_fdw_get_connections() ORDER BY 1;
+ server_name | valid 
+-------------+-------
+ loopback    | t
+(1 row)
+
+-- Closes loopback connection and returns true.
+SELECT * FROM postgres_fdw_disconnect_all();
+ postgres_fdw_disconnect_all 
+-----------------------------
+ t
+(1 row)
+
diff --git a/contrib/postgres_fdw/option.c b/contrib/postgres_fdw/option.c
index 64698c4da3..75126a84bc 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);
@@ -227,6 +228,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 fc243ce972..8539e69895 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2911,3 +2911,22 @@ SET postgres_fdw.keep_connections TO on;
 SELECT 1 FROM ft1 LIMIT 1;
 -- Should output loopback.
 SELECT * FROM postgres_fdw_get_connections() ORDER BY 1;
+
+-- ===================================================================
+-- 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 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 postgres_fdw_get_connections() ORDER BY 1;
+-- Closes loopback connection and returns true.
+SELECT * FROM postgres_fdw_disconnect_all();
diff --git a/doc/src/sgml/postgres-fdw.sgml b/doc/src/sgml/postgres-fdw.sgml
index 7d9921094a..d385d253b6 100644
--- a/doc/src/sgml/postgres-fdw.sgml
+++ b/doc/src/sgml/postgres-fdw.sgml
@@ -490,6 +490,35 @@ OPTIONS (ADD password_required 'false');
    </para>
 
   </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>
 
 <sect2>
@@ -592,6 +621,14 @@ postgres=# SELECT * FROM postgres_fdw_disconnect_all();
       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
@@ -635,6 +672,12 @@ postgres=# SELECT * FROM postgres_fdw_disconnect_all();
    local session doesn't keep remote connections that are made to the foreign
    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

