1 files changed, 71 insertions, 0 deletions

diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index b99705886fa..9b59ee840b9 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -549,10 +549,57 @@ 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);