0001-Output-affected-rows-in-EXPLAIN.patch

From c6cbc6fa9ddf24f29bc19ff115224dd76e351db1 Mon Sep 17 00:00:00 2001
From: Damir Belyalov <d.belyalov@postgrespro.ru>
Date: Tue, 5 Sep 2023 15:04:01 +0300
Subject: [PATCH] Output affected rows in EXPLAIN.

---
 src/backend/commands/explain.c | 10 +++++++++-
 src/backend/tcop/cmdtag.c      |  2 +-
 src/backend/tcop/pquery.c      |  8 +++++++-
 src/backend/tcop/utility.c     | 27 ++++++++++++++++++++++++++-
 src/bin/psql/common.c          |  5 +++--
 src/include/commands/explain.h |  3 ++-
 src/include/tcop/cmdtag.h      |  1 +
 src/include/tcop/cmdtaglist.h  |  3 +++
 src/interfaces/libpq/fe-exec.c |  4 +++-
 9 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
index 8570b14f62..453e545ba5 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -162,7 +162,7 @@ static void escape_yaml(StringInfo buf, const char *str);
  */
 void
 ExplainQuery(ParseState *pstate, ExplainStmt *stmt,
-			 ParamListInfo params, DestReceiver *dest)
+			 ParamListInfo params, DestReceiver *dest, uint64 *processed)
 {
 	ExplainState *es = NewExplainState();
 	TupOutputState *tstate;
@@ -173,6 +173,9 @@ ExplainQuery(ParseState *pstate, ExplainStmt *stmt,
 	bool		timing_set = false;
 	bool		summary_set = false;
 
+	if (processed)
+		*processed = 0;
+
 	/* Parse options list. */
 	foreach(lc, stmt->options)
 	{
@@ -311,6 +314,9 @@ ExplainQuery(ParseState *pstate, ExplainStmt *stmt,
 	end_tup_output(tstate);
 
 	pfree(es->str->data);
+
+	if (processed)
+		*processed = es->es_processed;
 }
 
 /*
@@ -649,6 +655,8 @@ ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into, ExplainState *es,
 	 */
 	INSTR_TIME_SET_CURRENT(starttime);
 
+	es->es_processed += queryDesc->estate->es_processed;
+
 	ExecutorEnd(queryDesc);
 
 	FreeQueryDesc(queryDesc);
diff --git a/src/backend/tcop/cmdtag.c b/src/backend/tcop/cmdtag.c
index 4bd713a0b4..9e6fdbd8af 100644
--- a/src/backend/tcop/cmdtag.c
+++ b/src/backend/tcop/cmdtag.c
@@ -146,7 +146,7 @@ BuildQueryCompletionString(char *buff, const QueryCompletion *qc,
 	 */
 	if (command_tag_display_rowcount(tag) && !nameonly)
 	{
-		if (tag == CMDTAG_INSERT)
+		if (tag == CMDTAG_INSERT || tag == CMDTAG_EXPLAIN_INSERT)
 		{
 			*bufp++ = ' ';
 			*bufp++ = '0';
diff --git a/src/backend/tcop/pquery.c b/src/backend/tcop/pquery.c
index 5565f200c3..ba0b33cc67 100644
--- a/src/backend/tcop/pquery.c
+++ b/src/backend/tcop/pquery.c
@@ -775,7 +775,13 @@ PortalRun(Portal portal, long count, bool isTopLevel, bool run_once,
 				if (qc && portal->qc.commandTag != CMDTAG_UNKNOWN)
 				{
 					CopyQueryCompletion(qc, &portal->qc);
-					qc->nprocessed = nprocessed;
+					if (portal->qc.commandTag == CMDTAG_EXPLAIN ||
+						portal->qc.commandTag == CMDTAG_EXPLAIN_INSERT ||
+						portal->qc.commandTag == CMDTAG_EXPLAIN_UPDATE ||
+						portal->qc.commandTag == CMDTAG_EXPLAIN_DELETE)
+						qc->nprocessed = portal->qc.nprocessed;
+					else
+						qc->nprocessed = nprocessed;
 				}
 
 				/* Mark portal not active */
diff --git a/src/backend/tcop/utility.c b/src/backend/tcop/utility.c
index e3ccf6c7f7..8975d046f9 100644
--- a/src/backend/tcop/utility.c
+++ b/src/backend/tcop/utility.c
@@ -867,7 +867,32 @@ standard_ProcessUtility(PlannedStmt *pstmt,
 			break;
 
 		case T_ExplainStmt:
-			ExplainQuery(pstate, (ExplainStmt *) parsetree, params, dest);
+			{
+				Query	   *query;
+				uint64		processed;
+				int			explainTag;
+
+				ExplainQuery(pstate, (ExplainStmt *) parsetree, params, dest, &processed);
+
+				query = castNode(Query, ((ExplainStmt *) parsetree)->query);
+				switch (query->commandType)
+				{
+					case CMD_INSERT:
+						explainTag = CMDTAG_EXPLAIN_INSERT;
+						break;
+					case CMD_UPDATE:
+						explainTag = CMDTAG_EXPLAIN_UPDATE;
+						break;
+					case CMD_DELETE:
+						explainTag = CMDTAG_EXPLAIN_DELETE;
+						break;
+					default:
+						explainTag = CMDTAG_EXPLAIN;
+						break;
+				}
+				if (qc)
+					SetQueryCompletion(qc, explainTag, processed);
+			}
 			break;
 
 		case T_AlterSystemStmt:
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index ede197bebe..a66d9127c5 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -987,8 +987,9 @@ PrintQueryResult(PGresult *result, bool last,
 			if (last || pset.show_all_results)
 			{
 				cmdstatus = PQcmdStatus(result);
-				if (strncmp(cmdstatus, "INSERT", 6) == 0 ||
-					strncmp(cmdstatus, "UPDATE", 6) == 0 ||
+				if (strncmp(cmdstatus, "EXPLAIN", 7) == 0 ||
+					strncmp(cmdstatus, "INSERT", 6) == 0  ||
+					strncmp(cmdstatus, "UPDATE", 6) == 0  ||
 					strncmp(cmdstatus, "DELETE", 6) == 0)
 					PrintQueryStatus(result, printStatusFout);
 			}
diff --git a/src/include/commands/explain.h b/src/include/commands/explain.h
index 3d3e632a0c..21fe5f7555 100644
--- a/src/include/commands/explain.h
+++ b/src/include/commands/explain.h
@@ -60,6 +60,7 @@ typedef struct ExplainState
 	bool		hide_workers;	/* set if we find an invisible Gather */
 	/* state related to the current plan node */
 	ExplainWorkersState *workers_state; /* needed if parallel plan */
+	uint64		es_processed;	/* sum of queryDesc->estate->es_processed */
 } ExplainState;
 
 /* Hook for plugins to get control in ExplainOneQuery() */
@@ -78,7 +79,7 @@ extern PGDLLIMPORT explain_get_index_name_hook_type explain_get_index_name_hook;
 
 
 extern void ExplainQuery(ParseState *pstate, ExplainStmt *stmt,
-						 ParamListInfo params, DestReceiver *dest);
+						 ParamListInfo params, DestReceiver *dest, uint64 *processed);
 
 extern ExplainState *NewExplainState(void);
 
diff --git a/src/include/tcop/cmdtag.h b/src/include/tcop/cmdtag.h
index 1e7514dcff..49f7ea85e7 100644
--- a/src/include/tcop/cmdtag.h
+++ b/src/include/tcop/cmdtag.h
@@ -30,6 +30,7 @@ typedef enum CommandTag
 typedef struct QueryCompletion
 {
 	CommandTag	commandTag;
+	CommandTag	explainCommandTag;
 	uint64		nprocessed;
 } QueryCompletion;
 
diff --git a/src/include/tcop/cmdtaglist.h b/src/include/tcop/cmdtaglist.h
index e738ac1c09..fdc570a304 100644
--- a/src/include/tcop/cmdtaglist.h
+++ b/src/include/tcop/cmdtaglist.h
@@ -178,6 +178,9 @@ PG_CMDTAG(CMDTAG_DROP_USER_MAPPING, "DROP USER MAPPING", true, false, false)
 PG_CMDTAG(CMDTAG_DROP_VIEW, "DROP VIEW", true, false, false)
 PG_CMDTAG(CMDTAG_EXECUTE, "EXECUTE", false, false, false)
 PG_CMDTAG(CMDTAG_EXPLAIN, "EXPLAIN", false, false, false)
+PG_CMDTAG(CMDTAG_EXPLAIN_INSERT, "INSERT", false, false, true)
+PG_CMDTAG(CMDTAG_EXPLAIN_UPDATE, "UPDATE", false, false, true)
+PG_CMDTAG(CMDTAG_EXPLAIN_DELETE, "DELETE", false, false, true)
 PG_CMDTAG(CMDTAG_FETCH, "FETCH", false, false, true)
 PG_CMDTAG(CMDTAG_GRANT, "GRANT", true, false, false)
 PG_CMDTAG(CMDTAG_GRANT_ROLE, "GRANT ROLE", false, false, false)
diff --git a/src/interfaces/libpq/fe-exec.c b/src/interfaces/libpq/fe-exec.c
index c6d80ec396..b85e9b8e04 100644
--- a/src/interfaces/libpq/fe-exec.c
+++ b/src/interfaces/libpq/fe-exec.c
@@ -3741,7 +3741,9 @@ PQcmdTuples(PGresult *res)
 	if (!res)
 		return "";
 
-	if (strncmp(res->cmdStatus, "INSERT ", 7) == 0)
+	if (strncmp(res->cmdStatus, "EXPLAIN ", 8) == 0)
+		p = res->cmdStatus + 8;
+	else if (strncmp(res->cmdStatus, "INSERT ", 7) == 0)
 	{
 		p = res->cmdStatus + 7;
 		/* INSERT: skip oid and space */
-- 
2.34.1