From c145f321b681ab6c3636849071ff30d5c6f6353c Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Fri, 19 Jul 2024 10:21:01 +0900
Subject: [PATCH] Propagate query IDs of utility statements in functions

For utility statements defined within a function, the query tree is
copied to a PlannedStmt as utility commands do not require planning.
However, the query ID was missing from the information passed down.

This leads to plugins relying on the query ID like pg_stat_statements to
not be able to track utility statements within function calls.  Tests
are added to check this behavior, depending on pg_stat_statements.track.

This is an old bug.  Now, query IDs for utilities are compiled using
their parsed trees rather than the query string since v16
(3db72ebcbe20), leading to less bloat with utilities, so backpatch down
only to this version.

Author: Anthonin Bonnefoy
Discussion: https://postgr.es/m/CAO6_XqrGp-uwBqi3vBPLuRULKkddjC7R5QZCgsFren=8E+m2Sg@mail.gmail.com
Backpatch-through: 16
---
 .../expected/level_tracking.out               | 45 +++++++++++++++++++
 .../pg_stat_statements/sql/level_tracking.sql | 22 +++++++++
 src/backend/executor/functions.c              |  1 +
 3 files changed, 68 insertions(+)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index fe01dd79af..d8dd8a2dee 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -67,6 +67,51 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | SET pg_stat_statements.track = 'all'
 (7 rows)
 
+-- Procedure with multiple utility statements.
+CREATE OR REPLACE PROCEDURE proc_with_utility_stmt()
+LANGUAGE SQL
+AS $$
+  SHOW pg_stat_statements.track;
+  show pg_stat_statements.track;
+  SHOW pg_stat_statements.track_utility;
+$$;
+SET pg_stat_statements.track_utility = TRUE;
+-- all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CALL proc_with_utility_stmt();
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | CALL proc_with_utility_stmt()
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     2 | SHOW pg_stat_statements.track
+ f        |     1 | SHOW pg_stat_statements.track_utility
+(4 rows)
+
+-- top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CALL proc_with_utility_stmt();
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | CALL proc_with_utility_stmt()
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
 -- 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..65a17147a5 100644
--- a/contrib/pg_stat_statements/sql/level_tracking.sql
+++ b/contrib/pg_stat_statements/sql/level_tracking.sql
@@ -33,6 +33,28 @@ END; $$;
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
 
+-- Procedure with multiple utility statements.
+CREATE OR REPLACE PROCEDURE proc_with_utility_stmt()
+LANGUAGE SQL
+AS $$
+  SHOW pg_stat_statements.track;
+  show pg_stat_statements.track;
+  SHOW pg_stat_statements.track_utility;
+$$;
+SET pg_stat_statements.track_utility = TRUE;
+-- all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CALL proc_with_utility_stmt();
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+-- top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CALL proc_with_utility_stmt();
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
 -- 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,