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

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index d924c87b41ed9988535dca1076a7e5a76483f0e0..0a059c4970e8327444dface69dbed729e56c4c65 100644 (file)
--- 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)
+
 -- PL/pgSQL function - top-level tracking.
 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 0c20b8ce69b71f0b4f5ff9d2d3d56b70e652e738..51b7bd360dee46a0aca017b5c8ac9098c853dbea 100644 (file)
--- 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;
+
 -- PL/pgSQL function - top-level tracking.
 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 89fcd3ec5fe35bca7aa6fdcd83aa07924065302a..81b0c029e7a4c5ff5127cb9d1fe3ea14dbe517c4 100644 (file)
--- 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,