Hi, 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. Regards, Anthonin
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)
