On Wed, Oct 23, 2024 at 11:24:11AM +0200, Anthonin Bonnefoy wrote:
> This also answers another issue I was wondering about. Should the
> child's parsestate inherit the location information when
> make_parsestate is called? That would be incorrect since this is used
> for sub-statement, pstate should reflect the size of the whole
> sub-statement. However, since this is unused, it is fine to leave the
> child parser with unset location data, which would in turn leave the
> statement's location unset in setQueryLocationAndLength.

Yeah, this argument sounds kind of right to me.

> Patch includes Micheal changes. I've left out 0002 for now to focus on 0001.

I've looked at this one again, and applied 0001.  The final result is
really nice, thanks for all your efforts here.  If this requires
tweaks in this release cycle, well, let's deal about them should they
show up.  At least the set of regression tests will show us what's
going on.

Attached is the remaining piece, for DECLARE and CTAS.  The
JumbleQuery() calls in ExecCreateTableAs() and ExplainOneUtility() for
CTAS queries are twins, showing the inner queries of CTAS
consistently.  DECLARE is covered by one call in ExplainOneUtility()
and one in PerformCursorOpen().

This should be OK as-is.  With the regression test coverage, it is
easy to see what changes.  Let's keep that around for a few more days.
--
Michael
From aa70b2ddc36ea33af39b113c361cdc71d49df26d Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonne...@datadoghq.com>
Date: Tue, 23 Jul 2024 08:26:49 +0200
Subject: [PATCH v14] Set query_id for queries contained in utility statement

Some utility statements contain queries that can be planned and
executed: EXPLAIN, CREATE TABLE AS and DECLARE CURSOR.

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,
and DeclareCursor before it is executed, adding matching jumbling class
in EXPLAIN as these queries are supported there.
---
 src/include/commands/explain.h                |  4 +-
 src/include/commands/prepare.h                |  4 +-
 src/backend/commands/createas.c               | 10 +++++
 src/backend/commands/explain.c                | 43 ++++++++++++-------
 src/backend/commands/portalcmds.c             | 10 +++++
 src/backend/commands/prepare.c                | 20 ++++-----
 src/test/regress/expected/explain.out         | 17 ++++++++
 src/test/regress/sql/explain.sql              |  4 ++
 .../expected/level_tracking.out               | 12 ++++--
 9 files changed, 89 insertions(+), 35 deletions(-)

diff --git a/src/include/commands/explain.h b/src/include/commands/explain.h
index 3ab0aae78f..aa5872bc15 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/prepare.h b/src/include/commands/prepare.h
index 61472c111d..e6fd400e02 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/backend/commands/createas.c b/src/backend/commands/createas.c
index 68ec122dbf..e4a627ad91 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.
diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
index 18a5af6b91..c81221cdbe 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/portalcmds.c b/src/backend/commands/portalcmds.c
index 4f6acf6719..ac52ca25e9 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 07257d4db9..a93f970a29 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/test/regress/expected/explain.out b/src/test/regress/expected/explain.out
index dcbdaa0388..d2eef8097c 100644
--- a/src/test/regress/expected/explain.out
+++ b/src/test/regress/expected/explain.out
@@ -662,6 +662,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 b861e2b53d..3ca285a1d7 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');
diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 489dc7143f..9aee9f5010 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -924,8 +924,9 @@ SELECT toplevel, calls, query FROM pg_stat_statements
           |       |   DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab
  t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1
  f        |     1 | SELECT $1
+ f        |     1 | SELECT * FROM stats_track_tab
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(4 rows)
+(5 rows)
 
 -- Explain analyze, top tracking.
 SET pg_stat_statements.track = 'top';
@@ -1047,9 +1048,10 @@ SELECT toplevel, calls, query FROM pg_stat_statements
 ----------+-------+-----------------------------------------------------------------
  t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT $1
  t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_2 AS EXECUTE test_prepare_pgss
+ f        |     1 | SELECT $1
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
  f        |     1 | select generate_series($1, $2)
-(4 rows)
+(5 rows)
 
 -- CREATE TABLE AS, top-level tracking.
 SET pg_stat_statements.track = 'top';
@@ -1089,8 +1091,9 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  toplevel | calls |                                   query                                   
 ----------+-------+---------------------------------------------------------------------------
  t        |     1 | EXPLAIN (COSTS OFF) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT $1
+ f        |     1 | SELECT $1
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(2 rows)
+(3 rows)
 
 -- EXPLAIN with CREATE TABLE AS - top-level tracking.
 SET pg_stat_statements.track = 'top';
@@ -1140,8 +1143,9 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | COMMIT
  t        |     1 | DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab
  t        |     1 | FETCH FORWARD 1 FROM foocur
+ f        |     1 | SELECT * from stats_track_tab
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(6 rows)
+(7 rows)
 
 -- DECLARE CURSOR, top-level tracking.
 SET pg_stat_statements.track = 'top';
-- 
2.45.2

Attachment: signature.asc
Description: PGP signature

Reply via email to