From b70163d1bd1279ca3f9b29166431cc64d25ca586 Mon Sep 17 00:00:00 2001 From: Anthonin Bonnefoy 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)