Correctly propagate queryId for utility stmt in function

From b70163d1bd1279ca3f9b29166431cc64d25ca586 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Thu, 18 Jul 2024 11:50:37 +0200
Subject: Correctly propagate queryId for utility stmt in function

For utility statements defined within a function, the queryTree is
copied to a plannedStmt as utility commands don't require planning.
However, the queryId is not propagated to the plannedStmt. This leads to
plugins relying on queryId like pg_stat_statements to not be able to
track utility statements within function calls.

This patch fixes the issue by correctly propagating queryId from the
cached queryTree to the plannedStmt.
---
 .../expected/level_tracking.out               | 28 +++++++++++++++++++
 .../pg_stat_statements/sql/level_tracking.sql | 12 ++++++++
 src/backend/executor/functions.c              |  1 +
 3 files changed, 41 insertions(+)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index fe01dd79af..f7b2ffceb3 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -67,6 +67,34 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | SET pg_stat_statements.track = 'all'
 (7 rows)
 
+-- Procedure with multiple utility statements - all level tracking.
+CREATE OR REPLACE PROCEDURE proc_with_utility_stmt()
+LANGUAGE SQL
+AS $$
+    SHOW pg_stat_statements.track;
+    SHOW pg_stat_statements.track_utility;
+$$;
+CALL proc_with_utility_stmt();
+SELECT toplevel, calls, query FROM pg_stat_statements
+  WHERE query LIKE '%SHOW%' ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                        query                         
+----------+-------+------------------------------------------------------
+ t        |     1 | CREATE OR REPLACE PROCEDURE proc_with_utility_stmt()+
+          |       | LANGUAGE SQL                                        +
+          |       | AS $$                                               +
+          |       |     SHOW pg_stat_statements.track;                  +
+          |       |     SHOW pg_stat_statements.track_utility;          +
+          |       | $$
+ f        |     1 | SHOW pg_stat_statements.track
+ f        |     1 | SHOW pg_stat_statements.track_utility
+(3 rows)
+
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
 -- DO block - top-level tracking without utility.
 SET pg_stat_statements.track = 'top';
 SET pg_stat_statements.track_utility = FALSE;
diff --git a/contrib/pg_stat_statements/sql/level_tracking.sql b/contrib/pg_stat_statements/sql/level_tracking.sql
index aa37408d52..115ff392a0 100644
--- a/contrib/pg_stat_statements/sql/level_tracking.sql
+++ b/contrib/pg_stat_statements/sql/level_tracking.sql
@@ -33,6 +33,18 @@ END; $$;
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
 
+-- Procedure with multiple utility statements - all level tracking.
+CREATE OR REPLACE PROCEDURE proc_with_utility_stmt()
+LANGUAGE SQL
+AS $$
+    SHOW pg_stat_statements.track;
+    SHOW pg_stat_statements.track_utility;
+$$;
+CALL proc_with_utility_stmt();
+SELECT toplevel, calls, query FROM pg_stat_statements
+  WHERE query LIKE '%SHOW%' ORDER BY query COLLATE "C", toplevel;
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+
 -- DO block - top-level tracking without utility.
 SET pg_stat_statements.track = 'top';
 SET pg_stat_statements.track_utility = FALSE;
diff --git a/src/backend/executor/functions.c b/src/backend/executor/functions.c
index 539cd0a999..692854e2b3 100644
--- a/src/backend/executor/functions.c
+++ b/src/backend/executor/functions.c
@@ -492,6 +492,7 @@ init_execution_state(List *queryTree_list,
 				stmt->utilityStmt = queryTree->utilityStmt;
 				stmt->stmt_location = queryTree->stmt_location;
 				stmt->stmt_len = queryTree->stmt_len;
+				stmt->queryId = queryTree->queryId;
 			}
 			else
 				stmt = pg_plan_query(queryTree,
-- 
2.39.3 (Apple Git-146)