psql: set SHELL_ERROR and SHELL_EXIT_CODE in more places.

Make the \g, \o, \w, and \copy commands set these variables
when closing a pipe.  We missed doing this in commit b0d8f2d98,
but it seems like a good idea.

There are some remaining places in psql that intentionally don't
update these variables after running a child program:
	* pager invocations
	* backtick evaluation within a prompt
	* \e (edit query buffer)

Corey Huinker and Tom Lane

Discussion: https://postgr.es/m/CADkLM=eSKwRGF-rnRqhtBORRtL49QsjcVUCa-kLxKTqxypsakw@mail.gmail.com

6 files changed, 44 insertions, 27 deletions

diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml
index 53875afbf0..f4f25d1b07 100644
--- a/doc/src/sgml/ref/psql-ref.sgml
+++ b/doc/src/sgml/ref/psql-ref.sgml
@@ -4277,8 +4277,12 @@ bar
         <para>
          <literal>true</literal> if the last shell command
          failed, <literal>false</literal> if it succeeded.
-         This applies to shell commands invoked via the <literal>\!</literal>
-         meta-command or backquote (<literal>`</literal>) expansion.
+         This applies to shell commands invoked via the <literal>\!</literal>,
+         <literal>\g</literal>, <literal>\o</literal>, <literal>\w</literal>,
+         and <literal>\copy</literal> meta-commands, as well as backquote
+         (<literal>`</literal>) expansion.  Note that
+         for <literal>\o</literal>, this variable is updated when the output
+         pipe is closed by the next <literal>\o</literal> command.
          See also <varname>SHELL_EXIT_CODE</varname>.
         </para>
        </listitem>
@@ -4292,8 +4296,12 @@ bar
          0&ndash;127 represent program exit codes, 128&ndash;255
          indicate termination by a signal, and -1 indicates failure
          to launch a program or to collect its exit status.
-         This applies to shell commands invoked via the <literal>\!</literal>
-         meta-command or backquote (<literal>`</literal>) expansion.
+         This applies to shell commands invoked via the <literal>\!</literal>,
+         <literal>\g</literal>, <literal>\o</literal>, <literal>\w</literal>,
+         and <literal>\copy</literal> meta-commands, as well as backquote
+         (<literal>`</literal>) expansion.  Note that
+         for <literal>\o</literal>, this variable is updated when the output
+         pipe is closed by the next <literal>\o</literal> command.
          See also <varname>SHELL_ERROR</varname>.
         </para>
        </listitem>
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index e8f583cac2..97f7d97220 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -2734,6 +2734,7 @@ exec_command_write(PsqlScanState scan_state, bool active_branch,
 					pg_log_error("%s: %s", fname, wait_result_to_str(result));
 					status = PSQL_CMD_ERROR;
 				}
