From f890ebb605c8a2aa27dcf87e1b36cbb48594ce59 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Fri, 4 Dec 2020 16:24:57 +0530
Subject: [PATCH v3] postgres_fdw function to discard cached connections

This patch introduces new function postgres_fdw_disconnect() when
called with a foreign server name discards the associated
connections with the server name. When called without any argument,
discards all the existing cached connections.

This patch also adds another function postgres_fdw_get_connections()
to get the list of all cached connections by corresponding foreign
server names.
---
 contrib/postgres_fdw/connection.c          | 164 +++++++++++++++++++++
 contrib/postgres_fdw/postgres_fdw--1.0.sql |  15 ++
 2 files changed, 179 insertions(+)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index ab3226287d..65695ebb67 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -14,6 +14,7 @@
 
 #include "access/htup_details.h"
 #include "access/xact.h"
+#include "catalog/pg_foreign_server.h"
 #include "catalog/pg_user_mapping.h"
 #include "commands/defrem.h"
 #include "mb/pg_wchar.h"
@@ -22,6 +23,7 @@
 #include "postgres_fdw.h"
 #include "storage/fd.h"
 #include "storage/latch.h"
+#include "utils/builtins.h"
 #include "utils/datetime.h"
 #include "utils/hsearch.h"
 #include "utils/inval.h"
@@ -73,6 +75,12 @@ static unsigned int prep_stmt_number = 0;
 /* tracks whether any work is needed in callback functions */
 static bool xact_got_connection = false;
 
+/*
+ * SQL functions
+ */
+PG_FUNCTION_INFO_V1(postgres_fdw_get_connections);
+PG_FUNCTION_INFO_V1(postgres_fdw_disconnect);
+
 /* prototypes of private functions */
 static void make_new_connection(ConnCacheEntry *entry, UserMapping *user);
 static PGconn *connect_pg_server(ForeignServer *server, UserMapping *user);
