Add psql variables to track success/failure of SQL queries.

tglsfdc · tglsfdc · commit 69835bc89888 · 2017-09-12T19:27:48.000-04:00
This patch adds ERROR, SQLSTATE, and ROW_COUNT, which are updated after every query, as well as LAST_ERROR_MESSAGE and LAST_ERROR_SQLSTATE, which are updated only when a query fails. The expected usage of these is for scripting. Fabien Coelho, reviewed by Pavel Stehule Discussion: https://postgr.es/m/alpine.DEB.2.20.1704042158020.12290@lancre
diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml
@@ -3517,6 +3517,16 @@ bar
         </listitem>
       </varlistentry>
 
+      <varlistentry>
+       <term><varname>ERROR</varname></term>
+       <listitem>
+        <para>
+         <literal>true</> if the last SQL query failed, <literal>false</> if
+         it succeeded.  See also <varname>SQLSTATE</>.
+        </para>
+       </listitem>
+      </varlistentry>
+
       <varlistentry>
         <term><varname>FETCH_COUNT</varname></term>
         <listitem>
@@ -3653,6 +3663,19 @@ bar
         </listitem>
       </varlistentry>
 
+      <varlistentry>
+       <term><varname>LAST_ERROR_MESSAGE</varname></term>
+       <term><varname>LAST_ERROR_SQLSTATE</varname></term>
+       <listitem>
+        <para>
+         The primary error message and associated SQLSTATE code for the most
+         recent failed query in the current <application>psql</> session, or
+         an empty string and <literal>00000</> if no error has occurred in
+         the current session.
+        </para>
+       </listitem>
+      </varlistentry>
+
       <varlistentry>
       <term>
        <varname>ON_ERROR_ROLLBACK</varname>
@@ -3732,6 +3755,16 @@ bar
         </listitem>
       </varlistentry>
 
+      <varlistentry>
+       <term><varname>ROW_COUNT</varname></term>
+       <listitem>
+        <para>
+         The number of rows returned or affected by the last SQL query, or 0
+         if the query failed or did not report a row count.
+        </para>
+       </listitem>
+      </varlistentry>
+
       <varlistentry>
         <term><varname>SERVER_VERSION_NAME</varname></term>
         <term><varname>SERVER_VERSION_NUM</varname></term>
@@ -3784,6 +3817,17 @@ bar
         </listitem>
       </varlistentry>
 
+      <varlistentry>
+       <term><varname>SQLSTATE</varname></term>
+       <listitem>
+        <para>
+         The error code (see <xref linkend="errcodes-appendix">) associated
+         with the last SQL query's failure, or <literal>00000</> if it
+         succeeded.
+        </para>
+       </listitem>
+      </varlistentry>
+
       <varlistentry>
         <term><varname>USER</varname></term>
         <listitem>
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
@@ -548,11 +548,58 @@ AcceptResult(const PGresult *result)
 }
 
 