+				SetShellResultVariables(result);
 			}
 			else
 			{
@@ -5119,20 +5120,7 @@ do_shell(const char *command)
 	else
 		result = system(command);
 
-	if (result == 0)
-	{
-		SetVariable(pset.vars, "SHELL_EXIT_CODE", "0");
-		SetVariable(pset.vars, "SHELL_ERROR", "false");
-	}
-	else
-	{
-		int			exit_code = wait_result_to_exit_code(result);
-		char		buf[32];
-
-		snprintf(buf, sizeof(buf), "%d", exit_code);
-		SetVariable(pset.vars, "SHELL_EXIT_CODE", buf);
-		SetVariable(pset.vars, "SHELL_ERROR", "true");
-	}
+	SetShellResultVariables(result);
 
 	if (result == 127 || result == -1)
 	{
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index f907f5d4e8..c0e6e8e6ed 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -103,7 +103,7 @@ setQFout(const char *fname)
 	if (pset.queryFout && pset.queryFout != stdout && pset.queryFout != stderr)
 	{
 		if (pset.queryFoutPipe)
-			pclose(pset.queryFout);
+			SetShellResultVariables(pclose(pset.queryFout));
 		else
 			fclose(pset.queryFout);
 	}
@@ -450,6 +450,26 @@ SetResultVariables(PGresult *result, bool success)
 
 
 /*
+ * Set special variables from a shell command result
+ * - SHELL_ERROR: true/false, whether command returned exit code 0
+ * - SHELL_EXIT_CODE: exit code according to shell conventions
+ *
+ * The argument is a wait status as returned by wait(2) or waitpid(2),
+ * which also applies to pclose(3) and system(3).
+ */
+void
+SetShellResultVariables(int wait_result)
+{
+	char		buf[32];
+
+	SetVariable(pset.vars, "SHELL_ERROR",
+				(wait_result == 0) ? "false" : "true");
+	snprintf(buf, sizeof(buf), "%d", wait_result_to_exit_code(wait_result));
+	SetVariable(pset.vars, "SHELL_EXIT_CODE", buf);
+}
+
+
+/*
  * ClearOrSaveResult
  *
  * If the result represents an error, remember it for possible display by
@@ -1652,7 +1672,7 @@ ExecQueryAndProcessResults(const char *query,
 	{
 		if (gfile_is_pipe)
 		{
-			pclose(gfile_fout);
+			SetShellResultVariables(pclose(gfile_fout));
 			restore_sigpipe_trap();
 		}
 		else
@@ -1870,7 +1890,7 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 		/* close \g argument file/pipe */
 		if (is_pipe)
 		{
-			pclose(fout);
+			SetShellResultVariables(pclose(fout));
 			restore_sigpipe_trap();
 		}
 		else
diff --git a/src/bin/psql/common.h b/src/bin/psql/common.h
index cc4c73d097..812b94a977 100644
--- a/src/bin/psql/common.h
+++ b/src/bin/psql/common.h
@@ -29,6 +29,8 @@ extern sigjmp_buf sigint_interrupt_jmp;
 
 extern void psql_setup_cancel_handler(void);
 
+extern void SetShellResultVariables(int wait_result);
+
 extern PGresult *PSQLexec(const char *query);
 extern int	PSQLexecWatch(const char *query, const printQueryOpt *opt, FILE *printQueryFout);
 
diff --git a/src/bin/psql/copy.c b/src/bin/psql/copy.c
index 110d2d7269..b3cc3d9a29 100644
--- a/src/bin/psql/copy.c
+++ b/src/bin/psql/copy.c
@@ -391,6 +391,7 @@ do_copy(const char *args)
 				}
 				success = false;
 			}
+			SetShellResultVariables(pclose_rc);
 			restore_sigpipe_trap();
 		}
 		else
diff --git a/src/bin/psql/psqlscanslash.l b/src/bin/psql/psqlscanslash.l
index fa2d0b2a5f..5c020f30b9 100644
--- a/src/bin/psql/psqlscanslash.l
+++ b/src/bin/psql/psqlscanslash.l
@@ -18,8 +18,8 @@
  */
 #include "postgres_fe.h"
 
+#include "common.h"
 #include "psqlscanslash.h"
-#include "settings.h"
 
 #include "common/logging.h"
 #include "fe_utils/conditional.h"
@@ -807,7 +807,7 @@ evaluate_backtick(PsqlScanState state)
 	if (fd)
 	{
 		/*
-		 * Although pclose's result always sets SHELL_EXIT_CODE, we
+		 * Although pclose's result always sets the shell result variables, we
 		 * historically have abandoned the backtick substitution only if it
 		 * returns -1.
 		 */
@@ -839,10 +839,8 @@ evaluate_backtick(PsqlScanState state)
 		appendBinaryPQExpBuffer(output_buf, cmd_output.data, cmd_output.len);
 	}
 
-	/* And finally, set the shell error variables */
-	snprintf(buf, sizeof(buf), "%d", wait_result_to_exit_code(exit_code));
-	SetVariable(pset.vars, "SHELL_EXIT_CODE", buf);
-	SetVariable(pset.vars, "SHELL_ERROR", (exit_code == 0) ? "false" : "true");
+	/* And finally, set the shell result variables */
+	SetShellResultVariables(exit_code);
 
 	termPQExpBuffer(&cmd_output);
 }