From 86a8078d2750a9ee42db5a341838fad6072252f5 Mon Sep 17 00:00:00 2001
From: Mark Dilger <mark.dilger@enterprisedb.com>
Date: Wed, 6 Jan 2021 13:25:29 -0800
Subject: [PATCH v32 4/9] Adding CurrentQueryHandler logic.

Extending the default set of PGresultHandlers and creating a
mechanism to switch between them using a new function
ResultHandlerSwitchTo, analogous to MemoryContextSwithTo.  In
addition to the exiting_handler already created in a prior commit
(which embeds the historical behavior from pg_dump), adding a
quiet_handler which cleans up and exits without logging anything,
and a noop_handler which does nothing, leaving the responsibility
for cleanup handling to the caller.
---
 src/fe_utils/query_utils.c         | 81 ++++++++++++++++++++++--------
 src/include/fe_utils/query_utils.h | 17 +++++++
 2 files changed, 76 insertions(+), 22 deletions(-)

diff --git a/src/fe_utils/query_utils.c b/src/fe_utils/query_utils.c
index b28750f4b2..355da6edaf 100644
--- a/src/fe_utils/query_utils.c
+++ b/src/fe_utils/query_utils.c
@@ -12,6 +12,12 @@
 #include "fe_utils/exit_utils.h"
 #include "fe_utils/query_utils.h"
 
+/*
+ * Global memory.
+ */
+
+PGresultHandler CurrentResultHandler = exiting_handler;
+
 /*
  * The exiting query result handler embeds the historical pg_dump behavior
  * under query error conditions, including exiting nicely.  The 'conn' object
@@ -29,7 +35,7 @@
  */
 PGresult *
 exiting_handler(PGresult *res, PGconn *conn, ExecStatusType expected_status,
-				  int expected_ntups, const char *query)
+				int expected_ntups, const char *query)
 {
 	if (PQresultStatus(res) != expected_status)
 	{
@@ -67,48 +73,79 @@ exiting_handler(PGresult *res, PGconn *conn, ExecStatusType expected_status,
 	return res;
 }
 
+/*
+ * Quietly cleans up and exits nicely unless the expected conditions were met.
+ */
+PGresult *
+quiet_handler(PGresult *res, PGconn *conn, ExecStatusType expected_status,
+			  int expected_ntups, const char *query)
+{
+	int ntups = PQntuples(res);
+
+	if ((PQresultStatus(res) != expected_status) ||
+		(expected_ntups == POSITIVE_NTUPS && ntups == 0) ||
+		(expected_ntups >= 0 && ntups != expected_ntups))
+	{
+		PQfinish(conn);
+		exit_nicely(1);
+	}
+
+	return res;
+}
+
+/*
+ * Does nothing other than returning the 'res' argument back to the caller.
+ * This handler is intended for callers who prefer to perform the error
+ * handling themselves.
+ */
+PGresult *
+noop_handler(PGresult *res, PGconn *conn, ExecStatusType expected_status,
+			 int expected_ntups, const char *query)
+{
+	return res;
+}
+
 /*
  * Executes the given SQL query statement.
  *
- * Invokes the exiting handler for any but PGRES_COMMAND_OK status.
+ * Expects a PGRES_COMMAND_OK status.
  */
 void
 ExecuteSqlStatement(PGconn *conn, const char *query)
 {
-	PQclear(exiting_handler(PQexec(conn, query),
-							conn,
-							PGRES_COMMAND_OK,
-							ANY_NTUPS,
-							query));
+	PQclear(CurrentResultHandler(PQexec(conn, query),
+								 conn,
+								 PGRES_COMMAND_OK,
+								 ANY_NTUPS,
+								 query));
 }
 
 /*
  * Executes the given SQL query.
  *
- * Invokes the exiting handler unless the given 'status' results.
- *
- * If successful, returns the query result.
+ * Expects the given status.
  */
 PGresult *
 ExecuteSqlQuery(PGconn *conn, const char *query, ExecStatusType status)
 {
-	return exiting_handler(PQexec(conn, query),
-						   conn,
-						   status,
-						   ANY_NTUPS,
-						   query);
+	return CurrentResultHandler(PQexec(conn, query),
+								conn,
+								status,
+								ANY_NTUPS,
+								query);
 }
 
 /*
- * Like ExecuteSqlQuery, but requires PGRES_TUPLES_OK status and
- * requires that exactly one row be returned.
+ * Executes the given SQL query.
+ *
+ * Expects a PGRES_TUPLES_OK status and precisely one row.
  */
 PGresult *
 ExecuteSqlQueryForSingleRow(PGconn *conn, const char *query)
 {
-	return exiting_handler(PQexec(conn, query),
-						   conn,
-						   PGRES_TUPLES_OK,
-						   1,
-						   query);
+	return CurrentResultHandler(PQexec(conn, query),
+								conn,
+								PGRES_TUPLES_OK,
+								1,
+								query);
 }
diff --git a/src/include/fe_utils/query_utils.h b/src/include/fe_utils/query_utils.h
index f03d17b1ed..80958e94fb 100644
--- a/src/include/fe_utils/query_utils.h
+++ b/src/include/fe_utils/query_utils.h
@@ -22,9 +22,26 @@ typedef PGresult *(*PGresultHandler) (PGresult *res,
 									  int expected_ntups,
 									  const char *query);
 
+extern PGresultHandler CurrentResultHandler;
+
+static inline PGresultHandler
+ResultHandlerSwitchTo(PGresultHandler handler)
+{
+	PGresultHandler old = CurrentResultHandler;
+
+	CurrentResultHandler = handler;
+	return old;
+}
+
 extern PGresult *exiting_handler(PGresult *res, PGconn *conn,
 								 ExecStatusType expected_status,
 								 int expected_ntups, const char *query);
+extern PGresult *quiet_handler(PGresult *res, PGconn *conn,
+							   ExecStatusType expected_status,
+							   int expected_ntups, const char *query);
+extern PGresult *noop_handler(PGresult *res, PGconn *conn,
+							  ExecStatusType expected_status,
+							  int expected_ntups, const char *query);
 
 extern void ExecuteSqlStatement(PGconn *conn, const char *query);
 extern PGresult *ExecuteSqlQuery(PGconn *conn, const char *query,
-- 
2.21.1 (Apple Git-122.3)