From cc4f2314f06ed7b17bdb87d5558c87992dc58452 Mon Sep 17 00:00:00 2001 From: Anthonin Bonnefoy Date: Tue, 23 Jul 2024 08:26:49 +0200 Subject: Set query_id for queries contained in utility statement Some utility statements like Explain, CreateTableAs, DeclareCursor and CreateMaterializedView contain a query which will be planned and executed. During post parse, only the top utility statement is jumbled, leaving the contained query without a set query_id. ExplainQuery does jumble the other three do not. This led to extensions relying on query_id like pg_stat_statements to not be able to track those nested queries as the query_id was 0. This patch fixes this by jumbling the nested query of CreateTableAs, DeclareCursor and CreateMaterializedView before it is executed. Additionally, explain's nested query can itself be a CreateTableAs, DeclareCursor or CreateMaterializedView which also needs to be jumbled. This is now done in ExplainOneUtility. --- .../expected/level_tracking.out | 18 +++++--- src/backend/commands/createas.c | 12 +++++- src/backend/commands/explain.c | 43 ++++++++++++------- src/backend/commands/matview.c | 25 ++++++++--- src/backend/commands/portalcmds.c | 10 +++++ src/backend/commands/prepare.c | 20 ++++----- src/backend/tcop/utility.c | 2 +- src/include/commands/explain.h | 4 +- src/include/commands/matview.h | 5 ++- src/include/commands/prepare.h | 4 +- src/test/regress/expected/explain.out | 17 ++++++++ src/test/regress/sql/explain.sql | 4 ++ 12 files changed, 116 insertions(+), 48 deletions(-) diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out index 74df1c3457d..28623ae21cc 100644 --- a/contrib/pg_stat_statements/expected/level_tracking.out +++ b/contrib/pg_stat_statements/expected/level_tracking.out @@ -764,7 +764,8 @@ SELECT toplevel, calls, query FROM pg_stat_statements t | 1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1 t | 1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t f | 1 | SELECT $1 -(4 rows) + f | 1 | SELECT * FROM stats_track_tab +(5 rows) -- Explain analyze, top tracking. SET pg_stat_statements.track = 'top'; @@ -810,7 +811,8 @@ SELECT toplevel, calls, query FROM pg_stat_statements ----------+-------+------------------------------------------------------------------------------------------------ t | 1 | CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series($1, $2) as id t | 1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t -(2 rows) + f | 1 | SELECT * FROM generate_series($1, $2) as id +(3 rows) -- Create Materialized View, top-level tracking. SET pg_stat_statements.track = 'top'; @@ -844,7 +846,8 @@ SELECT toplevel, calls, query FROM pg_stat_statements ----------+-------+---------------------------------------------------- t | 1 | REFRESH MATERIALIZED VIEW pgss_materialized_view t | 1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t -(2 rows) + f | 1 | REFRESH MATERIALIZED VIEW pgss_materialized_view; +(3 rows) -- Refresh Materialized View, top-level tracking. SET pg_stat_statements.track = 'top'; @@ -882,7 +885,8 @@ SELECT toplevel, calls, query FROM pg_stat_statements t | 1 | CREATE TEMPORARY TABLE pgss_ctas_2 AS EXECUTE test_prepare_pgss t | 1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t f | 1 | PREPARE test_prepare_pgss AS select generate_series($1, $2) -(4 rows) + f | 1 | SELECT $1 +(5 rows) -- Create Table As, top-level tracking. SET pg_stat_statements.track = 'top'; @@ -923,7 +927,8 @@ SELECT toplevel, calls, query FROM pg_stat_statements ----------+-------+--------------------------------------------------------------------------- t | 1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t t | 1 | explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT $1 -(2 rows) + f | 1 | SELECT $1 +(3 rows) -- Explain with Create Table As - top-level tracking. SET pg_stat_statements.track = 'top'; @@ -974,7 +979,8 @@ SELECT toplevel, calls, query FROM pg_stat_statements t | 1 | DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab t | 1 | FETCH FORWARD 1 FROM foocur t | 1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t -(6 rows) + f | 1 | SELECT * from stats_track_tab +(7 rows) -- Declare cursor, top-level tracking. SET pg_stat_statements.track = 'top'; diff --git a/src/backend/commands/createas.c b/src/backend/commands/createas.c index 0b629b1f79c..7f921cf9a99 100644 --- a/src/backend/commands/createas.c +++ b/src/backend/commands/createas.c @@ -38,6 +38,8 @@ #include "miscadmin.h" #include "nodes/makefuncs.h" #include "nodes/nodeFuncs.h" +#include "nodes/queryjumble.h" +#include "parser/analyze.h" #include "rewrite/rewriteHandler.h" #include "tcop/tcopprot.h" #include "utils/builtins.h" @@ -224,6 +226,7 @@ ExecCreateTableAs(ParseState *pstate, CreateTableAsStmt *stmt, { Query *query = castNode(Query, stmt->query); IntoClause *into = stmt->into; + JumbleState *jstate = NULL; bool is_matview = (into->viewQuery != NULL); bool do_refresh = false; DestReceiver *dest; @@ -238,6 +241,13 @@ ExecCreateTableAs(ParseState *pstate, CreateTableAsStmt *stmt, */ dest = CreateIntoRelDestReceiver(into); + /* Query contained by CTAS needs to be jumbled if requested */ + if (IsQueryIdEnabled()) + jstate = JumbleQuery(query); + + if (post_parse_analyze_hook) + (*post_parse_analyze_hook) (pstate, query, jstate); + /* * The contained Query could be a SELECT, or an EXECUTE utility command. * If the latter, we just pass it off to ExecuteQuery. @@ -284,7 +294,7 @@ ExecCreateTableAs(ParseState *pstate, CreateTableAsStmt *stmt, */ if (do_refresh) RefreshMatViewByOid(address.objectId, true, false, false, - pstate->p_sourcetext, qc); + pstate, qc); } else diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c index ee1bcb84e28..e805ca7ddac 100644 --- a/src/backend/commands/explain.c +++ b/src/backend/commands/explain.c @@ -71,8 +71,7 @@ typedef struct SerializeMetrics static void ExplainOneQuery(Query *query, int cursorOptions, IntoClause *into, ExplainState *es, - const char *queryString, ParamListInfo params, - QueryEnvironment *queryEnv); + ParseState *pstate, ParamListInfo params); static void ExplainPrintJIT(ExplainState *es, int jit_flags, JitInstrumentation *ji); static void ExplainPrintSerialize(ExplainState *es, @@ -350,7 +349,7 @@ ExplainQuery(ParseState *pstate, ExplainStmt *stmt, { ExplainOneQuery(lfirst_node(Query, l), CURSOR_OPT_PARALLEL_OK, NULL, es, - pstate->p_sourcetext, params, pstate->p_queryEnv); + pstate, params); /* Separate plans with an appropriate separator */ if (lnext(rewritten, l) != NULL) @@ -436,24 +435,22 @@ ExplainResultDesc(ExplainStmt *stmt) static void ExplainOneQuery(Query *query, int cursorOptions, IntoClause *into, ExplainState *es, - const char *queryString, ParamListInfo params, - QueryEnvironment *queryEnv) + ParseState *pstate, ParamListInfo params) { /* planner will not cope with utility statements */ if (query->commandType == CMD_UTILITY) { - ExplainOneUtility(query->utilityStmt, into, es, queryString, params, - queryEnv); + ExplainOneUtility(query->utilityStmt, into, es, pstate, params); return; } /* if an advisor plugin is present, let it manage things */ if (ExplainOneQuery_hook) (*ExplainOneQuery_hook) (query, cursorOptions, into, es, - queryString, params, queryEnv); + pstate->p_sourcetext, params, pstate->p_queryEnv); else standard_ExplainOneQuery(query, cursorOptions, into, es, - queryString, params, queryEnv); + pstate->p_sourcetext, params, pstate->p_queryEnv); } /* @@ -534,9 +531,10 @@ standard_ExplainOneQuery(Query *query, int cursorOptions, */ void ExplainOneUtility(Node *utilityStmt, IntoClause *into, ExplainState *es, - const char *queryString, ParamListInfo params, - QueryEnvironment *queryEnv) + ParseState *pstate, ParamListInfo params) { + JumbleState *jstate = NULL; + if (utilityStmt == NULL) return; @@ -547,6 +545,7 @@ ExplainOneUtility(Node *utilityStmt, IntoClause *into, ExplainState *es, * ExplainOneQuery. Copy to be safe in the EXPLAIN EXECUTE case. */ CreateTableAsStmt *ctas = (CreateTableAsStmt *) utilityStmt; + Query *ctas_query; List *rewritten; /* @@ -565,11 +564,16 @@ ExplainOneUtility(Node *utilityStmt, IntoClause *into, ExplainState *es, return; } - rewritten = QueryRewrite(castNode(Query, copyObject(ctas->query))); + ctas_query = castNode(Query, copyObject(ctas->query)); + if (IsQueryIdEnabled()) + jstate = JumbleQuery(ctas_query); + if (post_parse_analyze_hook) + (*post_parse_analyze_hook) (pstate, ctas_query, jstate); + rewritten = QueryRewrite(ctas_query); Assert(list_length(rewritten) == 1); ExplainOneQuery(linitial_node(Query, rewritten), CURSOR_OPT_PARALLEL_OK, ctas->into, es, - queryString, params, queryEnv); + pstate, params); } else if (IsA(utilityStmt, DeclareCursorStmt)) { @@ -582,17 +586,24 @@ ExplainOneUtility(Node *utilityStmt, IntoClause *into, ExplainState *es, * be created, however. */ DeclareCursorStmt *dcs = (DeclareCursorStmt *) utilityStmt; + Query *dcs_query; List *rewritten; - rewritten = QueryRewrite(castNode(Query, copyObject(dcs->query))); + dcs_query = castNode(Query, copyObject(dcs->query)); + if (IsQueryIdEnabled()) + jstate = JumbleQuery(dcs_query); + if (post_parse_analyze_hook) + (*post_parse_analyze_hook) (pstate, dcs_query, jstate); + + rewritten = QueryRewrite(dcs_query); Assert(list_length(rewritten) == 1); ExplainOneQuery(linitial_node(Query, rewritten), dcs->options, NULL, es, - queryString, params, queryEnv); + pstate, params); } else if (IsA(utilityStmt, ExecuteStmt)) ExplainExecuteQuery((ExecuteStmt *) utilityStmt, into, es, - queryString, params, queryEnv); + pstate, params); else if (IsA(utilityStmt, NotifyStmt)) { if (es->format == EXPLAIN_FORMAT_TEXT) diff --git a/src/backend/commands/matview.c b/src/backend/commands/matview.c index 010097873d1..7cc68338837 100644 --- a/src/backend/commands/matview.c +++ b/src/backend/commands/matview.c @@ -32,6 +32,7 @@ #include "executor/spi.h" #include "miscadmin.h" #include "pgstat.h" +#include "parser/analyze.h" #include "rewrite/rewriteHandler.h" #include "storage/lmgr.h" #include "tcop/tcopprot.h" @@ -60,7 +61,8 @@ static bool transientrel_receive(TupleTableSlot *slot, DestReceiver *self); static void transientrel_shutdown(DestReceiver *self); static void transientrel_destroy(DestReceiver *self); static uint64 refresh_matview_datafill(DestReceiver *dest, Query *query, - const char *queryString, bool is_create); + ParseState *pstate, const char *queryString, + bool is_create); static char *make_temptable_name_n(char *tempname, int n); static void refresh_by_match_merge(Oid matviewOid, Oid tempOid, Oid relowner, int save_sec_context); @@ -118,7 +120,7 @@ SetMatViewPopulatedState(Relation relation, bool newstate) * skipData field shows whether the clause was used. */ ObjectAddress -ExecRefreshMatView(RefreshMatViewStmt *stmt, const char *queryString, +ExecRefreshMatView(RefreshMatViewStmt *stmt, ParseState *pstate, QueryCompletion *qc) { Oid matviewOid; @@ -136,7 +138,7 @@ ExecRefreshMatView(RefreshMatViewStmt *stmt, const char *queryString, NULL); return RefreshMatViewByOid(matviewOid, false, stmt->skipData, - stmt->concurrent, queryString, qc); + stmt->concurrent, pstate, qc); } /* @@ -163,7 +165,7 @@ ExecRefreshMatView(RefreshMatViewStmt *stmt, const char *queryString, */ ObjectAddress RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData, - bool concurrent, const char *queryString, + bool concurrent, ParseState *pstate, QueryCompletion *qc) { Relation matviewRel; @@ -325,10 +327,11 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData, if (!skipData) { DestReceiver *dest; + const char *queryString = pstate->p_sourcetext; dest = CreateTransientRelDestReceiver(OIDNewHeap); - processed = refresh_matview_datafill(dest, dataQuery, queryString, - is_create); + processed = refresh_matview_datafill(dest, dataQuery, pstate, + queryString, is_create); } /* Make the matview match the newly generated data. */ @@ -403,17 +406,25 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData, */ static uint64 refresh_matview_datafill(DestReceiver *dest, Query *query, - const char *queryString, bool is_create) + ParseState *pstate, const char *queryString, + bool is_create) { List *rewritten; PlannedStmt *plan; QueryDesc *queryDesc; Query *copied_query; uint64 processed; + JumbleState *jstate = NULL; /* Lock and rewrite, using a copy to preserve the original query. */ copied_query = copyObject(query); AcquireRewriteLocks(copied_query, true, false); + + if (IsQueryIdEnabled()) + jstate = JumbleQuery(copied_query); + if (post_parse_analyze_hook) + (*post_parse_analyze_hook) (pstate, copied_query, jstate); + rewritten = QueryRewrite(copied_query); /* SELECT should never rewrite to more or less than one SELECT query */ diff --git a/src/backend/commands/portalcmds.c b/src/backend/commands/portalcmds.c index 4f6acf67198..ac52ca25e99 100644 --- a/src/backend/commands/portalcmds.c +++ b/src/backend/commands/portalcmds.c @@ -28,6 +28,8 @@ #include "executor/executor.h" #include "executor/tstoreReceiver.h" #include "miscadmin.h" +#include "nodes/queryjumble.h" +#include "parser/analyze.h" #include "rewrite/rewriteHandler.h" #include "tcop/pquery.h" #include "tcop/tcopprot.h" @@ -44,6 +46,7 @@ PerformCursorOpen(ParseState *pstate, DeclareCursorStmt *cstmt, ParamListInfo pa bool isTopLevel) { Query *query = castNode(Query, cstmt->query); + JumbleState *jstate = NULL; List *rewritten; PlannedStmt *plan; Portal portal; @@ -71,6 +74,13 @@ PerformCursorOpen(ParseState *pstate, DeclareCursorStmt *cstmt, ParamListInfo pa (errcode(ERRCODE_INSUFFICIENT_PRIVILEGE), errmsg("cannot create a cursor WITH HOLD within security-restricted operation"))); + /* Query contained by DeclareCursor needs to be jumbled if requested */ + if (IsQueryIdEnabled()) + jstate = JumbleQuery(query); + + if (post_parse_analyze_hook) + (*post_parse_analyze_hook) (pstate, query, jstate); + /* * Parse analysis was done already, but we still have to run the rule * rewriter. We do not do AcquireRewriteLocks: we assume the query either diff --git a/src/backend/commands/prepare.c b/src/backend/commands/prepare.c index 07257d4db94..a93f970a292 100644 --- a/src/backend/commands/prepare.c +++ b/src/backend/commands/prepare.c @@ -561,13 +561,12 @@ DropAllPreparedStatements(void) * "into" is NULL unless we are doing EXPLAIN CREATE TABLE AS EXECUTE, * in which case executing the query should result in creating that table. * - * Note: the passed-in queryString is that of the EXPLAIN EXECUTE, + * Note: the passed-in pstate's queryString is that of the EXPLAIN EXECUTE, * not the original PREPARE; we get the latter string from the plancache. */ void ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es, - const char *queryString, ParamListInfo params, - QueryEnvironment *queryEnv) + ParseState *pstate, ParamListInfo params) { PreparedStatement *entry; const char *query_string; @@ -610,10 +609,10 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es, /* Evaluate parameters, if any */ if (entry->plansource->num_params) { - ParseState *pstate; + ParseState *pstate_params; - pstate = make_parsestate(NULL); - pstate->p_sourcetext = queryString; + pstate_params = make_parsestate(NULL); + pstate_params->p_sourcetext = pstate->p_sourcetext; /* * Need an EState to evaluate parameters; must not delete it till end @@ -624,12 +623,12 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es, estate = CreateExecutorState(); estate->es_param_list_info = params; - paramLI = EvaluateParams(pstate, entry, execstmt->params, estate); + paramLI = EvaluateParams(pstate_params, entry, execstmt->params, estate); } /* Replan if needed, and acquire a transient refcount */ cplan = GetCachedPlan(entry->plansource, paramLI, - CurrentResourceOwner, queryEnv); + CurrentResourceOwner, pstate->p_queryEnv); INSTR_TIME_SET_CURRENT(planduration); INSTR_TIME_SUBTRACT(planduration, planstart); @@ -655,12 +654,11 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es, PlannedStmt *pstmt = lfirst_node(PlannedStmt, p); if (pstmt->commandType != CMD_UTILITY) - ExplainOnePlan(pstmt, into, es, query_string, paramLI, queryEnv, + ExplainOnePlan(pstmt, into, es, query_string, paramLI, pstate->p_queryEnv, &planduration, (es->buffers ? &bufusage : NULL), es->memory ? &mem_counters : NULL); else - ExplainOneUtility(pstmt->utilityStmt, into, es, query_string, - paramLI, queryEnv); + ExplainOneUtility(pstmt->utilityStmt, into, es, pstate, paramLI); /* No need for CommandCounterIncrement, as ExplainOnePlan did it */ diff --git a/src/backend/tcop/utility.c b/src/backend/tcop/utility.c index b2ea8125c92..4768b4f746b 100644 --- a/src/backend/tcop/utility.c +++ b/src/backend/tcop/utility.c @@ -1688,7 +1688,7 @@ ProcessUtilitySlow(ParseState *pstate, PG_TRY(2); { address = ExecRefreshMatView((RefreshMatViewStmt *) parsetree, - queryString, qc); + pstate, qc); } PG_FINALLY(2); { diff --git a/src/include/commands/explain.h b/src/include/commands/explain.h index 3ab0aae78f7..aa5872bc154 100644 --- a/src/include/commands/explain.h +++ b/src/include/commands/explain.h @@ -100,8 +100,8 @@ extern ExplainState *NewExplainState(void); extern TupleDesc ExplainResultDesc(ExplainStmt *stmt); extern void ExplainOneUtility(Node *utilityStmt, IntoClause *into, - ExplainState *es, const char *queryString, - ParamListInfo params, QueryEnvironment *queryEnv); + ExplainState *es, ParseState *pstate, + ParamListInfo params); extern void ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into, ExplainState *es, const char *queryString, diff --git a/src/include/commands/matview.h b/src/include/commands/matview.h index c8811e8fc71..6602640b400 100644 --- a/src/include/commands/matview.h +++ b/src/include/commands/matview.h @@ -17,16 +17,17 @@ #include "catalog/objectaddress.h" #include "nodes/params.h" #include "nodes/parsenodes.h" +#include "parser/parse_node.h" #include "tcop/dest.h" #include "utils/relcache.h" extern void SetMatViewPopulatedState(Relation relation, bool newstate); -extern ObjectAddress ExecRefreshMatView(RefreshMatViewStmt *stmt, const char *queryString, +extern ObjectAddress ExecRefreshMatView(RefreshMatViewStmt *stmt, ParseState *pstate, QueryCompletion *qc); extern ObjectAddress RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData, - bool concurrent, const char *queryString, + bool concurrent, ParseState *pstate, QueryCompletion *qc); extern DestReceiver *CreateTransientRelDestReceiver(Oid transientoid); diff --git a/src/include/commands/prepare.h b/src/include/commands/prepare.h index 61472c111d6..e6fd400e027 100644 --- a/src/include/commands/prepare.h +++ b/src/include/commands/prepare.h @@ -43,8 +43,8 @@ extern void ExecuteQuery(ParseState *pstate, DestReceiver *dest, QueryCompletion *qc); extern void DeallocateQuery(DeallocateStmt *stmt); extern void ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, - ExplainState *es, const char *queryString, - ParamListInfo params, QueryEnvironment *queryEnv); + ExplainState *es, ParseState *pstate, + ParamListInfo params); /* Low-level access to stored prepared statements */ extern void StorePreparedStatement(const char *stmt_name, diff --git a/src/test/regress/expected/explain.out b/src/test/regress/expected/explain.out index d01c304c24d..664050a9fa3 100644 --- a/src/test/regress/expected/explain.out +++ b/src/test/regress/expected/explain.out @@ -653,6 +653,23 @@ select explain_filter('explain (verbose) select * from int8_tbl i8'); Query Identifier: N (3 rows) +-- Test compute_query_id with utility statements containing plannable query +select explain_filter('explain (verbose) declare test_cur cursor for select * from int8_tbl'); + explain_filter +------------------------------------------------------------- + Seq Scan on public.int8_tbl (cost=N.N..N.N rows=N width=N) + Output: q1, q2 + Query Identifier: N +(3 rows) + +select explain_filter('explain (verbose) create table test_ctas as select 1'); + explain_filter +---------------------------------------- + Result (cost=N.N..N.N rows=N width=N) + Output: N + Query Identifier: N +(3 rows) + -- Test SERIALIZE option select explain_filter('explain (analyze,serialize) select * from int8_tbl i8'); explain_filter diff --git a/src/test/regress/sql/explain.sql b/src/test/regress/sql/explain.sql index b861e2b53d5..3ca285a1d7e 100644 --- a/src/test/regress/sql/explain.sql +++ b/src/test/regress/sql/explain.sql @@ -163,6 +163,10 @@ select explain_filter('explain (verbose) select * from t1 where pg_temp.mysin(f1 set compute_query_id = on; select explain_filter('explain (verbose) select * from int8_tbl i8'); +-- Test compute_query_id with utility statements containing plannable query +select explain_filter('explain (verbose) declare test_cur cursor for select * from int8_tbl'); +select explain_filter('explain (verbose) create table test_ctas as select 1'); + -- Test SERIALIZE option select explain_filter('explain (analyze,serialize) select * from int8_tbl i8'); select explain_filter('explain (analyze,serialize text,buffers,timing off) select * from int8_tbl i8'); -- 2.39.5 (Apple Git-154)