From fef0e52706c9136659f0f8846bae289f0dbfb469 Mon Sep 17 00:00:00 2001 From: coreyhuinker Date: Fri, 4 Nov 2022 04:45:39 -0400 Subject: [PATCH] POC: expose shell exit code as a psql variable --- src/bin/psql/command.c | 4 ++++ src/bin/psql/help.c | 2 ++ src/bin/psql/psqlscanslash.l | 28 +++++++++++++++++++++++++--- src/test/regress/expected/psql.out | 11 +++++++++++ src/test/regress/sql/psql.sql | 7 +++++++ 5 files changed, 49 insertions(+), 3 deletions(-) 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..30e6f5dcd4 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,25 @@ evaluate_backtick(PsqlScanState state) } while (!feof(fd)); } - if (fd && pclose(fd) == -1) + if (fd) { - pg_log_error("%s: %m", cmd); - error = true; + 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)) { + //If you need to act upon the process stopping, do it here. + exit_code=WSTOPSIG(exit_code); + } } if (PQExpBufferDataBroken(cmd_output)) @@ -826,5 +846,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); } diff --git a/src/test/regress/expected/psql.out b/src/test/regress/expected/psql.out index b4cb6ffb5b..06aed04327 100644 --- a/src/test/regress/expected/psql.out +++ b/src/test/regress/expected/psql.out @@ -1306,6 +1306,17 @@ execute q; +----+-------------+ deallocate q; +-- test SHELL_EXIT_CODE +\! nosuchcommand +sh: line 1: nosuchcommand: command not found +\echo :SHELL_EXIT_CODE +127 +\set nosuchvar `nosuchcommand` +sh: line 1: nosuchcommand: command not found +\! nosuchcommand +sh: line 1: nosuchcommand: command not found +\echo :SHELL_EXIT_CODE +127 -- test single-line header and data prepare q as select repeat('x',2*n) as "0123456789abcdef", repeat('y',20-2*n) as "0123456789" from generate_series(1,10) as n; \pset linestyle ascii diff --git a/src/test/regress/sql/psql.sql b/src/test/regress/sql/psql.sql index 2da9665a19..623e7a0538 100644 --- a/src/test/regress/sql/psql.sql +++ b/src/test/regress/sql/psql.sql @@ -291,6 +291,13 @@ execute q; deallocate q; +-- test SHELL_EXIT_CODE +\! nosuchcommand +\echo :SHELL_EXIT_CODE +\set nosuchvar `nosuchcommand` +\! nosuchcommand +\echo :SHELL_EXIT_CODE + -- test single-line header and data prepare q as select repeat('x',2*n) as "0123456789abcdef", repeat('y',20-2*n) as "0123456789" from generate_series(1,10) as n; -- 2.38.1