v9-0002-Allow-SIGINT-to-cancel-psql-database-reconnection.patch

From 779e642df8013b3fc8c63f9faf370cbbf2950392 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Tue, 30 Jan 2024 15:53:10 -0600
Subject: [PATCH v9 2/2] Allow SIGINT to cancel psql database reconnections

After installing the SIGINT handler in psql, SIGINT can no longer cancel
database reconnections. For instance, if the user starts a reconnection
and then needs to do some form of interaction (ie psql is polling),
there is no way to cancel the reconnection process currently.

Use PQconnectStartParams() in order to insert a cancel_pressed check
into the polling loop.
---
 src/bin/psql/command.c | 50 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 49 insertions(+), 1 deletion(-)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 9b0fa041f73..00a20e47328 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -159,6 +159,7 @@ static void discard_query_text(PsqlScanState scan_state, ConditionalStack cstack
 static bool copy_previous_query(PQExpBuffer query_buf, PQExpBuffer previous_buf);
 static bool do_connect(enum trivalue reuse_previous_specification,
 					   char *dbname, char *user, char *host, char *port);
+static void wait_until_connected(PGconn *conn);
 static bool do_edit(const char *filename_arg, PQExpBuffer query_buf,
 					int lineno, bool discard_on_quit, bool *edited);
 static bool do_shell(const char *command);
@@ -3595,11 +3596,12 @@ do_connect(enum trivalue reuse_previous_specification,
 		values[paramnum] = NULL;
 
 		/* Note we do not want libpq to re-expand the dbname parameter */
-		n_conn = PQconnectdbParams(keywords, values, false);
+		n_conn = PQconnectStartParams(keywords, values, false);
 
 		pg_free(keywords);
 		pg_free(values);
 
+		wait_until_connected(n_conn);
 		if (PQstatus(n_conn) == CONNECTION_OK)
 			break;
 
@@ -3748,6 +3750,52 @@ do_connect(enum trivalue reuse_previous_specification,
 	return true;
 }
 
+static void
+wait_until_connected(PGconn *conn)
+{
+	bool		for_read = false;
+
+	while (true)
+	{
+		int			rc;
+		int			sock;
+		time_t		timeout;
+
+		/*
+		 * On every iteration of the state machine, let's check if the user has
+		 * requested a cancellation of the connection process.
+		 */
+		if (cancel_pressed)
+			break;
+
+		sock = PQsocket(conn);
+		if (sock == -1)
+			break;
+
+		/* Use a 1 second timeout */
+		timeout = time(NULL) + 1;
+		rc = PQsocketPoll(sock, for_read, !for_read, timeout);
+		if (rc == -1)
+			return;
+
+		switch (PQconnectPoll(conn))
+		{
+			case PGRES_POLLING_OK:
+			case PGRES_POLLING_FAILED:
+				return;
+			case PGRES_POLLING_READING:
+				for_read = true;
+				continue;
+			case PGRES_POLLING_WRITING:
+				for_read = false;
+				continue;
+			case PGRES_POLLING_ACTIVE:
+				pg_unreachable();
+		}
+	}
+
+	pg_unreachable();
+}
 
 void
 connection_warnings(bool in_startup)
-- 
Tristan Partin
Neon (https://neon.tech)