From 45d0a4bbbcc8a3c0e8d377a1c2b6240056a4479d Mon Sep 17 00:00:00 2001
From: Sagar Shedge <sagar.shedge92@gmail.com>
Date: Mon, 24 Feb 2025 14:44:34 +0000
Subject: [PATCH v1] Extend postgres_fdw_get_connections to return remote
 backend pid

This commit adds a "remote_backend_pid" output column to
the postgres_fdw_get_connections function, returning the process ID of
remote backend handling connection. If connection marked as invalid then
API returns process ID of remote backend which handled this connection.

Postgres backend (Coordinator) serving global session can establish connections
with multiple backends (worker). Using coordinator backend PID and worker backend
PID information, users can build a global session/transaction view and also use
this information for monitoring and reporting purpose.
---
 contrib/postgres_fdw/connection.c             | 20 ++++++++++--
 .../postgres_fdw/expected/postgres_fdw.out    | 32 +++++++++++--------
 .../postgres_fdw/postgres_fdw--1.1--1.2.sql   |  2 +-
 contrib/postgres_fdw/sql/postgres_fdw.sql     | 12 ++++---
 doc/src/sgml/postgres-fdw.sgml                | 22 +++++++++----
 5 files changed, 60 insertions(+), 28 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 8a8d3b4481f..0eb55bdfeb9 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -2111,8 +2111,8 @@ pgfdw_finish_abort_cleanup(List *pending_entries, List *cancel_requested,
 
 /* Number of output arguments (columns) for various API versions */
 #define POSTGRES_FDW_GET_CONNECTIONS_COLS_V1_1	2
-#define POSTGRES_FDW_GET_CONNECTIONS_COLS_V1_2	5
-#define POSTGRES_FDW_GET_CONNECTIONS_COLS	5	/* maximum of above */
+#define POSTGRES_FDW_GET_CONNECTIONS_COLS_V1_2	6
+#define POSTGRES_FDW_GET_CONNECTIONS_COLS	6	/* maximum of above */
 
 /*
  * Internal function used by postgres_fdw_get_connections variants.
@@ -2128,13 +2128,16 @@ pgfdw_finish_abort_cleanup(List *pending_entries, List *cancel_requested,
  *
  * For API version 1.2 and later, this function takes an input parameter
  * to check a connection status and returns the following
- * additional values along with the three values from version 1.1:
+ * additional values along with the four values from version 1.1:
  *
  * - user_name - the local user name of the active connection. In case the
  *   user mapping is dropped but the connection is still active, then the
  *   user name will be NULL in the output.
  * - used_in_xact - true if the connection is used in the current transaction.
  * - closed - true if the connection is closed.
+ * - remote_backend_pid - process ID of remote backend handling connection. In
+ *   case the connection status is not OK or no backend assigned to this
+ *   connection then remote backend pid will be NULL in the output.
  *
  * No records are returned when there are no cached connections at all.
  */
@@ -2261,6 +2264,7 @@ postgres_fdw_get_connections_internal(FunctionCallInfo fcinfo,
 		if (api_version >= PGFDW_V1_2)
 		{
 			bool		check_conn = PG_GETARG_BOOL(0);
+			int 		remote_backend_pid;
 
 			/* Is this connection used in the current transaction? */
 			values[i++] = BoolGetDatum(entry->xact_depth > 0);
@@ -2273,6 +2277,16 @@ postgres_fdw_get_connections_internal(FunctionCallInfo fcinfo,
 				values[i++] = BoolGetDatum(pgfdw_conn_check(entry->conn) != 0);
 			else
 				nulls[i++] = true;
+
+			/*
+			 * Return non-zero remote backend ID. Otherwise, return NULL to indicate
+			 * connection status is not OK or no backend assigned to this connection.
+			 */
+			remote_backend_pid = PQbackendPID(entry->conn);
+			if (remote_backend_pid != 0)
+				values[i++] = Int32GetDatum(remote_backend_pid);
+			else
+				nulls[i++] = true;
 		}
 
 		tuplestore_putvalues(rsinfo->setResult, rsinfo->setDesc, values, nulls);
diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index daa3b1d7a6d..5d4188407f4 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -10589,13 +10589,15 @@ drop cascades to foreign table ft7
 -- List all the existing cached connections. loopback and loopback3
 -- should be output as invalid connections. Also the server name and user name
 -- for loopback3 should be NULL because both server and user mapping were
--- dropped.
-SELECT server_name, user_name = CURRENT_USER as "user_name = CURRENT_USER", valid, used_in_xact, closed
+-- dropped. remote_backend_pid will continue to return available as it fetch remote
+-- server backend pid from cached connections.
+SELECT server_name, user_name = CURRENT_USER as "user_name = CURRENT_USER", valid, used_in_xact, closed,
+CASE WHEN remote_backend_pid = ANY(SELECT pid FROM pg_stat_activity WHERE backend_type = 'client backend' AND pid <> pg_backend_pid()) then 'available' ELSE 'not available' END AS remote_backend_pid
 FROM postgres_fdw_get_connections() ORDER BY 1;
- server_name | user_name = CURRENT_USER | valid | used_in_xact | closed 
--------------+--------------------------+-------+--------------+--------
- loopback    | t                        | f     | t            | 
-             |                          | f     | t            | 
+ server_name | user_name = CURRENT_USER | valid | used_in_xact | closed | remote_backend_pid 
+-------------+--------------------------+-------+--------------+--------+--------------------
+ loopback    | t                        | f     | t            |        | available
+             |                          | f     | t            |        | available
 (2 rows)
 
 -- The invalid connections get closed in pgfdw_xact_callback during commit.
@@ -12436,11 +12438,12 @@ SELECT 1 FROM ft1 LIMIT 1;
 
 -- Since the remote server is still connected, "closed" should be FALSE,
 -- or NULL if the connection status check is not available.
-SELECT CASE WHEN closed IS NOT true THEN 1 ELSE 0 END
+SELECT server_name, CASE WHEN closed IS NOT true THEN 'false' ELSE 'true' END AS remote_conn_closed,
+CASE WHEN remote_backend_pid = (SELECT pid FROM pg_stat_activity WHERE application_name = 'fdw_conn_check') THEN 'available' ELSE 'not available' END AS remote_backend_pid
   FROM postgres_fdw_get_connections(true);
- case 
-------
-    1
+ server_name | remote_conn_closed | remote_backend_pid 
+-------------+--------------------+--------------------
+ loopback    | false              | available
 (1 row)
 
 -- After terminating the remote backend, since the connection is closed,
@@ -12450,11 +12453,12 @@ DO $$ BEGIN
 PERFORM pg_terminate_backend(pid, 180000) FROM pg_stat_activity
   WHERE application_name = 'fdw_conn_check';
 END $$;
-SELECT CASE WHEN closed IS NOT false THEN 1 ELSE 0 END
+SELECT server_name, CASE WHEN closed IS NOT false THEN 'true' ELSE 'false' END AS remote_conn_closed,
+CASE WHEN remote_backend_pid IS NOT NULL THEN 'available' ELSE 'not available' END AS remote_backend_pid
   FROM postgres_fdw_get_connections(true);
- case 
-------
-    1
+ server_name | remote_conn_closed | remote_backend_pid 
+-------------+--------------------+--------------------
+ loopback    | true               | available
 (1 row)
 
 -- Clean up
diff --git a/contrib/postgres_fdw/postgres_fdw--1.1--1.2.sql b/contrib/postgres_fdw/postgres_fdw--1.1--1.2.sql
index 81aad4fcdaa..720401b3c21 100644
--- a/contrib/postgres_fdw/postgres_fdw--1.1--1.2.sql
+++ b/contrib/postgres_fdw/postgres_fdw--1.1--1.2.sql
@@ -12,7 +12,7 @@ DROP FUNCTION postgres_fdw_get_connections ();
 CREATE FUNCTION postgres_fdw_get_connections (
     IN check_conn boolean DEFAULT false, OUT server_name text,
     OUT user_name text, OUT valid boolean, OUT used_in_xact boolean,
-    OUT closed boolean)
+    OUT closed boolean, OUT remote_backend_pid INT4)
 RETURNS SETOF record
 AS 'MODULE_PATHNAME', 'postgres_fdw_get_connections_1_2'
 LANGUAGE C STRICT PARALLEL RESTRICTED;
diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index 1598d9e0862..cee12c0f92e 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -3410,8 +3410,10 @@ DROP SERVER loopback3 CASCADE;
 -- List all the existing cached connections. loopback and loopback3
 -- should be output as invalid connections. Also the server name and user name
 -- for loopback3 should be NULL because both server and user mapping were
--- dropped.
-SELECT server_name, user_name = CURRENT_USER as "user_name = CURRENT_USER", valid, used_in_xact, closed
+-- dropped. remote_backend_pid will continue to return available as it fetch remote
+-- server backend pid from cached connections.
+SELECT server_name, user_name = CURRENT_USER as "user_name = CURRENT_USER", valid, used_in_xact, closed,
+CASE WHEN remote_backend_pid = ANY(SELECT pid FROM pg_stat_activity WHERE backend_type = 'client backend' AND pid <> pg_backend_pid()) then 'available' ELSE 'not available' END AS remote_backend_pid
 FROM postgres_fdw_get_connections() ORDER BY 1;
 -- The invalid connections get closed in pgfdw_xact_callback during commit.
 COMMIT;
@@ -4288,7 +4290,8 @@ SELECT 1 FROM ft1 LIMIT 1;
 
 -- Since the remote server is still connected, "closed" should be FALSE,
 -- or NULL if the connection status check is not available.
-SELECT CASE WHEN closed IS NOT true THEN 1 ELSE 0 END
+SELECT server_name, CASE WHEN closed IS NOT true THEN 'false' ELSE 'true' END AS remote_conn_closed,
+CASE WHEN remote_backend_pid = (SELECT pid FROM pg_stat_activity WHERE application_name = 'fdw_conn_check') THEN 'available' ELSE 'not available' END AS remote_backend_pid
   FROM postgres_fdw_get_connections(true);
 
 -- After terminating the remote backend, since the connection is closed,
@@ -4298,7 +4301,8 @@ DO $$ BEGIN
 PERFORM pg_terminate_backend(pid, 180000) FROM pg_stat_activity
   WHERE application_name = 'fdw_conn_check';
 END $$;
-SELECT CASE WHEN closed IS NOT false THEN 1 ELSE 0 END
+SELECT server_name, CASE WHEN closed IS NOT false THEN 'true' ELSE 'false' END AS remote_conn_closed,
+CASE WHEN remote_backend_pid IS NOT NULL THEN 'available' ELSE 'not available' END AS remote_backend_pid
   FROM postgres_fdw_get_connections(true);
 
 -- Clean up
diff --git a/doc/src/sgml/postgres-fdw.sgml b/doc/src/sgml/postgres-fdw.sgml
index d2998c13d5d..ff1ec8fb121 100644
--- a/doc/src/sgml/postgres-fdw.sgml
+++ b/doc/src/sgml/postgres-fdw.sgml
@@ -854,7 +854,7 @@ OPTIONS (ADD password_required 'false');
     <term><function>postgres_fdw_get_connections(
       IN check_conn boolean DEFAULT false, OUT server_name text,
       OUT user_name text, OUT valid boolean, OUT used_in_xact boolean,
-      OUT closed boolean)
+      OUT closed boolean, OUT remote_backend_pid INT4)
       returns setof record</function></term>
     <listitem>
      <para>
@@ -882,11 +882,11 @@ OPTIONS (ADD password_required 'false');
       Example usage of the function:
 <screen>
 postgres=# SELECT * FROM postgres_fdw_get_connections(true);
- server_name | user_name | valid | used_in_xact | closed
--------------+-----------+-------+--------------+--------
- loopback1   | postgres  | t     | t            | f
- loopback2   | public    | t     | t            | f
- loopback3   |           | f     | t            | f
+ server_name | user_name | valid | used_in_xact | closed | remote_backend_pid
+-------------+-----------+-------+--------------+-----------------------------
+ loopback1   | postgres  | t     | t            | f      |            1353340
+ loopback2   | public    | t     | t            | f      |            1353120
+ loopback3   |           | f     | t            | f      |            1353156
 </screen>
       The output columns are described in
       <xref linkend="postgres-fdw-get-connections-columns"/>.
@@ -951,6 +951,16 @@ postgres=# SELECT * FROM postgres_fdw_get_connections(true);
          is not available on this platform.
         </entry>
        </row>
+       <row>
+        <entry><structfield>remote_backend_pid</structfield></entry>
+        <entry><type>int4</type></entry>
+        <entry>
+         Process ID of remote backend handling connection. If <literal>valid</literal>
+         is set to <literal>false</literal> (i.e., marked as invalid), this will be
+         process ID of remote backend which handled this connection. <literal>NULL</literal>
+         is returned if no backend assigned to this connection or connection status is bad.
+        </entry>
+       </row>
       </tbody>
      </tgroup>
     </table>
-- 
2.43.0

