From f493be68e085fec4a7fb9518bb53d3e6da2acd18 Mon Sep 17 00:00:00 2001
From: Srinath Reddy Sadipiralla <srinath.reddy@zohocorp.com>
Date: Fri, 8 Nov 2024 23:53:46 +0530
Subject: [PATCH] Kill psql process with single process implementation

---
 src/bin/psql/input.c | 105 ++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 103 insertions(+), 2 deletions(-)

diff --git a/src/bin/psql/input.c b/src/bin/psql/input.c
index 01b7ef74c3..0f50e9c866 100644
--- a/src/bin/psql/input.c
+++ b/src/bin/psql/input.c
@@ -18,6 +18,8 @@
 #include "input.h"
 #include "settings.h"
 #include "tab-complete.h"
+#include <pthread.h>
+#include <sys/socket.h>
 
 #ifndef WIN32
 #define PSQLHISTORY ".psql_history"
@@ -47,8 +49,107 @@ static int	history_lines_added;
 #define NL_IN_HISTORY	0x01
 #endif
 
+static char *line_buffer = NULL;
+static bool received_line = false;
+char *full_query_buffer = NULL;
+
 static void finishInput(void);
 
+static void handle_line(char *line)
+{
+	if (line == NULL)
+	{
+		rl_callback_handler_remove();
+		exit(0);
+	}
+
+	if (line_buffer != NULL)
+	{
+		free(line_buffer);
+		line_buffer = NULL;
+	}
+
+	if (full_query_buffer == NULL)
+	{
+		full_query_buffer = strdup(line);
+	}
+	else
+	{
+		char *temp = malloc(strlen(full_query_buffer) + strlen(line) + 2);
+		sprintf(temp, "%s\n%s", full_query_buffer, line);
+		free(full_query_buffer);
+		full_query_buffer = temp;
+	}
+
+	if (full_query_buffer[strlen(full_query_buffer) - 1] == ';')
+	{
+		line_buffer = strdup(full_query_buffer);
+		free(full_query_buffer);
+		full_query_buffer = NULL;
+		received_line = true;
+	}
+	else
+	{
+		line_buffer = strdup(line);
+		received_line = true;
+	}
+	free(line);
+	line = NULL;
+}
+
+static char *monitor_stdin_psql_fds(const char *prompt)
+{
+	fd_set read_fds;
+	int max_fd;
+	struct timeval timeout = {};
+	int retval;
+
+	pgsocket psql_fd = PQsocket(pset.db);
+	int terminal_fd = fileno(stdin);
+
+	rl_callback_handler_install(prompt, handle_line);
+
+	while (1)
+	{
+		FD_ZERO(&read_fds);
+		FD_SET(terminal_fd, &read_fds);
+		FD_SET(psql_fd, &read_fds);
+		max_fd = (terminal_fd > psql_fd) ? terminal_fd : psql_fd;
+
+		retval = select(max_fd + 1, &read_fds, NULL, NULL, &timeout);
+
+		if (retval == -1 && errno != EINTR && errno != EAGAIN)
+		{
+			rl_callback_handler_remove();
+			return NULL;
+		}
+
+		if (FD_ISSET(terminal_fd, &read_fds))
+		{
+
+			rl_callback_read_char();
+			if (received_line)
+			{
+				char *result = line_buffer;
+				line_buffer = NULL;
+				received_line = false;
+				return result;
+			}
+		}
+
+		if (FD_ISSET(psql_fd, &read_fds))
+		{
+			char buf;
+			int len = recv(psql_fd, &buf, 1, 0);
+			if (len <= 0 && errno != EINTR && errno != EAGAIN)
+			{
+				rl_callback_handler_remove();
+				/* If recv fails, we consider it as a trigger to terminate the process */
+				return NULL;
+			}
+		}
+	}
+}
 
 /*
  * gets_interactive()
@@ -88,8 +189,8 @@ gets_interactive(const char *prompt, PQExpBuffer query_buf)
 		/* Enable SIGINT to longjmp to sigint_interrupt_jmp */
 		sigint_interrupt_enabled = true;
 
-		result = readline(prompt);
-
+		result = monitor_stdin_psql_fds(prompt);
+		
 		/* Disable SIGINT again */
 		sigint_interrupt_enabled = false;
 
-- 
2.39.3 (Apple Git-146)