v9-0003-Track-location-in-nested-explain-statement.patch

From dbe7d9c8e72a92d1199d5e645ba69501f4253c55 Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Thu, 17 Oct 2024 11:28:23 +0800
Subject: [PATCH v9 3/3] Track location in nested explain statement

---
 .../expected/level_tracking.out               | 92 +++++++++----------
 src/backend/commands/copy.c                   |  2 +-
 src/backend/commands/copyto.c                 |  7 +-
 src/backend/parser/analyze.c                  | 16 ++++
 src/backend/parser/gram.y                     |  6 ++
 src/include/commands/copy.h                   |  3 +-
 src/include/nodes/parsenodes.h                |  3 +
 src/include/parser/parse_node.h               |  2 +
 .../test_copy_callbacks/test_copy_callbacks.c |  2 +-
 9 files changed, 82 insertions(+), 51 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 8c5c1b5b55..36584ddefa 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -221,19 +221,19 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | explain (costs off) TABLE stats_track_tab
  t        |     1 | explain (costs off) UPDATE stats_track_tab SET x=$1 WHERE x=$2
  t        |     1 | explain (costs off) VALUES($1)
- f        |     1 | explain (costs off) (SELECT $1, $2);
- f        |     1 | explain (costs off) (TABLE test_table);
- f        |     1 | explain (costs off) (VALUES($1, $2));
- f        |     1 | explain (costs off) DELETE FROM stats_track_tab;
- f        |     1 | explain (costs off) INSERT INTO stats_track_tab VALUES (($1));
- f        |     1 | explain (costs off) MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($1, $2) id) ON x = id+
+ f        |     1 | (SELECT $1, $2)
+ f        |     1 | (TABLE test_table)
+ f        |     1 | (VALUES($1, $2))
+ f        |     1 | DELETE FROM stats_track_tab
+ f        |     1 | INSERT INTO stats_track_tab VALUES (($1))
+ f        |     1 | MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($1, $2) id) ON x = id                    +
           |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                   +
-          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id);
- f        |     1 | explain (costs off) SELECT $1 UNION SELECT $2;
- f        |     1 | explain (costs off) SELECT $1;
- f        |     1 | explain (costs off) TABLE stats_track_tab;
- f        |     1 | explain (costs off) UPDATE stats_track_tab SET x=$1 WHERE x=$2;
- f        |     1 | explain (costs off) VALUES($1);
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ f        |     1 | SELECT $1
+ f        |     1 | SELECT $1 UNION SELECT $2
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | UPDATE stats_track_tab SET x=$1 WHERE x=$2
+ f        |     1 | VALUES($1)
 (23 rows)
 
 -- Explain - top-level tracking.
@@ -480,8 +480,8 @@ explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY toplevel desc, query COLLATE "C";
- toplevel | calls |                                                              query                                                               
-----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ toplevel | calls |                                                   query                                                    
+----------+-------+------------------------------------------------------------------------------------------------------------
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
  t        |     1 | explain (costs off) (SELECT $1, $2, $3)
  t        |     1 | explain (costs off) (SELECT $1, $2, $3) UNION SELECT $4, $5, $6
@@ -492,8 +492,8 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | explain (costs off) DELETE FROM stats_track_tab WHERE x=$1
  t        |     1 | explain (costs off) INSERT INTO stats_track_tab VALUES ($1), ($2)
  t        |     1 | explain (costs off) INSERT INTO stats_track_tab VALUES (($1))
- t        |     1 | explain (costs off) MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($1, $2) id) ON x = id                      +
-          |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                                         +
+ t        |     1 | explain (costs off) MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($1, $2) id) ON x = id+
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                   +
           |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
  t        |     1 | explain (costs off) SELECT $1
  t        |     1 | explain (costs off) SELECT $1, $2
@@ -503,28 +503,26 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | explain (costs off) UPDATE stats_track_tab SET x=$1
  t        |     1 | explain (costs off) UPDATE stats_track_tab SET x=$1 WHERE x=$2
  t        |     1 | explain (costs off) VALUES($1)