+/*
+ * Set special variables from a query result
+ * - ERROR: true/false, whether an error occurred on this query
+ * - SQLSTATE: code of error, or "00000" if no error, or "" if unknown
+ * - ROW_COUNT: how many rows were returned or affected, or "0"
+ * - LAST_ERROR_SQLSTATE: same for last error
+ * - LAST_ERROR_MESSAGE: message of last error
+ *
+ * Note: current policy is to apply this only to the results of queries
+ * entered by the user, not queries generated by slash commands.
+ */
+static void
+SetResultVariables(PGresult *results, bool success)
+{
+	if (success)
+	{
+		const char *ntuples = PQcmdTuples(results);
+
+		SetVariable(pset.vars, "ERROR", "false");
+		SetVariable(pset.vars, "SQLSTATE", "00000");
+		SetVariable(pset.vars, "ROW_COUNT", *ntuples ? ntuples : "0");
+	}
+	else
+	{
+		const char *code = PQresultErrorField(results, PG_DIAG_SQLSTATE);
+		const char *mesg = PQresultErrorField(results, PG_DIAG_MESSAGE_PRIMARY);
+
+		SetVariable(pset.vars, "ERROR", "true");
+
+		/*
+		 * If there is no SQLSTATE code, use an empty string.  This can happen
+		 * for libpq-detected errors (e.g., lost connection, ENOMEM).
+		 */
+		if (code == NULL)
+			code = "";
+		SetVariable(pset.vars, "SQLSTATE", code);
+		SetVariable(pset.vars, "ROW_COUNT", "0");
+		SetVariable(pset.vars, "LAST_ERROR_SQLSTATE", code);
+		SetVariable(pset.vars, "LAST_ERROR_MESSAGE", mesg ? mesg : "");
+	}
+}
+
+
 /*
  * ClearOrSaveResult
  *
  * If the result represents an error, remember it for possible display by
  * \errverbose.  Otherwise, just PQclear() it.
+ *
+ * Note: current policy is to apply this to the results of all queries,
+ * including "back door" queries, for debugging's sake.  It's OK to use
+ * PQclear() directly on results known to not be error results, however.
  */
 static void
 ClearOrSaveResult(PGresult *result)
@@ -1107,6 +1154,8 @@ ProcessResult(PGresult **results)
 		first_cycle = false;
 	}
 
+	SetResultVariables(*results, success);
+
 	/* may need this to recover from conn loss during COPY */
 	if (!first_cycle && !CheckConnection())
 		return false;
@@ -1526,6 +1575,7 @@ DescribeQuery(const char *query, double *elapsed_msec)
 	if (PQresultStatus(results) != PGRES_COMMAND_OK)
 	{
 		psql_error("%s", PQerrorMessage(pset.db));
+		SetResultVariables(results, false);
 		ClearOrSaveResult(results);
 		return false;
 	}
@@ -1599,6 +1649,7 @@ DescribeQuery(const char *query, double *elapsed_msec)
 					_("The command has no result, or the result has no columns.\n"));
 	}
 
+	SetResultVariables(results, OK);
 	ClearOrSaveResult(results);
 
 	return OK;
@@ -1626,6 +1677,7 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 	bool		is_pipe;
 	bool		is_pager = false;
 	bool		started_txn = false;
+	int64		total_tuples = 0;
 	int			ntuples;
 	int			fetch_count;
 	char		fetch_cmd[64];
@@ -1663,6 +1715,8 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 	results = PQexec(pset.db, buf.data);
 	OK = AcceptResult(results) &&
 		(PQresultStatus(results) == PGRES_COMMAND_OK);
+	if (!OK)
+		SetResultVariables(results, OK);
 	ClearOrSaveResult(results);
 	termPQExpBuffer(&buf);
 	if (!OK)
@@ -1738,6 +1792,7 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 
 			OK = AcceptResult(results);
 			Assert(!OK);
+			SetResultVariables(results, OK);
 			ClearOrSaveResult(results);
 			break;
 		}
@@ -1755,6 +1810,7 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 		 */
 
 		ntuples = PQntuples(results);
+		total_tuples += ntuples;
 
 		if (ntuples < fetch_count)
 		{
@@ -1816,6 +1872,21 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 		ClosePager(fout);
 	}
 
+	if (OK)
+	{
+		/*
+		 * We don't have a PGresult here, and even if we did it wouldn't have
+		 * the right row count, so fake SetResultVariables().  In error cases,
+		 * we already set the result variables above.
+		 */
+		char		buf[32];
+
+		SetVariable(pset.vars, "ERROR", "false");
+		SetVariable(pset.vars, "SQLSTATE", "00000");
+		snprintf(buf, sizeof(buf), INT64_FORMAT, total_tuples);
+		SetVariable(pset.vars, "ROW_COUNT", buf);
+	}
+
 cleanup:
 	if (pset.timing)
 		INSTR_TIME_SET_CURRENT(before);
