From 443be6f64ca89bbd6367a011f2a98aa9324f27a9 Mon Sep 17 00:00:00 2001
From: coreyhuinker <corey.huinker@gmail.com>
Date: Wed, 21 Dec 2022 00:24:47 -0500
Subject: [PATCH 1/2] Make the exit code of shell commands executed via psql
 visible via the variable SHELL_EXIT_CODE.

---
 doc/src/sgml/ref/psql-ref.sgml |  9 +++++++++
 src/bin/psql/command.c         |  4 ++++
 src/bin/psql/help.c            |  2 ++
 src/bin/psql/psqlscanslash.l   | 24 +++++++++++++++++++++---
 4 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml
index 8a5285da9a..d0c80b4528 100644
--- a/doc/src/sgml/ref/psql-ref.sgml
+++ b/doc/src/sgml/ref/psql-ref.sgml
@@ -4260,6 +4260,15 @@ bar
         </listitem>
       </varlistentry>
 
+      <varlistentry>
+       <term><varname>SHELL_EXIT_CODE</varname></term>
+       <listitem>
+        <para>
+         The exit code return by the last shell command. 0 means no error.
+        </para>
+       </listitem>
+      </varlistentry>
+
       <varlistentry>
         <term><varname>SHOW_ALL_RESULTS</varname></term>
         <listitem>
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index de6a3a71f8..f6d6a489a9 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -4998,6 +4998,7 @@ static bool
 do_shell(const char *command)
 {
 	int			result;
+	char		exit_code_buf[32];
 
 	fflush(NULL);
 	if (!command)
@@ -5025,6 +5026,9 @@ do_shell(const char *command)
 	else
 		result = system(command);
 
+	snprintf(exit_code_buf, sizeof(exit_code_buf), "%d", WEXITSTATUS(result));
+	SetVariable(pset.vars, "SHELL_EXIT_CODE", exit_code_buf);
+
 	if (result == 127 || result == -1)
 	{
 		pg_log_error("\\!: failed");
diff --git a/src/bin/psql/help.c b/src/bin/psql/help.c
index b4e0ec2687..caf13e2ed2 100644
--- a/src/bin/psql/help.c
+++ b/src/bin/psql/help.c
@@ -455,6 +455,8 @@ helpVariables(unsigned short int pager)
 		  "    show all results of a combined query (\\;) instead of only the last\n");
 	HELP0("  SHOW_CONTEXT\n"
 		  "    controls display of message context fields [never, errors, always]\n");
+	HELP0("  SHELL_EXIT_CODE\n"
+		  "    Exit code of the last shell command\n");
 	HELP0("  SINGLELINE\n"
 		  "    if set, end of line terminates SQL commands (same as -S option)\n");
 	HELP0("  SINGLESTEP\n"
diff --git a/src/bin/psql/psqlscanslash.l b/src/bin/psql/psqlscanslash.l
index a467b72144..ecbc3b9d8b 100644
--- a/src/bin/psql/psqlscanslash.l
+++ b/src/bin/psql/psqlscanslash.l
@@ -27,6 +27,8 @@
 
 %{
 #include "fe_utils/psqlscan_int.h"
+#include "settings.h"
+#include "variables.h"
 
 /*
  * We must have a typedef YYSTYPE for yylex's first argument, but this lexer
@@ -774,6 +776,8 @@ evaluate_backtick(PsqlScanState state)
 	bool		error = false;
 	char		buf[512];
 	size_t		result;
+	int			exit_code = 0;
+	char		exit_code_buf[32];
 
 	initPQExpBuffer(&cmd_output);
 
@@ -783,6 +787,7 @@ evaluate_backtick(PsqlScanState state)
 	{
 		pg_log_error("%s: %m", cmd);
 		error = true;
+		exit_code = -1;
 	}
 
 	if (!error)
@@ -800,10 +805,21 @@ evaluate_backtick(PsqlScanState state)
 		} while (!feof(fd));
 	}
 
-	if (fd && pclose(fd) == -1)
+	if (fd)
 	{
-		pg_log_error("%s: %m", cmd);
-		error = true;
+		/* Capture exit code for SHELL_EXIT_CODE */
+		exit_code = pclose(fd);
+		if (exit_code == -1)
+		{
+			pg_log_error("%s: %m", cmd);
+			error = true;
+		}
+		if (WIFEXITED(exit_code))
+			exit_code=WEXITSTATUS(exit_code);
+		else if(WIFSIGNALED(exit_code))
+			exit_code=WTERMSIG(exit_code);
+		else if(WIFSTOPPED(exit_code))
+			exit_code=WSTOPSIG(exit_code);
 	}
 
 	if (PQExpBufferDataBroken(cmd_output))
@@ -826,5 +842,7 @@ evaluate_backtick(PsqlScanState state)
 		appendBinaryPQExpBuffer(output_buf, cmd_output.data, cmd_output.len);
 	}
 
+	snprintf(exit_code_buf, sizeof(exit_code_buf), "%d", exit_code);
+	SetVariable(pset.vars, "SHELL_EXIT_CODE", exit_code_buf);
 	termPQExpBuffer(&cmd_output);
 }
-- 
2.38.1