- f        |     1 | explain (costs off) (SELECT $1, $2, $3); explain (costs off) (SELECT 1, 2, 3, 4);
- f        |     1 | explain (costs off) (SELECT 1, 2, 3); explain (costs off) (SELECT $1, $2, $3, $4);
- f        |     1 | explain (costs off) DELETE FROM stats_track_tab; explain (costs off) DELETE FROM stats_track_tab WHERE x=$1;
- f        |     1 | explain (costs off) DELETE FROM stats_track_tab; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
- f        |     1 | explain (costs off) INSERT INTO stats_track_tab VALUES (($1)); explain (costs off) INSERT INTO stats_track_tab VALUES (1), (2);
- f        |     1 | explain (costs off) INSERT INTO stats_track_tab VALUES ((1)); explain (costs off) INSERT INTO stats_track_tab VALUES ($1), ($2);
- f        |     1 | explain (costs off) MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($1, $2) id) ON x = id                      +
-          |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                                         +
-          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id); explain (costs off) SELECT 1, 2, 3, 4, 5;
- f        |     1 | explain (costs off) MERGE INTO stats_track_tab USING (SELECT id FROM generate_series(1, 10) id) ON x = id                       +
-          |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                                         +
-          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id); explain (costs off) SELECT $1, $2, $3, $4, $5;
- f        |     1 | explain (costs off) SELECT $1, $2 UNION SELECT $3, $4; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
- f        |     1 | explain (costs off) SELECT $1; explain (costs off) SELECT 1, 2;
- f        |     1 | explain (costs off) SELECT 1, 2 UNION SELECT 3, 4; explain (costs off) (SELECT $1, $2, $3) UNION SELECT $4, $5, $6;
- f        |     1 | explain (costs off) SELECT 1; explain (costs off) SELECT $1, $2;
- f        |     1 | explain (costs off) TABLE stats_track_tab; explain (costs off) (TABLE test_table);
- f        |     1 | explain (costs off) TABLE stats_track_tab; explain (costs off) (TABLE test_table);
- f        |     1 | explain (costs off) UPDATE stats_track_tab SET x=$1 WHERE x=$2; explain (costs off) UPDATE stats_track_tab SET x=1;
- f        |     1 | explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1; explain (costs off) UPDATE stats_track_tab SET x=$1;
- f        |     1 | explain (costs off) VALUES($1); explain (costs off) (VALUES(1, 2));
- f        |     1 | explain (costs off) VALUES(1); explain (costs off) (VALUES($1, $2));
+ f        |     1 | (SELECT $1, $2, $3)
+ f        |     1 | (SELECT $1, $2, $3) UNION SELECT $4, $5, $6
+ f        |     1 | (SELECT $1, $2, $3, $4)
+ f        |     1 | (TABLE test_table)
+ f        |     1 | (VALUES($1, $2))
+ f        |     1 | DELETE FROM stats_track_tab
+ f        |     1 | DELETE FROM stats_track_tab WHERE x=$1
+ f        |     1 | INSERT INTO stats_track_tab VALUES ($1), ($2)
+ f        |     1 | INSERT INTO stats_track_tab VALUES (($1))
+ f        |     1 | MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($1, $2) id) ON x = id                    +
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                   +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ f        |     1 | SELECT $1
+ f        |     1 | SELECT $1, $2
+ f        |     1 | SELECT $1, $2 UNION SELECT $3, $4
+ f        |     1 | SELECT $1, $2, $3, $4, $5
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | UPDATE stats_track_tab SET x=$1
+ f        |     1 | UPDATE stats_track_tab SET x=$1 WHERE x=$2
+ f        |     1 | VALUES($1)
 (37 rows)
 
 -- Explain - top-level tracking with multi statement.
@@ -762,15 +760,15 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | explain (costs off) WITH a AS (select $1) SELECT $2
  t        |     1 | explain (costs off) WITH a AS (select $1) SELECT $2 UNION SELECT $3
  t        |     1 | explain (costs off) WITH a AS (select $1) UPDATE stats_track_tab SET x=$2 WHERE x=$3
- f        |     1 | explain (costs off) (WITH a AS (select $1) (SELECT $2, $3));
- f        |     1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab;
- f        |     1 | explain (costs off) WITH a AS (select $1) INSERT INTO stats_track_tab VALUES (($2));
- f        |     1 | explain (costs off) WITH a AS (select $1) MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($2, $3) id) ON x = id+
+ f        |     1 | (WITH a AS (select $1) (SELECT $2, $3))
+ f        |     1 | WITH a AS (select $1) DELETE FROM stats_track_tab
+ f        |     1 | WITH a AS (select $1) INSERT INTO stats_track_tab VALUES (($2))
+ f        |     1 | WITH a AS (select $1) MERGE INTO stats_track_tab USING (SELECT id FROM generate_series($2, $3) id) ON x = id                    +
           |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                                         +
-          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id);
- f        |     1 | explain (costs off) WITH a AS (select $1) SELECT $2 UNION SELECT $3;
- f        |     1 | explain (costs off) WITH a AS (select $1) SELECT $2;
- f        |     1 | explain (costs off) WITH a AS (select $1) UPDATE stats_track_tab SET x=$2 WHERE x=$3;
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ f        |     1 | WITH a AS (select $1) SELECT $2
+ f        |     1 | WITH a AS (select $1) SELECT $2 UNION SELECT $3
+ f        |     1 | WITH a AS (select $1) UPDATE stats_track_tab SET x=$2 WHERE x=$3
 (15 rows)
 
 -- Explain with CTE - top-level tracking