diff --git a/src/bin/psql/help.c b/src/bin/psql/help.c
@@ -337,7 +337,7 @@ helpVariables(unsigned short int pager)
 	 * Windows builds currently print one more line than non-Windows builds.
 	 * Using the larger number is fine.
 	 */
-	output = PageOutput(147, pager ? &(pset.popt.topt) : NULL);
+	output = PageOutput(156, pager ? &(pset.popt.topt) : NULL);
 
 	fprintf(output, _("List of specially treated variables\n\n"));
 
@@ -360,6 +360,8 @@ helpVariables(unsigned short int pager)
 					  "    if set to \"noexec\", just show them without execution\n"));
 	fprintf(output, _("  ENCODING\n"
 					  "    current client character set encoding\n"));
+	fprintf(output, _("  ERROR\n"
+					  "    true if last query failed, else false\n"));
 	fprintf(output, _("  FETCH_COUNT\n"
 					  "    the number of result rows to fetch and display at a time (0 = unlimited)\n"));
 	fprintf(output, _("  HISTCONTROL\n"
@@ -374,6 +376,9 @@ helpVariables(unsigned short int pager)
 					  "    number of EOFs needed to terminate an interactive session\n"));
 	fprintf(output, _("  LASTOID\n"
 					  "    value of the last affected OID\n"));
+	fprintf(output, _("  LAST_ERROR_MESSAGE\n"
+					  "  LAST_ERROR_SQLSTATE\n"
+					  "    message and SQLSTATE of last error, or empty string and \"00000\" if none\n"));
 	fprintf(output, _("  ON_ERROR_ROLLBACK\n"
 					  "    if set, an error doesn't stop a transaction (uses implicit savepoints)\n"));
 	fprintf(output, _("  ON_ERROR_STOP\n"
@@ -388,6 +393,8 @@ helpVariables(unsigned short int pager)
 					  "    specifies the prompt used during COPY ... FROM STDIN\n"));
 	fprintf(output, _("  QUIET\n"
 					  "    run quietly (same as -q option)\n"));
+	fprintf(output, _("  ROW_COUNT\n"
+					  "    number of rows returned or affected by last query, or 0\n"));
 	fprintf(output, _("  SERVER_VERSION_NAME\n"
 					  "  SERVER_VERSION_NUM\n"
 					  "    server's version (in short string or numeric format)\n"));
@@ -397,6 +404,8 @@ helpVariables(unsigned short int pager)
 					  "    if set, end of line terminates SQL commands (same as -S option)\n"));
 	fprintf(output, _("  SINGLESTEP\n"
 					  "    single-step mode (same as -s option)\n"));
+	fprintf(output, _("  SQLSTATE\n"
+					  "    SQLSTATE of last query, or \"00000\" if no error\n"));
 	fprintf(output, _("  USER\n"
 					  "    the currently connected database user\n"));
 	fprintf(output, _("  VERBOSITY\n"
diff --git a/src/bin/psql/startup.c b/src/bin/psql/startup.c
@@ -165,6 +165,10 @@ main(int argc, char *argv[])
 	SetVariable(pset.vars, "VERSION_NAME", PG_VERSION);
 	SetVariable(pset.vars, "VERSION_NUM", CppAsString2(PG_VERSION_NUM));
 
+	/* Initialize variables for last error */
+	SetVariable(pset.vars, "LAST_ERROR_MESSAGE", "");
+	SetVariable(pset.vars, "LAST_ERROR_SQLSTATE", "00000");
+
 	/* Default values for variables (that don't match the result of \unset) */
 	SetVariableBool(pset.vars, "AUTOCOMMIT");
 	SetVariable(pset.vars, "PROMPT1", DEFAULT_PROMPT1);
diff --git a/src/test/regress/expected/psql.out b/src/test/regress/expected/psql.out
diff --git a/src/test/regress/sql/psql.sql b/src/test/regress/sql/psql.sql
