0001-WIP-PQexecMulti.patch

From 947d9d98507d6c93b547ac17fef41c1870c0d577 Mon Sep 17 00:00:00 2001
From: Peter Eisentraut <peter@eisentraut.org>
Date: Thu, 27 Jan 2022 14:09:52 +0100
Subject: [PATCH] WIP: PQexecMulti

---
 src/bin/psql/common.c            | 44 +++++++++++++++++++-------------
 src/interfaces/libpq/exports.txt |  1 +
 src/interfaces/libpq/fe-exec.c   | 37 +++++++++++++++++++++++++++
 src/interfaces/libpq/libpq-fe.h  |  1 +
 4 files changed, 65 insertions(+), 18 deletions(-)

diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index 3503605a7d..deaaa54f10 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -1195,7 +1195,6 @@ bool
 SendQuery(const char *query)
 {
 	bool		timing = pset.timing;
-	PGresult   *results;
 	PGTransactionStatusType transaction_status;
 	double		elapsed_msec = 0;
 	bool		OK = false;
@@ -1247,15 +1246,17 @@ SendQuery(const char *query)
 		!pset.autocommit &&
 		!command_no_begin(query))
 	{
-		results = PQexec(pset.db, "BEGIN");
-		if (PQresultStatus(results) != PGRES_COMMAND_OK)
+		PGresult   *result;
+
+		result = PQexec(pset.db, "BEGIN");
+		if (PQresultStatus(result) != PGRES_COMMAND_OK)
 		{
 			pg_log_info("%s", PQerrorMessage(pset.db));
-			ClearOrSaveResult(results);
+			ClearOrSaveResult(result);
 			ResetCancelConn();
 			goto sendquery_cleanup;
 		}
-		ClearOrSaveResult(results);
+		ClearOrSaveResult(result);
 		transaction_status = PQtransactionStatus(pset.db);
 	}
 
@@ -1264,15 +1265,17 @@ SendQuery(const char *query)
 		(pset.cur_cmd_interactive ||
 		 pset.on_error_rollback == PSQL_ERROR_ROLLBACK_ON))
 	{
-		results = PQexec(pset.db, "SAVEPOINT pg_psql_temporary_savepoint");
-		if (PQresultStatus(results) != PGRES_COMMAND_OK)
+		PGresult   *result;
+
+		result = PQexec(pset.db, "SAVEPOINT pg_psql_temporary_savepoint");
+		if (PQresultStatus(result) != PGRES_COMMAND_OK)
 		{
 			pg_log_info("%s", PQerrorMessage(pset.db));
-			ClearOrSaveResult(results);
+			ClearOrSaveResult(result);
 			ResetCancelConn();
 			goto sendquery_cleanup;
 		}
-		ClearOrSaveResult(results);
+		ClearOrSaveResult(result);
 		on_error_rollback_savepoint = true;
 	}
 
@@ -1281,7 +1284,6 @@ SendQuery(const char *query)
 		/* Describe query's result columns, without executing it */
 		OK = DescribeQuery(query, &elapsed_msec);
 		ResetCancelConn();
-		results = NULL;			/* PQclear(NULL) does nothing */
 	}
 	else if (pset.fetch_count <= 0 || pset.gexec_flag ||
 			 pset.crosstab_flag || !is_select_command(query))
@@ -1289,15 +1291,19 @@ SendQuery(const char *query)
 		/* Default fetch-it-all-and-print mode */
 		instr_time	before,
 					after;
+		PGresult  **results;
+		PGresult  **res;
 
 		if (timing)
 			INSTR_TIME_SET_CURRENT(before);
 
-		results = PQexec(pset.db, query);
+		results = PQexecMulti(pset.db, query);
 
 		/* these operations are included in the timing result: */
 		ResetCancelConn();
-		OK = ProcessResult(&results);
+		OK = false;
+		for (res = results; *res; res++)
+			OK |= ProcessResult(res);
 
 		if (timing)
 		{
@@ -1307,15 +1313,18 @@ SendQuery(const char *query)
 		}
 
 		/* but printing results isn't: */
-		if (OK && results)
-			OK = PrintQueryResults(results);
+		if (OK)
+			for (res = results; *res; res++)
+				OK |= PrintQueryResults(*res);
+
+		for (res = results; *res; res++)
+			ClearOrSaveResult(*res);
 	}
 	else
 	{
 		/* Fetch-in-segments mode */
 		OK = ExecQueryUsingCursor(query, &elapsed_msec);
 		ResetCancelConn();
-		results = NULL;			/* PQclear(NULL) does nothing */
 	}
 
 	if (!OK && pset.echo == PSQL_ECHO_ERRORS)
