From d3b0ae12b612a049db871c80acc245055c3bde2c Mon Sep 17 00:00:00 2001
From: Sami Imseih <simseih@amazon.com>
Date: Tue, 30 Apr 2024 21:38:59 -0500
Subject: [PATCH 1/1] Fix queryId reportig in Extended Query Protocol

---
 src/backend/tcop/pquery.c           | 12 ++++++++++++
 src/backend/utils/cache/plancache.c |  8 ++++++++
 2 files changed, 20 insertions(+)

diff --git a/src/backend/tcop/pquery.c b/src/backend/tcop/pquery.c
index 0c45fcf318..275a5d3b8d 100644
--- a/src/backend/tcop/pquery.c
+++ b/src/backend/tcop/pquery.c
@@ -24,6 +24,7 @@
 #include "pg_trace.h"
 #include "tcop/pquery.h"
 #include "tcop/utility.h"
+#include "utils/backend_status.h"
 #include "utils/memutils.h"
 #include "utils/snapmgr.h"
 
@@ -882,6 +883,14 @@ PortalRunSelect(Portal portal,
 	/* Caller messed up if we have neither a ready query nor held data. */
 	Assert(queryDesc || portal->holdStore);
 
+	/*
+	 * Report the queryId which may have been reset by pgstat_report_activity
+	 * in the execute message of extended query protocol. We do the same
+	 * for PortalRunMulti.
+	 */
+	if (queryDesc && queryDesc->plannedstmt)
+		pgstat_report_query_id(queryDesc->plannedstmt->queryId, false);
+
 	/*
 	 * Force the queryDesc destination to the right thing.  This supports
 	 * MOVE, for example, which will pass in dest = DestNone.  This is okay to
@@ -1210,6 +1219,9 @@ PortalRunMulti(Portal portal,
 	if (altdest->mydest == DestRemoteExecute)
 		altdest = None_Receiver;
 
+	if (portal->stmts)
+		pgstat_report_query_id(linitial_node(PlannedStmt, portal->stmts)->queryId, false);
+
 	/*
 	 * Loop to handle the individual queries generated from a single parsetree
 	 * by analysis and rewrite.
diff --git a/src/backend/utils/cache/plancache.c b/src/backend/utils/cache/plancache.c
index 5af1a168ec..762ff0a8ad 100644
--- a/src/backend/utils/cache/plancache.c
+++ b/src/backend/utils/cache/plancache.c
@@ -66,6 +66,7 @@
 #include "storage/lmgr.h"
 #include "tcop/pquery.h"
 #include "tcop/utility.h"
+#include "utils/backend_status.h"
 #include "utils/inval.h"
 #include "utils/memutils.h"
 #include "utils/resowner.h"
@@ -1182,6 +1183,13 @@ GetCachedPlan(CachedPlanSource *plansource, ParamListInfo boundParams,
 	/* Make sure the querytree list is valid and we have parse-time locks */
 	qlist = RevalidateCachedQuery(plansource, queryEnv);
 
+	/*
+	 * In case the query tree underwent re-analysis, we may have a new
+	 * queryId to report.
+	 */
+	if (qlist)
+		pgstat_report_query_id(linitial_node(Query, qlist)->queryId, false);
+
 	/* Decide whether to use a custom plan */
 	customplan = choose_custom_plan(plansource, boundParams);
 
-- 
2.39.3 (Apple Git-146)