From 10c4dffefeecac50e114c3cc96ff43e87f77b964 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Fri, 27 Nov 2020 06:53:18 +0530
Subject: [PATCH v1] 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.
---
 contrib/postgres_fdw/connection.c          | 98 ++++++++++++++++++++++
 contrib/postgres_fdw/postgres_fdw--1.0.sql | 10 +++
 2 files changed, 108 insertions(+)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index ab3226287d..7e7233650c 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -22,6 +22,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 +74,11 @@ 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_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 +100,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 +1332,94 @@ exit:	;
 		*result = last_res;
 	return timed_out;
 }
+
+/*
+ * 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 != NULL)
+		{
+			uint32	hashvalue;
+
+			hashvalue =
+				GetSysCacheHashValue1(FOREIGNSERVEROID,
+									  ObjectIdGetDatum(server->serverid));
+
+			if (ConnectionHash != NULL)
+				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 != NULL)
+			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 != NULL)
+		{
+			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..9f70ed1c32 100644
--- a/contrib/postgres_fdw/postgres_fdw--1.0.sql
+++ b/contrib/postgres_fdw/postgres_fdw--1.0.sql
@@ -16,3 +16,13 @@ LANGUAGE C STRICT;
 CREATE FOREIGN DATA WRAPPER postgres_fdw
   HANDLER postgres_fdw_handler
   VALIDATOR postgres_fdw_validator;
+
+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