@@ -1341,6 +1350,7 @@ SendQuery(const char *query)
 
 			case PQTRANS_INTRANS:
 
+#if FIXME
 				/*
 				 * Do nothing if they are messing with savepoints themselves:
 				 * If the user did COMMIT AND CHAIN, RELEASE or ROLLBACK, our
@@ -1354,6 +1364,7 @@ SendQuery(const char *query)
 					 strcmp(PQcmdStatus(results), "ROLLBACK") == 0))
 					svptcmd = NULL;
 				else
+#endif
 					svptcmd = "RELEASE pg_psql_temporary_savepoint";
 				break;
 
@@ -1379,7 +1390,6 @@ SendQuery(const char *query)
 				ClearOrSaveResult(svptres);
 				OK = false;
 
-				PQclear(results);
 				ResetCancelConn();
 				goto sendquery_cleanup;
 			}
@@ -1387,8 +1397,6 @@ SendQuery(const char *query)
 		}
 	}
 
-	ClearOrSaveResult(results);
-
 	/* Possible microtiming output */
 	if (timing)
 		PrintTiming(elapsed_msec);
diff --git a/src/interfaces/libpq/exports.txt b/src/interfaces/libpq/exports.txt
index e8bcc88370..878e430b92 100644
--- a/src/interfaces/libpq/exports.txt
+++ b/src/interfaces/libpq/exports.txt
@@ -186,3 +186,4 @@ PQpipelineStatus          183
 PQsetTraceFlags           184
 PQmblenBounded            185
 PQsendFlushRequest        186
+PQexecMulti               187
diff --git a/src/interfaces/libpq/fe-exec.c b/src/interfaces/libpq/fe-exec.c
index 59121873d2..90f00ad401 100644
--- a/src/interfaces/libpq/fe-exec.c
+++ b/src/interfaces/libpq/fe-exec.c
@@ -70,6 +70,7 @@ static void parseInput(PGconn *conn);
 static PGresult *getCopyResult(PGconn *conn, ExecStatusType copytype);
 static bool PQexecStart(PGconn *conn);
 static PGresult *PQexecFinish(PGconn *conn);
+static PGresult **PQexecFinishMulti(PGconn *conn);
 static int	PQsendDescribe(PGconn *conn, char desc_type,
 						   const char *desc_target);
 static int	check_field_number(const PGresult *res, int field_num);
@@ -2199,6 +2200,16 @@ PQexec(PGconn *conn, const char *query)
 	return PQexecFinish(conn);
 }
 
+PGresult **
+PQexecMulti(PGconn *conn, const char *query)
+{
+	if (!PQexecStart(conn))
+		return NULL;
+	if (!PQsendQuery(conn, query))
+		return NULL;
+	return PQexecFinishMulti(conn);
+}
+
 /*
  * PQexecParams
  *		Like PQexec, but use extended query protocol so we can pass parameters
@@ -2369,6 +2380,32 @@ PQexecFinish(PGconn *conn)
 	return lastResult;
 }
 
+static PGresult **
+PQexecFinishMulti(PGconn *conn)
+{
+	int			count = 0;
+	PGresult   *result;
+	PGresult  **ret = NULL;
+
+	while ((result = PQgetResult(conn)) != NULL)
+	{
+		count++;
+		ret = realloc(ret, count * sizeof(PGresult*));
+		ret[count - 1] = result;
+
+		if (result->resultStatus == PGRES_COPY_IN ||
+			result->resultStatus == PGRES_COPY_OUT ||
+			result->resultStatus == PGRES_COPY_BOTH ||
+			conn->status == CONNECTION_BAD)
+			break;
+	}
+
+	ret = realloc(ret, (count + 1) * sizeof(PGresult*));
+	ret[count] = NULL;
+
+	return ret;
+}
+
 /*
  * PQdescribePrepared
  *	  Obtain information about a previously prepared statement
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index 20eb855abc..61b3e9bf21 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -418,6 +418,7 @@ extern void PQsetTraceFlags(PGconn *conn, int flags);
 
 /* Simple synchronous query */
 extern PGresult *PQexec(PGconn *conn, const char *query);
+extern PGresult **PQexecMulti(PGconn *conn, const char *query);
 extern PGresult *PQexecParams(PGconn *conn,
 							  const char *command,
 							  int nParams,
-- 
2.34.1