@@ -882,7 +880,7 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab
  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 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1;
+ f        |     1 | SELECT $1
 (4 rows)
 
 -- Explain analyze, top tracking.
diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c
index 0b093dbb2a..6794a2f16d 100644
--- a/src/backend/commands/copy.c
+++ b/src/backend/commands/copy.c
@@ -312,7 +312,7 @@ DoCopy(ParseState *pstate, const CopyStmt *stmt,
 
 		cstate = BeginCopyTo(pstate, rel, query, relid,
 							 stmt->filename, stmt->is_program,
-							 NULL, stmt->attlist, stmt->options);
+							 NULL, stmt->attlist, stmt->options, stmt->location, stmt->stmt_len);
 		*processed = DoCopyTo(cstate);	/* copy from database to file */
 		EndCopyTo(cstate);
 	}
diff --git a/src/backend/commands/copyto.c b/src/backend/commands/copyto.c
index 463083e645..a38f688ffa 100644
--- a/src/backend/commands/copyto.c
+++ b/src/backend/commands/copyto.c
@@ -355,7 +355,9 @@ BeginCopyTo(ParseState *pstate,
 			bool is_program,
 			copy_data_dest_cb data_dest_cb,
 			List *attnamelist,
-			List *options)
+			List *options,
+			int	location,
+			int	stmt_len)
 {
 	CopyToState cstate;
 	bool		pipe = (filename == NULL && data_dest_cb == NULL);
@@ -485,6 +487,9 @@ BeginCopyTo(ParseState *pstate,
 
 		query = linitial_node(Query, rewritten);
 
+		query->stmt_location = location;
+		query->stmt_len = stmt_len;
+
 		/* The grammar allows SELECT INTO, but we don't support that */
 		if (query->utilityStmt != NULL &&
 			IsA(query->utilityStmt, CreateTableAsStmt))
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index 2d3d8fcf76..23d61c9837 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -113,6 +113,8 @@ parse_analyze_fixedparams(RawStmt *parseTree, const char *sourceText,
 	Assert(sourceText != NULL); /* required as of 8.4 */
 
 	pstate->p_sourcetext = sourceText;
+	pstate->p_stmt_len = parseTree->stmt_len;
+	pstate->p_stmt_location = parseTree->stmt_location;
 
 	if (numParams > 0)
 		setup_parse_fixed_parameters(pstate, paramTypes, numParams);
@@ -153,6 +155,8 @@ parse_analyze_varparams(RawStmt *parseTree, const char *sourceText,
 	Assert(sourceText != NULL); /* required as of 8.4 */
 
 	pstate->p_sourcetext = sourceText;
+	pstate->p_stmt_len = parseTree->stmt_len;
+	pstate->p_stmt_location = parseTree->stmt_location;
 
 	setup_parse_variable_parameters(pstate, paramTypes, numParams);
 
@@ -195,6 +199,8 @@ parse_analyze_withcb(RawStmt *parseTree, const char *sourceText,
 	Assert(sourceText != NULL); /* required as of 8.4 */
 
 	pstate->p_sourcetext = sourceText;
+	pstate->p_stmt_len = parseTree->stmt_len;
+	pstate->p_stmt_location = parseTree->stmt_location;
 	pstate->p_queryEnv = queryEnv;
 	(*parserSetup) (pstate, parserSetupArg);
 
@@ -2968,6 +2974,7 @@ static Query *
 transformExplainStmt(ParseState *pstate, ExplainStmt *stmt)
 {
 	Query	   *result;
+	Query	   *explained_query;
 	bool		generic_plan = false;
 	Oid		   *paramTypes = NULL;
 	int			numParams = 0;
@@ -2996,6 +3003,15 @@ transformExplainStmt(ParseState *pstate, ExplainStmt *stmt)
 	/* transform contained query, allowing SELECT INTO */
 	stmt->query = (Node *) transformOptionalSelectInto(pstate, stmt->query);
 
+	explained_query = (Query *) stmt->query;
+	explained_query->stmt_location = stmt->location;
+
+	/*
+	 * the being explained query stmt_len is top level query stmt_len minus the
+	 * being EXPLAIN nested query's beginning position.
+	*/
+	explained_query->stmt_len = pstate->p_stmt_location + pstate->p_stmt_len - stmt->location;
+
 	/* make sure all is well with parameter types */
 	if (generic_plan)
 		check_variable_parameters(pstate, (Query *) stmt->query);
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 4bab2117d9..97bc52de3b 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -3390,6 +3390,8 @@ CopyStmt:	COPY opt_binary qualified_name opt_column_list
 					n->is_program = $6;
 					n->filename = $7;
 					n->options = $9;
+					n->location = @3;
+					n->stmt_len = @4 - @3;
 
 					if (n->is_program && n->filename == NULL)
 						ereport(ERROR,
@@ -11958,6 +11960,7 @@ ExplainStmt:
 					ExplainStmt *n = makeNode(ExplainStmt);
 
 					n->query = $2;
+					n->location = @2;
 					n->options = NIL;
 					$$ = (Node *) n;
 				}
@@ -11966,6 +11969,7 @@ ExplainStmt:
 					ExplainStmt *n = makeNode(ExplainStmt);
 
 					n->query = $4;
+					n->location = @4;
 					n->options = list_make1(makeDefElem("analyze", NULL, @2));
 					if ($3)
 						n->options = lappend(n->options,
@@ -11977,6 +11981,7 @@ ExplainStmt:
 					ExplainStmt *n = makeNode(ExplainStmt);
 
 					n->query = $3;
+					n->location = @3;
 					n->options = list_make1(makeDefElem("verbose", NULL, @2));
 					$$ = (Node *) n;
 				}
@@ -11985,6 +11990,7 @@ ExplainStmt:
 					ExplainStmt *n = makeNode(ExplainStmt);
 
 					n->query = $5;
+					n->location = @5;
 					n->options = $3;
 					$$ = (Node *) n;
 				}
diff --git a/src/include/commands/copy.h b/src/include/commands/copy.h
index 4002a7f538..0440ff2345 100644
--- a/src/include/commands/copy.h
+++ b/src/include/commands/copy.h
@@ -121,7 +121,8 @@ extern DestReceiver *CreateCopyDestReceiver(void);
  */
 extern CopyToState BeginCopyTo(ParseState *pstate, Relation rel, RawStmt *raw_query,
 							   Oid queryRelId, const char *filename, bool is_program,
-							   copy_data_dest_cb data_dest_cb, List *attnamelist, List *options);
+							   copy_data_dest_cb data_dest_cb, List *attnamelist, List *options,
+							   int location, int stmt_len);
 extern void EndCopyTo(CopyToState cstate);
 extern uint64 DoCopyTo(CopyToState cstate);
 extern List *CopyGetAttnums(TupleDesc tupDesc, Relation rel,
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index c92cef3d16..bb22ecc6b3 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -2596,6 +2596,8 @@ typedef struct CopyStmt
 	char	   *filename;		/* filename, or NULL for STDIN/STDOUT */
 	List	   *options;		/* List of DefElem nodes */
 	Node	   *whereClause;	/* WHERE condition (or NULL) */
+	ParseLoc	location;		/* the nest query location. COPY TO only */
+	int			stmt_len;		/* the nest query length. COPY TO only */
 } CopyStmt;
 
 /* ----------------------
@@ -3887,6 +3889,7 @@ typedef struct ExplainStmt
 	NodeTag		type;
 	Node	   *query;			/* the query (see comments above) */
 	List	   *options;		/* list of DefElem nodes */
+	ParseLoc	location;		/* location of the statement being explained */
 } ExplainStmt;
 
 /* ----------------------
diff --git a/src/include/parser/parse_node.h b/src/include/parser/parse_node.h
index 543df56814..ba572b3aea 100644
--- a/src/include/parser/parse_node.h
+++ b/src/include/parser/parse_node.h
@@ -193,6 +193,8 @@ struct ParseState
 {
 	ParseState *parentParseState;	/* stack link */
 	const char *p_sourcetext;	/* source text, or NULL if not available */
+	ParseLoc	p_stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	p_stmt_len;		/* length in bytes; 0 means "rest of string" */
 	List	   *p_rtable;		/* range table so far */
 	List	   *p_rteperminfos; /* list of RTEPermissionInfo nodes for each
 								 * RTE_RELATION entry in rtable */
diff --git a/src/test/modules/test_copy_callbacks/test_copy_callbacks.c b/src/test/modules/test_copy_callbacks/test_copy_callbacks.c
index 0bbd2aa6bb..453e53e466 100644
--- a/src/test/modules/test_copy_callbacks/test_copy_callbacks.c
+++ b/src/test/modules/test_copy_callbacks/test_copy_callbacks.c
@@ -38,7 +38,7 @@ test_copy_to_callback(PG_FUNCTION_ARGS)
 	int64		processed;
 
 	cstate = BeginCopyTo(NULL, rel, NULL, RelationGetRelid(rel), NULL, false,
-						 to_cb, NIL, NIL);
+						 to_cb, NIL, NIL, 0, 0);
 	processed = DoCopyTo(cstate);
 	EndCopyTo(cstate);
 
-- 
2.34.1