@@ -94,6 +102,7 @@ static bool pgfdw_exec_cleanup_query(PGconn *conn, const char *query,
 static bool pgfdw_get_cleanup_result(PGconn *conn, TimestampTz endtime,
 									 PGresult **result);
 static bool UserMappingPasswordRequired(UserMapping *user);
+static bool disconnect_cached_connections(uint32 hashvalue, bool all);
 
 /*
  * Get a PGconn which can be used to execute queries on the remote PostgreSQL
@@ -1325,3 +1334,158 @@ exit:	;
 		*result = last_res;
 	return timed_out;
 }
+
+/*
+ * List all cached connections by corresponding foreign server names.
+ *
+ * Takes no paramaters as input. Returns an array of all foreign server names
+ * with which connections exist. Returns NULL, in case there are no cached
+ * connections at all.
+ */
+Datum
+postgres_fdw_get_connections(PG_FUNCTION_ARGS)
+{
+	ArrayBuildState *astate = NULL;
+
+	if (ConnectionHash)
+	{
+		HASH_SEQ_STATUS	scan;
+		ConnCacheEntry	*entry;
+
+		hash_seq_init(&scan, ConnectionHash);
+		while ((entry = (ConnCacheEntry *) hash_seq_search(&scan)))
+		{
+			Form_pg_user_mapping umap;
+			HeapTuple	umaptup;
+			Form_pg_foreign_server fsrv;
+			HeapTuple	fsrvtup;
+
+			/* We only look for active and open remote connections. */
+			if (entry->invalidated || !entry->conn)
+				continue;
+
+			umaptup = SearchSysCache1(USERMAPPINGOID,
+									  ObjectIdGetDatum(entry->key));
+
+			if (!HeapTupleIsValid(umaptup))
+				elog(ERROR, "cache lookup failed for user mapping with OID %u",
+					 entry->key);
+
+			umap = (Form_pg_user_mapping) GETSTRUCT(umaptup);
+
+			fsrvtup = SearchSysCache1(FOREIGNSERVEROID,
+									  ObjectIdGetDatum(umap->umserver));
+
+			if (!HeapTupleIsValid(fsrvtup))
+				elog(ERROR, "cache lookup failed for foreign server with OID %u",
+					 umap->umserver);
+
+			fsrv = (Form_pg_foreign_server) GETSTRUCT(fsrvtup);
+
+			/* stash away current value */
+			astate = accumArrayResult(astate,
+									  CStringGetTextDatum(NameStr(fsrv->srvname)),
+									  false, TEXTOID, CurrentMemoryContext);
+
+			ReleaseSysCache(umaptup);
+			ReleaseSysCache(fsrvtup);
+		}
+	}
+
+	if (astate)
+		PG_RETURN_ARRAYTYPE_P(makeArrayResult(astate,
+											  CurrentMemoryContext));
+	else
+		PG_RETURN_NULL();
+}
+
+/*
+ * Disconnects the cached connections.
+ *
+ * If server name is provided as input, it disconnects the associated cached
+ * connections. Otherwise all the cached connections are disconnected and the
+ * cache is destroyed.
+ *
+ * Returns false if the cache is empty or if the cache is non empty and server
+ * name is provided and it exists but it has no associated entry in the cache.
+ * An error is emitted when the given foreign server does not exist.
+ * In all other cases true is returned.
+ */
+Datum
+postgres_fdw_disconnect(PG_FUNCTION_ARGS)
+{
+	bool	result = false;
+
+	if (PG_NARGS() == 1)
+	{
+		char			*servername = NULL;
+		ForeignServer	*server = NULL;
+
+		servername = text_to_cstring(PG_GETARG_TEXT_PP(0));
+		server = GetForeignServerByName(servername, true);
+
+		if (server)
+		{
+			uint32	hashvalue;
+
+			hashvalue =
+				GetSysCacheHashValue1(FOREIGNSERVEROID,
+									  ObjectIdGetDatum(server->serverid));
+
+			if (ConnectionHash)
+				result = disconnect_cached_connections(hashvalue, false);
+		}
+		else
+			ereport(ERROR,
+					(errcode(ERRCODE_CONNECTION_DOES_NOT_EXIST),
+					 errmsg("foreign server \"%s\" does not exist", servername)));
+	}
+	else
+	{
+		if (ConnectionHash)
+			result = disconnect_cached_connections(0, true);
+	}
+
+	PG_RETURN_BOOL(result);
+}
+
+/*
+ * Workhorse to disconnect the cached connections.
+ *
+ * If all is true, all the cached connections are disconnected and cache is
+ * destroyed. Otherwise, the entries with the given hashvalue are disconnected.
+ *
+ * Returns true in the following cases: either the cache is destroyed now or
+ * at least a single entry with the given hashvalue exists. In all other cases
+ * false is returned.
+ */
+bool disconnect_cached_connections(uint32 hashvalue, bool all)
+{
+	bool			result = false;
+	HASH_SEQ_STATUS	scan;
+	ConnCacheEntry	*entry;
+
+	/*
+	 * We do not come here if the ConnectionHash is NULL. We handle it in the
+	 * caller.
+	 */
+	hash_seq_init(&scan, ConnectionHash);
+	while ((entry = (ConnCacheEntry *) hash_seq_search(&scan)))
+	{
+		if ((all || entry->server_hashvalue == hashvalue) &&
+			entry->conn)
+		{
+			disconnect_pg_server(entry);
+			result = true;
+		}
+	}
+
+	if (all)
+	{
+		hash_destroy(ConnectionHash);
+		ConnectionHash = NULL;
+		result = true;
+	}
+
+	return result;
+}
diff --git a/contrib/postgres_fdw/postgres_fdw--1.0.sql b/contrib/postgres_fdw/postgres_fdw--1.0.sql
index a0f0fc1bf4..edbb5911d2 100644
--- a/contrib/postgres_fdw/postgres_fdw--1.0.sql
+++ b/contrib/postgres_fdw/postgres_fdw--1.0.sql
@@ -16,3 +16,18 @@ LANGUAGE C STRICT;
 CREATE FOREIGN DATA WRAPPER postgres_fdw
   HANDLER postgres_fdw_handler
   VALIDATOR postgres_fdw_validator;
+
+CREATE FUNCTION postgres_fdw_get_connections ()
+RETURNS text[]
+AS 'MODULE_PATHNAME','postgres_fdw_get_connections'
+LANGUAGE C STRICT PARALLEL RESTRICTED;
+
+CREATE FUNCTION postgres_fdw_disconnect ()
+RETURNS bool
+AS 'MODULE_PATHNAME','postgres_fdw_disconnect'
+LANGUAGE C STRICT PARALLEL RESTRICTED;
+
+CREATE FUNCTION postgres_fdw_disconnect (text)
+RETURNS bool
+AS 'MODULE_PATHNAME','postgres_fdw_disconnect'
+LANGUAGE C STRICT PARALLEL RESTRICTED;
-- 
2.25.1

