Set query_id for query contained in utility statement

From fb57a065e4ad9feb3792680e7173ba6cff619d0e Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
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 and DeclareCursor
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
contained query but CreateTableAs and DeclareCursor 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. Additionally, explain's nested
query can itself be a CreateTableAs or DeclareCursor which also needs to
be jumbled which is now done in ExplainQuery.
---
 .../expected/level_tracking.out               | 212 ++++++++++++++++++
 .../pg_stat_statements/sql/level_tracking.sql |  71 ++++++
 src/backend/commands/createas.c               |  10 +
 src/backend/commands/explain.c                |  21 +-
 src/backend/commands/portalcmds.c             |  10 +
 src/test/regress/expected/explain.out         |  17 ++
 src/test/regress/sql/explain.sql              |   4 +
 7 files changed, 344 insertions(+), 1 deletion(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index d8dd8a2deea..50f81a805c9 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -112,6 +112,218 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (2 rows)
 
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_test AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_test AS SELECT $1
+ f        |     1 | CREATE TEMPORARY TABLE pgss_test AS SELECT $1;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Create Table As using prepared stmt, all-level tracking.
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_test2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                             query                              
+----------+-------+----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_test2 AS EXECUTE test_prepare_pgss
+ f        |     1 | PREPARE test_prepare_pgss AS select generate_series(1, 10)
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                          query                           
+----------+-------+----------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ t        |     1 | DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab
+ f        |     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
+(7 rows)
+
+-- Explain analyze, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                              query                               
+----------+-------+------------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1
+ f        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Explain analyze with declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                                     query                                                      
+----------+-------+----------------------------------------------------------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab
+ f        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TABLE pgss_test_2 AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | CREATE TABLE pgss_test_2 AS SELECT $1
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Create Table As using prepared stmt, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_test3 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                             query                              
+----------+-------+----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_test3 AS EXECUTE test_prepare_pgss
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                          query                          
+----------+-------+---------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ 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)
+
+-- Explain analyze, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Explain analyze with declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                                     query                                                     
+----------+-------+---------------------------------------------------------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
 -- 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 65a17147a5a..385785aabbb 100644
--- a/contrib/pg_stat_statements/sql/level_tracking.sql
+++ b/contrib/pg_stat_statements/sql/level_tracking.sql
@@ -55,6 +55,77 @@ CALL proc_with_utility_stmt();
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
 
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_test AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Create Table As using prepared stmt, all-level tracking.
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_test2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze with declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TABLE pgss_test_2 AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Create Table As using prepared stmt, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_test3 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze with declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
 -- 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/commands/createas.c b/src/backend/commands/createas.c
index c71ff801888..dc28cb9188c 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 5771aabf40a..4a922481583 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -308,8 +308,27 @@ ExplainQuery(ParseState *pstate, ExplainStmt *stmt,
 
 	query = castNode(Query, stmt->query);
 	if (IsQueryIdEnabled())
+	{
 		jstate = JumbleQuery(query);
 
+		/*
+		 * explain can contain an utility statement with a plannable query, we
+		 * need to jumble it to set the nested query's id
+		 */
+		if (query->utilityStmt != NULL)
+		{
+			Query	   *nested_query = NULL;
+			Node	   *utilityStmt = query->utilityStmt;
+
+			if (IsA(utilityStmt, CreateTableAsStmt))
+				nested_query = castNode(Query, ((CreateTableAsStmt *) utilityStmt)->query);
+			else if (IsA(utilityStmt, DeclareCursorStmt))
+				nested_query = castNode(Query, ((DeclareCursorStmt *) utilityStmt)->query);
+			if (nested_query)
+				JumbleQuery(nested_query);
+		}
+	}
+
 	if (post_parse_analyze_hook)
 		(*post_parse_analyze_hook) (pstate, query, jstate);
 
@@ -319,7 +338,7 @@ ExplainQuery(ParseState *pstate, ExplainStmt *stmt,
 	 * came straight from the parser, or suitable locks were acquired by
 	 * plancache.c.
 	 */
-	rewritten = QueryRewrite(castNode(Query, stmt->query));
+	rewritten = QueryRewrite(query);
 
 	/* emit opening boilerplate */
 	ExplainBeginOutput(es);
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/test/regress/expected/explain.out b/src/test/regress/expected/explain.out
index 6585c6a69ef..e01ed569989 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 c7055f850c5..5796137c7ff 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.3 (Apple Git-146)

From 8719d2184affb9ecd69b4c1e7c204108b76d1b76 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
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 and DeclareCursor
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
contained query but CreateTableAs and DeclareCursor 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. Additionally, explain's nested
query can itself be a CreateTableAs or DeclareCursor which also needs to
be jumbled which is now done in ExplainOneUtility.
---
 .../expected/level_tracking.out               | 255 ++++++++++++++++++
 .../pg_stat_statements/sql/level_tracking.sql |  83 ++++++
 src/backend/commands/createas.c               |  10 +
 src/backend/commands/explain.c                |  45 ++--
 src/backend/commands/portalcmds.c             |  10 +
 src/backend/commands/prepare.c                |  18 +-
 src/include/commands/explain.h                |   4 +-
 src/include/commands/prepare.h                |   4 +-
 src/test/regress/expected/explain.out         |  17 ++
 src/test/regress/sql/explain.sql              |   4 +
 10 files changed, 415 insertions(+), 35 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index d8dd8a2deea..6291f06d711 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -112,6 +112,261 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (2 rows)
 
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_test AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_test AS SELECT $1
+ f        |     1 | CREATE TEMPORARY TABLE pgss_test AS SELECT $1;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Create Table As using prepared stmt, all-level tracking.
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_test2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                             query                              
+----------+-------+----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_test2 AS EXECUTE test_prepare_pgss
+ f        |     1 | PREPARE test_prepare_pgss AS select generate_series(1, 10)
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                          query                           
+----------+-------+----------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ t        |     1 | DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab
+ f        |     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
+(7 rows)
+
+-- Explain analyze, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                              query                               
+----------+-------+------------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1
+ f        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Explain analyze with declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                                     query                                                      
+----------+-------+----------------------------------------------------------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab
+ f        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Explain with ctas, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                        query                                        
+----------+-------+-------------------------------------------------------------------------------------
+ t        |     1 | EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT $1
+ f        |     1 | EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT $1;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TABLE pgss_test_4 AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | CREATE TABLE pgss_test_4 AS SELECT $1
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Create Table As using prepared stmt, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_test5 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                             query                              
+----------+-------+----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_test5 AS EXECUTE test_prepare_pgss
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                          query                          
+----------+-------+---------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ 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)
+
+-- Explain analyze, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Explain analyze with declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                                     query                                                     
+----------+-------+---------------------------------------------------------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Explain with ctas, top-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                       query                                        
+----------+-------+------------------------------------------------------------------------------------
+ t        |     1 | EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT $1
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
 -- 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 65a17147a5a..2fc7054a9c6 100644
--- a/contrib/pg_stat_statements/sql/level_tracking.sql
+++ b/contrib/pg_stat_statements/sql/level_tracking.sql
@@ -55,6 +55,89 @@ CALL proc_with_utility_stmt();
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
 
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_test AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Create Table As using prepared stmt, all-level tracking.
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_test2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze with declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain with ctas, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TABLE pgss_test_4 AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Create Table As using prepared stmt, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_test5 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze with declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain with ctas, top-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
 -- 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/commands/createas.c b/src/backend/commands/createas.c
index 0b629b1f79c..90bb964cd89 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 11df4a04d43..cd42fb1b5a0 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,
@@ -319,7 +318,7 @@ ExplainQuery(ParseState *pstate, ExplainStmt *stmt,
 	 * came straight from the parser, or suitable locks were acquired by
 	 * plancache.c.
 	 */
-	rewritten = QueryRewrite(castNode(Query, stmt->query));
+	rewritten = QueryRewrite(query);
 
 	/* emit opening boilerplate */
 	ExplainBeginOutput(es);
@@ -342,7 +341,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)
@@ -428,24 +427,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);
 }
 
 /*
@@ -526,9 +523,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;
 
@@ -539,6 +537,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;
 
 		/*
@@ -557,11 +556,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))
 	{
@@ -574,17 +578,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 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..e141ecf9577 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -566,11 +566,9 @@ DropAllPreparedStatements(void)
  */
 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;
 	CachedPlan *cplan;
 	List	   *plan_list;
 	ListCell   *p;
@@ -605,16 +603,9 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 	if (!entry->plansource->fixed_result)
 		elog(ERROR, "EXPLAIN EXECUTE does not support variable-result cached plans");
 
-	query_string = entry->plansource->query_string;
-
 	/* Evaluate parameters, if any */
 	if (entry->plansource->num_params)
 	{
-		ParseState *pstate;
-
-		pstate = make_parsestate(NULL);
-		pstate->p_sourcetext = queryString;
-
 		/*
 		 * Need an EState to evaluate parameters; must not delete it till end
 		 * of query, in case parameters are pass-by-reference.  Note that the
@@ -629,7 +620,7 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 
 	/* 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 +646,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, pstate->p_sourcetext, 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/include/commands/explain.h b/src/include/commands/explain.h
index 9b8b351d9a2..0fd3933a17a 100644
--- a/src/include/commands/explain.h
+++ b/src/include/commands/explain.h
@@ -98,8 +98,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 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 6585c6a69ef..e01ed569989 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 c7055f850c5..5796137c7ff 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.3 (Apple Git-146)

From ab4ebed0c5bef271c5e2c56bc61a6d6053834ffd Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
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 and DeclareCursor
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
contained query but CreateTableAs and DeclareCursor 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. Additionally, explain's nested
query can itself be a CreateTableAs or DeclareCursor which also needs to
be jumbled. This is now done in ExplainOneUtility.
---
 .../expected/level_tracking.out               | 255 ++++++++++++++++++
 .../pg_stat_statements/sql/level_tracking.sql |  83 ++++++
 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/include/commands/explain.h                |   4 +-
 src/include/commands/prepare.h                |   4 +-
 src/test/regress/expected/explain.out         |  17 ++
 src/test/regress/sql/explain.sql              |   4 +
 10 files changed, 419 insertions(+), 31 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index d8dd8a2deea..6291f06d711 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -112,6 +112,261 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (2 rows)
 
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_test AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_test AS SELECT $1
+ f        |     1 | CREATE TEMPORARY TABLE pgss_test AS SELECT $1;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Create Table As using prepared stmt, all-level tracking.
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_test2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                             query                              
+----------+-------+----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_test2 AS EXECUTE test_prepare_pgss
+ f        |     1 | PREPARE test_prepare_pgss AS select generate_series(1, 10)
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                          query                           
+----------+-------+----------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ t        |     1 | DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab
+ f        |     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
+(7 rows)
+
+-- Explain analyze, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                              query                               
+----------+-------+------------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1
+ f        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Explain analyze with declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                                     query                                                      
+----------+-------+----------------------------------------------------------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab
+ f        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Explain with ctas, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                        query                                        
+----------+-------+-------------------------------------------------------------------------------------
+ t        |     1 | EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT $1
+ f        |     1 | EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT $1;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TABLE pgss_test_4 AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | CREATE TABLE pgss_test_4 AS SELECT $1
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Create Table As using prepared stmt, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_test5 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                             query                              
+----------+-------+----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_test5 AS EXECUTE test_prepare_pgss
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                          query                          
+----------+-------+---------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ 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)
+
+-- Explain analyze, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Explain analyze with declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                                     query                                                     
+----------+-------+---------------------------------------------------------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Explain with ctas, top-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                       query                                        
+----------+-------+------------------------------------------------------------------------------------
+ t        |     1 | EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT $1
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
 -- 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 65a17147a5a..2fc7054a9c6 100644
--- a/contrib/pg_stat_statements/sql/level_tracking.sql
+++ b/contrib/pg_stat_statements/sql/level_tracking.sql
@@ -55,6 +55,89 @@ CALL proc_with_utility_stmt();
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
 
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_test AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Create Table As using prepared stmt, all-level tracking.
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_test2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze with declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain with ctas, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TABLE pgss_test_4 AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Create Table As using prepared stmt, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_test5 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze with declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain with ctas, top-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
 -- 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/commands/createas.c b/src/backend/commands/createas.c
index 0b629b1f79c..90bb964cd89 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 11df4a04d43..81e460000a3 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,
@@ -342,7 +341,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)
@@ -428,24 +427,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);
 }
 
 /*
@@ -526,9 +523,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;
 
@@ -539,6 +537,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;
 
 		/*
@@ -557,11 +556,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))
 	{
@@ -574,17 +578,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 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/include/commands/explain.h b/src/include/commands/explain.h
index 9b8b351d9a2..0fd3933a17a 100644
--- a/src/include/commands/explain.h
+++ b/src/include/commands/explain.h
@@ -98,8 +98,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 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 6585c6a69ef..e01ed569989 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 c7055f850c5..5796137c7ff 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.3 (Apple Git-146)

From a36066d24611528cb319f9d929d4278ec37c7d21 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Thu, 3 Oct 2024 08:52:02 +0200
Subject: Track location to extract relevant part in nested statement

Previously, Query generated through transform would have unset
stmt_location. Extensions relying on the statement location to extract
the relevant part of the statement would fallback to use the whole
statement instead, thus showing the same string in the top and
nested level which was a source of confusion.

This patch fixes the issue by keeping track of the statement locations
and propagate it to Query during transform, allowing pgss to only show
the relevant part of the query for nested query.
---
 .../expected/level_tracking.out               | 159 ++++++++++++++++++
 .../pg_stat_statements/sql/level_tracking.sql |  26 +++
 src/backend/parser/analyze.c                  |  10 ++
 src/backend/parser/gram.y                     |  17 +-
 src/backend/parser/parse_merge.c              |   2 +
 src/include/nodes/parsenodes.h                |   5 +
 6 files changed, 214 insertions(+), 5 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index bb65e98ce09..5649c0b7c10 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -67,6 +67,165 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | SET pg_stat_statements.track = $1
 (7 rows)
 
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+-- Explain - all-level tracking.
+SET pg_stat_statements.track = 'all';
+explain (costs off) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+-- Check we correctly capture substring with CTE
+explain (costs off) WITH a AS (select 4) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) 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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                                              query                                                               
+----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ 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 | SELECT $1 UNION SELECT $2;
+ f        |     1 | SELECT $1;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | SET pg_stat_statements.track = $1
+ f        |     1 | UPDATE stats_track_tab SET x=$1 WHERE x=$2;
+ 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 | WITH a AS (select $1) SELECT $2 UNION SELECT $3;
+ f        |     1 | WITH a AS (select $1) SELECT $2;
+ f        |     1 | WITH a AS (select $1) UPDATE stats_track_tab SET x=$2 WHERE x=$3;
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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                                                                                         +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1 UNION SELECT $2
+ t        |     1 | explain (costs off) UPDATE stats_track_tab SET x=$1 WHERE x=$2
+ t        |     1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab
+ t        |     1 | explain (costs off) WITH a AS (select $1) INSERT INTO stats_track_tab VALUES (($2))
+ t        |     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+
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                                         +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ 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
+(26 rows)
+
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
 -- Procedure with multiple utility statements.
 CREATE OR REPLACE PROCEDURE proc_with_utility_stmt()
 LANGUAGE SQL
diff --git a/contrib/pg_stat_statements/sql/level_tracking.sql b/contrib/pg_stat_statements/sql/level_tracking.sql
index 65a17147a5a..1b431c9928c 100644
--- a/contrib/pg_stat_statements/sql/level_tracking.sql
+++ b/contrib/pg_stat_statements/sql/level_tracking.sql
@@ -32,6 +32,32 @@ BEGIN
 END; $$;
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+
+-- Explain - all-level tracking.
+SET pg_stat_statements.track = 'all';
+explain (costs off) SELECT 1;
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) DELETE FROM stats_track_tab;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((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 UNION SELECT 2;
+
+-- Check we correctly capture substring with CTE
+explain (costs off) WITH a AS (select 4) SELECT 1;
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+explain (costs off) WITH a AS (select 4) 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) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
 
 -- Procedure with multiple utility statements.
 CREATE OR REPLACE PROCEDURE proc_with_utility_stmt()
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index e901203424d..1cf489eb3a0 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -518,6 +518,7 @@ transformDeleteStmt(ParseState *pstate, DeleteStmt *stmt)
 	Node	   *qual;
 
 	qry->commandType = CMD_DELETE;
+	qry->stmt_location = stmt->location;
 
 	/* process the WITH clause independently of all else */
 	if (stmt->withClause)
@@ -525,6 +526,7 @@ transformDeleteStmt(ParseState *pstate, DeleteStmt *stmt)
 		qry->hasRecursive = stmt->withClause->recursive;
 		qry->cteList = transformWithClause(pstate, stmt->withClause);
 		qry->hasModifyingCTE = pstate->p_hasModifyingCTE;
+		qry->stmt_location = stmt->withClause->location;
 	}
 
 	/* set up range table with just the result rel */
@@ -606,6 +608,7 @@ transformInsertStmt(ParseState *pstate, InsertStmt *stmt)
 	Assert(pstate->p_ctenamespace == NIL);
 
 	qry->commandType = CMD_INSERT;
+	qry->stmt_location = stmt->location;
 	pstate->p_is_insert = true;
 
 	/* process the WITH clause independently of all else */
@@ -614,6 +617,7 @@ transformInsertStmt(ParseState *pstate, InsertStmt *stmt)
 		qry->hasRecursive = stmt->withClause->recursive;
 		qry->cteList = transformWithClause(pstate, stmt->withClause);
 		qry->hasModifyingCTE = pstate->p_hasModifyingCTE;
+		qry->stmt_location = stmt->withClause->location;
 	}
 
 	qry->override = stmt->override;
@@ -1347,6 +1351,7 @@ transformSelectStmt(ParseState *pstate, SelectStmt *stmt)
 	ListCell   *l;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->location;
 
 	/* process the WITH clause independently of all else */
 	if (stmt->withClause)
@@ -1354,6 +1359,7 @@ transformSelectStmt(ParseState *pstate, SelectStmt *stmt)
 		qry->hasRecursive = stmt->withClause->recursive;
 		qry->cteList = transformWithClause(pstate, stmt->withClause);
 		qry->hasModifyingCTE = pstate->p_hasModifyingCTE;
+		qry->stmt_location = stmt->withClause->location;
 	}
 
 	/* Complain if we get called from someplace where INTO is not allowed */
@@ -1730,6 +1736,7 @@ transformSetOperationStmt(ParseState *pstate, SelectStmt *stmt)
 	int			tllen;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->location;
 
 	/*
 	 * Find leftmost leaf SelectStmt.  We currently only need to do this in
@@ -1784,6 +1791,7 @@ transformSetOperationStmt(ParseState *pstate, SelectStmt *stmt)
 		qry->hasRecursive = withClause->recursive;
 		qry->cteList = transformWithClause(pstate, withClause);
 		qry->hasModifyingCTE = pstate->p_hasModifyingCTE;
+		qry->stmt_location = withClause->location;
 	}
 
 	/*
@@ -2429,6 +2437,7 @@ transformUpdateStmt(ParseState *pstate, UpdateStmt *stmt)
 	Node	   *qual;
 
 	qry->commandType = CMD_UPDATE;
+	qry->stmt_location = stmt->location;
 	pstate->p_is_insert = false;
 
 	/* process the WITH clause independently of all else */
@@ -2437,6 +2446,7 @@ transformUpdateStmt(ParseState *pstate, UpdateStmt *stmt)
 		qry->hasRecursive = stmt->withClause->recursive;
 		qry->cteList = transformWithClause(pstate, stmt->withClause);
 		qry->hasModifyingCTE = pstate->p_hasModifyingCTE;
+		qry->stmt_location = stmt->withClause->location;
 	}
 
 	qry->resultRelation = setTargetTable(pstate, stmt->relation,
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 4aa8646af7b..dcb2588b396 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -190,7 +190,7 @@ static void insertSelectOptions(SelectStmt *stmt,
 								SelectLimit *limitClause,
 								WithClause *withClause,
 								core_yyscan_t yyscanner);
-static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg);
+static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location);
 static Node *doNegate(Node *n, int location);
 static void doNegateFloat(Float *v);
 static Node *makeAndExpr(Node *lexpr, Node *rexpr, int location);
@@ -12170,6 +12170,7 @@ InsertStmt:
 					$5->onConflictClause = $6;
 					$5->returningList = $7;
 					$5->withClause = $1;
+					$5->location = @2;
 					$$ = (Node *) $5;
 				}
 		;
@@ -12323,6 +12324,7 @@ DeleteStmt: opt_with_clause DELETE_P FROM relation_expr_opt_alias
 					n->whereClause = $6;
 					n->returningList = $7;
 					n->withClause = $1;
+					n->location = @2;
 					$$ = (Node *) n;
 				}
 		;
@@ -12397,6 +12399,7 @@ UpdateStmt: opt_with_clause UPDATE relation_expr_opt_alias
 					n->whereClause = $7;
 					n->returningList = $8;
 					n->withClause = $1;
+					n->location = @2;
 					$$ = (Node *) n;
 				}
 		;
@@ -12474,6 +12477,7 @@ MergeStmt:
 					m->joinCondition = $8;
 					m->mergeWhenClauses = $9;
 					m->returningList = $10;
+					m->location = @2;
 
 					$$ = (Node *) m;
 				}
@@ -12836,6 +12840,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->location = @1;
 					$$ = (Node *) n;
 				}
 			| SELECT distinct_clause target_list
@@ -12853,6 +12858,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->location = @1;
 					$$ = (Node *) n;
 				}
 			| values_clause							{ $$ = $1; }
@@ -12877,15 +12883,15 @@ simple_select:
 				}
 			| select_clause UNION set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause INTERSECT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause EXCEPT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 		;
 
@@ -18967,7 +18973,7 @@ insertSelectOptions(SelectStmt *stmt,
 }
 
 static Node *
-makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
+makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location)
 {
 	SelectStmt *n = makeNode(SelectStmt);
 
@@ -18975,6 +18981,7 @@ makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
 	n->all = all;
 	n->larg = (SelectStmt *) larg;
 	n->rarg = (SelectStmt *) rarg;
+	n->location = location;
 	return (Node *) n;
 }
 
diff --git a/src/backend/parser/parse_merge.c b/src/backend/parser/parse_merge.c
index 87df79027d7..2389cc00513 100644
--- a/src/backend/parser/parse_merge.c
+++ b/src/backend/parser/parse_merge.c
@@ -118,6 +118,7 @@ transformMergeStmt(ParseState *pstate, MergeStmt *stmt)
 	Assert(pstate->p_ctenamespace == NIL);
 
 	qry->commandType = CMD_MERGE;
+	qry->stmt_location = stmt->location;
 	qry->hasRecursive = false;
 
 	/* process the WITH clause independently of all else */
@@ -130,6 +131,7 @@ transformMergeStmt(ParseState *pstate, MergeStmt *stmt)
 
 		qry->cteList = transformWithClause(pstate, stmt->withClause);
 		qry->hasModifyingCTE = pstate->p_hasModifyingCTE;
+		qry->stmt_location = stmt->withClause->location;
 	}
 
 	/*
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 1c314cd9074..aa73e2b2428 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -2045,6 +2045,7 @@ typedef struct InsertStmt
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
 	OverridingKind override;	/* OVERRIDING clause */
+	ParseLoc	location;		/* token location, or -1 if unknown */
 } InsertStmt;
 
 /* ----------------------
@@ -2059,6 +2060,7 @@ typedef struct DeleteStmt
 	Node	   *whereClause;	/* qualifications */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	location;		/* token location, or -1 if unknown */
 } DeleteStmt;
 
 /* ----------------------
@@ -2074,6 +2076,7 @@ typedef struct UpdateStmt
 	List	   *fromClause;		/* optional from clause for more tables */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	location;		/* token location, or -1 if unknown */
 } UpdateStmt;
 
 /* ----------------------
@@ -2089,6 +2092,7 @@ typedef struct MergeStmt
 	List	   *mergeWhenClauses;	/* list of MergeWhenClause(es) */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	location;		/* token location, or -1 if unknown */
 } MergeStmt;
 
 /* ----------------------
@@ -2158,6 +2162,7 @@ typedef struct SelectStmt
 	bool		all;			/* ALL specified? */
 	struct SelectStmt *larg;	/* left child */
 	struct SelectStmt *rarg;	/* right child */
+	ParseLoc	location;		/* name's token location */
 	/* Eventually add fields for CORRESPONDING spec here */
 } SelectStmt;
 
-- 
2.39.3 (Apple Git-146)

From 74d897c037b6dcdf642a93d4d8e2aeadf80beb97 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Thu, 3 Oct 2024 10:21:44 +0200
Subject: Use view's definition as query string on a materialized view refresh

When creating a materialized view, the first refresh will have the
"Select" part of the statement as a query string. On subsequent refresh,
the "REFRESH MATERIALIZED" utility statement will be passed as query
string. This causes pgss to track both the top query and nested query as
a refresh.

This patch changes the query string on a refresh to fetch the view
definition instead. This will allow pgss to display the correct
statement when tracking refresh's nested query.
---
 .../expected/level_tracking.out               | 38 +++++++++++++++++++
 .../pg_stat_statements/sql/level_tracking.sql | 14 +++++++
 src/backend/commands/matview.c                | 23 +++++++----
 src/backend/utils/adt/ruleutils.c             | 13 +++++++
 src/include/utils/ruleutils.h                 |  2 +-
 5 files changed, 82 insertions(+), 8 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index c6b011caa42..06e64dab346 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -271,6 +271,44 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (2 rows)
 
+-- Create Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                             query                                              
+----------+-------+------------------------------------------------------------------------------------------------
+ t        |     1 | CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series($1, $2) as id
+ f        |     1 | SELECT * FROM generate_series($1, $2) as id;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Refresh Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
+ t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
+ f        |     1 | SELECT * FROM pg_catalog.pg_rewrite WHERE ev_class = $1 AND rulename = $2
+ f        |     1 | SELECT id                                                                +
+          |       |    FROM generate_series(1, 5) id(id);
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(4 rows)
+
 -- Create Table As, all-level tracking.
 SET pg_stat_statements.track = 'all';
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
diff --git a/contrib/pg_stat_statements/sql/level_tracking.sql b/contrib/pg_stat_statements/sql/level_tracking.sql
index accd4dfce26..f1af0cfd743 100644
--- a/contrib/pg_stat_statements/sql/level_tracking.sql
+++ b/contrib/pg_stat_statements/sql/level_tracking.sql
@@ -81,6 +81,20 @@ CALL proc_with_utility_stmt();
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
 
+-- Create Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Refresh Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
 -- Create Table As, all-level tracking.
 SET pg_stat_statements.track = 'all';
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
diff --git a/src/backend/commands/matview.c b/src/backend/commands/matview.c
index 7cc68338837..784fd9b0425 100644
--- a/src/backend/commands/matview.c
+++ b/src/backend/commands/matview.c
@@ -39,6 +39,7 @@
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
+#include "utils/ruleutils.h"
 #include "utils/snapmgr.h"
 #include "utils/syscache.h"
 
@@ -61,8 +62,7 @@ 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,
-									   ParseState *pstate, const char *queryString,
-									   bool is_create);
+									   ParseState *pstate, 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);
@@ -327,11 +327,20 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData,
 	if (!skipData)
 	{
 		DestReceiver *dest;
-		const char *queryString = pstate->p_sourcetext;
+		ParseState *refresh_pstate = pstate;
+
+		/*
+		 * On refresh, the queryString will be the refresh utility statement.
+		 * We need to fetch the query string from the view definition.
+		 */
+		if (!is_create)
+		{
+			refresh_pstate = make_parsestate(NULL);
+			refresh_pstate->p_sourcetext = pg_get_viewdef_string(matviewOid, false);
+		}
 
 		dest = CreateTransientRelDestReceiver(OIDNewHeap);
-		processed = refresh_matview_datafill(dest, dataQuery, pstate,
-											 queryString, is_create);
+		processed = refresh_matview_datafill(dest, dataQuery, refresh_pstate, is_create);
 	}
 
 	/* Make the matview match the newly generated data. */
@@ -406,8 +415,7 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData,
  */
 static uint64
 refresh_matview_datafill(DestReceiver *dest, Query *query,
-						 ParseState *pstate, const char *queryString,
-						 bool is_create)
+						 ParseState *pstate, bool is_create)
 {
 	List	   *rewritten;
 	PlannedStmt *plan;
@@ -415,6 +423,7 @@ refresh_matview_datafill(DestReceiver *dest, Query *query,
 	Query	   *copied_query;
 	uint64		processed;
 	JumbleState *jstate = NULL;
+	const char *queryString = pstate->p_sourcetext;
 
 	/* Lock and rewrite, using a copy to preserve the original query. */
 	copied_query = copyObject(query);
diff --git a/src/backend/utils/adt/ruleutils.c b/src/backend/utils/adt/ruleutils.c
index 2177d17e278..0bf85cbb759 100644
--- a/src/backend/utils/adt/ruleutils.c
+++ b/src/backend/utils/adt/ruleutils.c
@@ -776,6 +776,19 @@ pg_get_viewdef_name_ext(PG_FUNCTION_ARGS)
 	PG_RETURN_TEXT_P(string_to_text(res));
 }
 
+/*
+ * Internal version of pg_get_viewdef
+ */
+char *
+pg_get_viewdef_string(Oid viewoid, bool pretty)
+{
+	int			prettyFlags;
+
+	prettyFlags = GET_PRETTY_FLAGS(pretty);
+
+	return pg_get_viewdef_worker(viewoid, prettyFlags, WRAP_COLUMN_DEFAULT);
+}
+
 /*
  * Common code for by-OID and by-name variants of pg_get_viewdef
  */
diff --git a/src/include/utils/ruleutils.h b/src/include/utils/ruleutils.h
index 161fb5ef02c..72177b9dce3 100644
--- a/src/include/utils/ruleutils.h
+++ b/src/include/utils/ruleutils.h
@@ -29,7 +29,7 @@ extern char *pg_get_indexdef_columns(Oid indexrelid, bool pretty);
 extern char *pg_get_indexdef_columns_extended(Oid indexrelid,
 											  bits16 flags);
 extern char *pg_get_querydef(Query *query, bool pretty);
-
+extern char *pg_get_viewdef_string(Oid viewoid, bool pretty);
 extern char *pg_get_partkeydef_columns(Oid relid, bool pretty);
 extern char *pg_get_partconstrdef_string(Oid partitionId, char *aliasname);
 
-- 
2.39.3 (Apple Git-146)

From 89d2ba4873836f7a85f8064425921c47e2fb17bc Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
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               | 255 ++++++++++++++++++
 .../pg_stat_statements/sql/level_tracking.sql |  83 ++++++
 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 +
 13 files changed, 442 insertions(+), 42 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 5649c0b7c10..c6b011caa42 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -271,6 +271,261 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (2 rows)
 
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_test AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_test AS SELECT $1
+ f        |     1 | SELECT $1;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Create Table As using prepared stmt, all-level tracking.
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_test2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                             query                              
+----------+-------+----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_test2 AS EXECUTE test_prepare_pgss
+ f        |     1 | PREPARE test_prepare_pgss AS select generate_series(1, 10)
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                          query                          
+----------+-------+---------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ 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
+(7 rows)
+
+-- Explain analyze, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1
+ f        |     1 | SELECT $1;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Explain analyze with declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                                     query                                                     
+----------+-------+---------------------------------------------------------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab
+ f        |     1 | SELECT * FROM stats_track_tab;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Explain with ctas, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                       query                                        
+----------+-------+------------------------------------------------------------------------------------
+ t        |     1 | EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT $1
+ f        |     1 | SELECT $1;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TABLE pgss_test_4 AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | CREATE TABLE pgss_test_4 AS SELECT $1
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Create Table As using prepared stmt, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_test5 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                             query                              
+----------+-------+----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_test5 AS EXECUTE test_prepare_pgss
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                          query                          
+----------+-------+---------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ 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)
+
+-- Explain analyze, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Explain analyze with declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                                     query                                                     
+----------+-------+---------------------------------------------------------------------------------------------------------------
+ t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Explain with ctas, top-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+ toplevel | calls |                                       query                                        
+----------+-------+------------------------------------------------------------------------------------
+ t        |     1 | EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT $1
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
 -- 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 1b431c9928c..accd4dfce26 100644
--- a/contrib/pg_stat_statements/sql/level_tracking.sql
+++ b/contrib/pg_stat_statements/sql/level_tracking.sql
@@ -81,6 +81,89 @@ CALL proc_with_utility_stmt();
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
 
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_test AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Create Table As using prepared stmt, all-level tracking.
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_test2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze with declare cursor, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain with ctas, all-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TABLE pgss_test_4 AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Create Table As using prepared stmt, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_test5 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain analyze with declare cursor, top tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
+-- Explain with ctas, top-level tracking.
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (COSTS OFF, SUMMARY OFF, TIMING OFF) CREATE TABLE pgss_test_3 AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C", toplevel;
+
 -- 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/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.3 (Apple Git-146)

diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index ca8e9201f7..7adf0d0975 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -3004,6 +3004,12 @@ transformExplainStmt(ParseState *pstate, ExplainStmt *stmt)
 			setup_parse_variable_parameters(pstate, &paramTypes, &numParams);
 	}
 
+	if (IsA(stmt->query, SelectStmt))
+	{
+		SelectStmt *sel_stmt = (SelectStmt *) stmt->query;
+		sel_stmt->location = stmt->location;
+	}
+
 	/* transform contained query, allowing SELECT INTO */
 	stmt->query = (Node *) transformOptionalSelectInto(pstate, stmt->query);
 
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 5a800f1e6b..7f6793f226 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -11978,6 +11978,7 @@ ExplainStmt:
 					ExplainStmt *n = makeNode(ExplainStmt);
 
 					n->query = $2;
+					n->location = @2;
 					n->options = NIL;;
 					$$ = (Node *) n;
 				}
@@ -11986,6 +11987,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,
@@ -11997,6 +11999,7 @@ ExplainStmt:
 					ExplainStmt *n = makeNode(ExplainStmt);
 
 					n->query = $3;
+					n->location = @3;
 					n->options = list_make1(makeDefElem("verbose", NULL, @2));
 					$$ = (Node *) n;
 				}
@@ -12005,6 +12008,7 @@ ExplainStmt:
 					ExplainStmt *n = makeNode(ExplainStmt);
 
 					n->query = $5;
+					n->location = @5;
 					n->options = $3;
 					$$ = (Node *) n;
 				}
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index aa73e2b242..d30869595e 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -3893,6 +3893,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;
 
 /* ----------------------

From 4d17ac1cb5c86add99881c780fd6d07b43ec5f2c Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Thu, 3 Oct 2024 10:21:44 +0200
Subject: Use view's definition as query string on a materialized view refresh

When creating a materialized view, the first refresh will have the
"Select" part of the statement as a query string. On subsequent refresh,
the "REFRESH MATERIALIZED" utility statement will be passed as query
string. This causes pgss to track both the top query and nested query as
a refresh.

This patch changes the query string on a refresh to fetch the view
definition instead. This will allow pgss to display the correct
statement when tracking refresh's nested query.
---
 .../expected/level_tracking.out               | 10 ++++----
 src/backend/commands/matview.c                | 24 +++++++++++++------
 src/backend/utils/adt/ruleutils.c             | 13 ++++++++++
 src/include/utils/ruleutils.h                 |  2 +-
 4 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 28623ae21cc..320cb93b597 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -842,12 +842,14 @@ SELECT pg_stat_statements_reset() IS NOT NULL AS t;
 REFRESH MATERIALIZED VIEW pgss_materialized_view;
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY toplevel desc, query COLLATE "C";
- toplevel | calls |                       query                        
-----------+-------+----------------------------------------------------
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
  t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
- f        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view;
-(3 rows)
+ f        |     1 | SELECT * FROM pg_catalog.pg_rewrite WHERE ev_class = $1 AND rulename = $2
+ f        |     1 | SELECT id                                                                +
+          |       |    FROM generate_series(1, 5) id(id);
+(4 rows)
 
 -- Refresh Materialized View, top-level tracking.
 SET pg_stat_statements.track = 'top';
diff --git a/src/backend/commands/matview.c b/src/backend/commands/matview.c
index 7cc68338837..92d5dd43864 100644
--- a/src/backend/commands/matview.c
+++ b/src/backend/commands/matview.c
@@ -39,6 +39,7 @@
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
+#include "utils/ruleutils.h"
 #include "utils/snapmgr.h"
 #include "utils/syscache.h"
 
@@ -61,8 +62,7 @@ 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,
-									   ParseState *pstate, const char *queryString,
-									   bool is_create);
+									   ParseState *pstate, 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);
@@ -327,11 +327,21 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData,
 	if (!skipData)
 	{
 		DestReceiver *dest;
-		const char *queryString = pstate->p_sourcetext;
+		ParseState *refresh_pstate = pstate;
+
+		/*
+		 * On refresh, the pstate's source text will be the refresh utility
+		 * statement. We need to fetch the the view definition to get the
+		 * query executed by the refresh.
+		 */
+		if (!is_create)
+		{
+			refresh_pstate = make_parsestate(NULL);
+			refresh_pstate->p_sourcetext = pg_get_viewdef_string(matviewOid, false);
+		}
 
 		dest = CreateTransientRelDestReceiver(OIDNewHeap);
-		processed = refresh_matview_datafill(dest, dataQuery, pstate,
-											 queryString, is_create);
+		processed = refresh_matview_datafill(dest, dataQuery, refresh_pstate, is_create);
 	}
 
 	/* Make the matview match the newly generated data. */
@@ -406,8 +416,7 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData,
  */
 static uint64
 refresh_matview_datafill(DestReceiver *dest, Query *query,
-						 ParseState *pstate, const char *queryString,
-						 bool is_create)
+						 ParseState *pstate, bool is_create)
 {
 	List	   *rewritten;
 	PlannedStmt *plan;
@@ -415,6 +424,7 @@ refresh_matview_datafill(DestReceiver *dest, Query *query,
 	Query	   *copied_query;
 	uint64		processed;
 	JumbleState *jstate = NULL;
+	const char *queryString = pstate->p_sourcetext;
 
 	/* Lock and rewrite, using a copy to preserve the original query. */
 	copied_query = copyObject(query);
diff --git a/src/backend/utils/adt/ruleutils.c b/src/backend/utils/adt/ruleutils.c
index 2177d17e278..0bf85cbb759 100644
--- a/src/backend/utils/adt/ruleutils.c
+++ b/src/backend/utils/adt/ruleutils.c
@@ -776,6 +776,19 @@ pg_get_viewdef_name_ext(PG_FUNCTION_ARGS)
 	PG_RETURN_TEXT_P(string_to_text(res));
 }
 
+/*
+ * Internal version of pg_get_viewdef
+ */
+char *
+pg_get_viewdef_string(Oid viewoid, bool pretty)
+{
+	int			prettyFlags;
+
+	prettyFlags = GET_PRETTY_FLAGS(pretty);
+
+	return pg_get_viewdef_worker(viewoid, prettyFlags, WRAP_COLUMN_DEFAULT);
+}
+
 /*
  * Common code for by-OID and by-name variants of pg_get_viewdef
  */
diff --git a/src/include/utils/ruleutils.h b/src/include/utils/ruleutils.h
index 161fb5ef02c..72177b9dce3 100644
--- a/src/include/utils/ruleutils.h
+++ b/src/include/utils/ruleutils.h
@@ -29,7 +29,7 @@ extern char *pg_get_indexdef_columns(Oid indexrelid, bool pretty);
 extern char *pg_get_indexdef_columns_extended(Oid indexrelid,
 											  bits16 flags);
 extern char *pg_get_querydef(Query *query, bool pretty);
-
+extern char *pg_get_viewdef_string(Oid viewoid, bool pretty);
 extern char *pg_get_partkeydef_columns(Oid relid, bool pretty);
 extern char *pg_get_partconstrdef_string(Oid partitionId, char *aliasname);
 
-- 
2.39.5 (Apple Git-154)

From baf6ab05908d75e8da88ad7affac8c28b5ead248 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Thu, 3 Oct 2024 08:52:02 +0200
Subject: Track location to extract relevant part in nested statement

Previously, Query generated through transform would have unset
stmt_location. Extensions relying on the statement location to extract
the relevant part of the statement would fallback to use the whole
statement instead, thus showing the same string in the top and
nested level which was a source of confusion.

This patch fixes the issue by keeping track of the statement locations
and propagate it to Query during transform, allowing pgss to only show
the relevant part of the query for nested query.
---
 .../expected/level_tracking.out               | 76 +++++++++----------
 src/backend/parser/analyze.c                  | 29 +++++++
 src/backend/parser/gram.y                     | 34 +++++++--
 src/backend/parser/parse_merge.c              |  2 +
 src/include/nodes/parsenodes.h                |  6 ++
 src/include/parser/parse_node.h               |  2 +
 6 files changed, 104 insertions(+), 45 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 21a6c3ba7bd..74df1c3457d 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -192,15 +192,15 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | explain (costs off) SELECT $1
  t        |     1 | explain (costs off) SELECT $1 UNION SELECT $2
  t        |     1 | explain (costs off) UPDATE stats_track_tab SET x=$1 WHERE x=$2
- f        |     1 | explain (costs off) (SELECT $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 | 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) UPDATE stats_track_tab SET x=$1 WHERE x=$2;
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ f        |     1 | SELECT $1
+ f        |     1 | SELECT $1 UNION SELECT $2
+ f        |     1 | SELECT $1, $2
+ f        |     1 | UPDATE stats_track_tab SET x=$1 WHERE x=$2
 (15 rows)
 
 -- Explain - top-level tracking.
@@ -397,8 +397,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
@@ -407,8 +407,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
@@ -416,24 +416,22 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | explain (costs off) SELECT $1, $2, $3, $4, $5
  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
- 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) 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 | (SELECT $1, $2, $3) UNION SELECT $4, $5, $6
+ 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
+ f        |     1 | SELECT $1, $2, $3, $4
+ f        |     1 | SELECT $1, $2, $3, $4, $5
+ f        |     1 | UPDATE stats_track_tab SET x=$1
+ f        |     1 | UPDATE stats_track_tab SET x=$1 WHERE x=$2
 (29 rows)
 
 -- Explain - top-level tracking with multi statement.
@@ -645,15 +643,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
@@ -765,7 +763,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/parser/analyze.c b/src/backend/parser/analyze.c
index e901203424d..98ff3cd2ea5 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);
 
@@ -518,6 +524,8 @@ transformDeleteStmt(ParseState *pstate, DeleteStmt *stmt)
 	Node	   *qual;
 
 	qry->commandType = CMD_DELETE;
+	qry->stmt_location = stmt->location;
+	qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location);
 
 	/* process the WITH clause independently of all else */
 	if (stmt->withClause)
@@ -606,6 +614,8 @@ transformInsertStmt(ParseState *pstate, InsertStmt *stmt)
 	Assert(pstate->p_ctenamespace == NIL);
 
 	qry->commandType = CMD_INSERT;
+	qry->stmt_location = stmt->location;
+	qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location);
 	pstate->p_is_insert = true;
 
 	/* process the WITH clause independently of all else */
@@ -1347,6 +1357,21 @@ transformSelectStmt(ParseState *pstate, SelectStmt *stmt)
 	ListCell   *l;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->location;
+	if (stmt->stmt_len > 0)
+
+		/*
+		 * If the select statement is within parentheses, stmt_len will be set
+		 * and represent the length of the select within parentheses
+		 */
+		qry->stmt_len = stmt->stmt_len;
+	else
+
+		/*
+		 * Otherwise, we fallback to computing the length from the
+		 * ParseState's length and location
+		 */
+		qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location);
 
 	/* process the WITH clause independently of all else */
 	if (stmt->withClause)
@@ -1730,6 +1755,8 @@ transformSetOperationStmt(ParseState *pstate, SelectStmt *stmt)
 	int			tllen;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->location;
+	qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location);
 
 	/*
 	 * Find leftmost leaf SelectStmt.  We currently only need to do this in
@@ -2429,6 +2456,8 @@ transformUpdateStmt(ParseState *pstate, UpdateStmt *stmt)
 	Node	   *qual;
 
 	qry->commandType = CMD_UPDATE;
+	qry->stmt_location = stmt->location;
+	qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location);
 	pstate->p_is_insert = false;
 
 	/* process the WITH clause independently of all else */
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 4aa8646af7b..0133187bb91 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -190,7 +190,7 @@ static void insertSelectOptions(SelectStmt *stmt,
 								SelectLimit *limitClause,
 								WithClause *withClause,
 								core_yyscan_t yyscanner);
-static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg);
+static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location);
 static Node *doNegate(Node *n, int location);
 static void doNegateFloat(Float *v);
 static Node *makeAndExpr(Node *lexpr, Node *rexpr, int location);
@@ -12170,6 +12170,9 @@ InsertStmt:
 					$5->onConflictClause = $6;
 					$5->returningList = $7;
 					$5->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					$5->location = @$;
 					$$ = (Node *) $5;
 				}
 		;
@@ -12323,6 +12326,9 @@ DeleteStmt: opt_with_clause DELETE_P FROM relation_expr_opt_alias
 					n->whereClause = $6;
 					n->returningList = $7;
 					n->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					n->location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12397,6 +12403,9 @@ UpdateStmt: opt_with_clause UPDATE relation_expr_opt_alias
 					n->whereClause = $7;
 					n->returningList = $8;
 					n->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					n->location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12474,6 +12483,9 @@ MergeStmt:
 					m->joinCondition = $8;
 					m->mergeWhenClauses = $9;
 					m->returningList = $10;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					m->location = @$;
 
 					$$ = (Node *) m;
 				}
@@ -12714,7 +12726,12 @@ SelectStmt: select_no_parens			%prec UMINUS
 		;
 
 select_with_parens:
-			'(' select_no_parens ')'				{ $$ = $2; }
+			'(' select_no_parens ')'
+				{
+					SelectStmt *n = (SelectStmt *) $2;
+					n->stmt_len = @3 - @2;
+					$$ = $2;
+				}
 			| '(' select_with_parens ')'			{ $$ = $2; }
 		;
 
@@ -12836,6 +12853,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->location = @1;
 					$$ = (Node *) n;
 				}
 			| SELECT distinct_clause target_list
@@ -12853,6 +12871,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->location = @1;
 					$$ = (Node *) n;
 				}
 			| values_clause							{ $$ = $1; }
@@ -12877,15 +12896,15 @@ simple_select:
 				}
 			| select_clause UNION set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause INTERSECT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause EXCEPT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 		;
 
@@ -18963,11 +18982,13 @@ insertSelectOptions(SelectStmt *stmt,
 					 errmsg("multiple WITH clauses not allowed"),
 					 parser_errposition(exprLocation((Node *) withClause))));
 		stmt->withClause = withClause;
+		/* Update SelectStmt's location to the start of the with clause */
+		stmt->location = withClause->location;
 	}
 }
 
 static Node *
-makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
+makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location)
 {
 	SelectStmt *n = makeNode(SelectStmt);
 
@@ -18975,6 +18996,7 @@ makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
 	n->all = all;
 	n->larg = (SelectStmt *) larg;
 	n->rarg = (SelectStmt *) rarg;
+	n->location = location;
 	return (Node *) n;
 }
 
diff --git a/src/backend/parser/parse_merge.c b/src/backend/parser/parse_merge.c
index 87df79027d7..f5f34116a61 100644
--- a/src/backend/parser/parse_merge.c
+++ b/src/backend/parser/parse_merge.c
@@ -118,6 +118,8 @@ transformMergeStmt(ParseState *pstate, MergeStmt *stmt)
 	Assert(pstate->p_ctenamespace == NIL);
 
 	qry->commandType = CMD_MERGE;
+	qry->stmt_location = stmt->location;
+	qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location);
 	qry->hasRecursive = false;
 
 	/* process the WITH clause independently of all else */
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 1c314cd9074..b47fe03cdea 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -2045,6 +2045,7 @@ typedef struct InsertStmt
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
 	OverridingKind override;	/* OVERRIDING clause */
+	ParseLoc	location;		/* token location, or -1 if unknown */
 } InsertStmt;
 
 /* ----------------------
@@ -2059,6 +2060,7 @@ typedef struct DeleteStmt
 	Node	   *whereClause;	/* qualifications */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	location;		/* token location, or -1 if unknown */
 } DeleteStmt;
 
 /* ----------------------
@@ -2074,6 +2076,7 @@ typedef struct UpdateStmt
 	List	   *fromClause;		/* optional from clause for more tables */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	location;		/* token location, or -1 if unknown */
 } UpdateStmt;
 
 /* ----------------------
@@ -2089,6 +2092,7 @@ typedef struct MergeStmt
 	List	   *mergeWhenClauses;	/* list of MergeWhenClause(es) */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	location;		/* token location, or -1 if unknown */
 } MergeStmt;
 
 /* ----------------------
@@ -2158,6 +2162,8 @@ typedef struct SelectStmt
 	bool		all;			/* ALL specified? */
 	struct SelectStmt *larg;	/* left child */
 	struct SelectStmt *rarg;	/* right child */
+	ParseLoc	location;		/* token location, or -1 if unknown */
+	ParseLoc	stmt_len;
 	/* Eventually add fields for CORRESPONDING spec here */
 } SelectStmt;
 
diff --git a/src/include/parser/parse_node.h b/src/include/parser/parse_node.h
index 543df568147..ba572b3aea2 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 */
-- 
2.39.5 (Apple Git-154)

From e232cc34249af05c9670544fe919843d3b9f6789 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Mon, 7 Oct 2024 10:45:49 +0200
Subject: Add tests covering pgss nested queries

What pgss reports for nested statements can be confusing. Some
statements like CreateTableAs, DeclareCursor and CreateMaterializedView
don't jumble the nested query and thus won't report it in pgss. Explain
explicitely Jumble the nested query and will appear in pgss. However,
the reported query string will be the same as the top level statement.

Multi statements queries are also not handled correctly. The statement
length used in the top RawStmt is not propagated in the nested
statement. Thus, it fallbacks to the "use the whole query string"
default which is reported to pgss.

This patch creates tests to cover the current behaviour.
---
 .../expected/level_tracking.out               | 895 ++++++++++++++++++
 .../pg_stat_statements/sql/level_tracking.sql | 188 +++-
 2 files changed, 1082 insertions(+), 1 deletion(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index bb65e98ce09..21a6c3ba7bd 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -112,6 +112,901 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (2 rows)
 
+-- Explain - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                   query                                                    
+----------+-------+------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (SELECT $1, $2)
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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                                                                   +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1 UNION SELECT $2
+ t        |     1 | explain (costs off) UPDATE stats_track_tab SET x=$1 WHERE x=$2
+ f        |     1 | explain (costs off) (SELECT $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+
+          |       |     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) UPDATE stats_track_tab SET x=$1 WHERE x=$2;
+(15 rows)
+
+-- Explain - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                   query                                                    
+----------+-------+------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (SELECT $1, $2)
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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                                                                   +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1 UNION SELECT $2
+ t        |     1 | explain (costs off) UPDATE stats_track_tab SET x=$1 WHERE x=$2
+(8 rows)
+
+-- Explain - all-level tracking with multi statement.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES (1), (2);
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+           QUERY PLAN            
+---------------------------------
+ Insert on stats_track_tab
+   ->  Values Scan on "*VALUES*"
+(2 rows)
+
+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;
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+         QUERY PLAN         
+----------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+           QUERY PLAN            
+---------------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2), (3)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ 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
+ t        |     1 | explain (costs off) (SELECT $1, $2, $3, $4)
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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                                                                                         +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1, $2
+ t        |     1 | explain (costs off) SELECT $1, $2 UNION SELECT $3, $4
+ t        |     1 | explain (costs off) SELECT $1, $2, $3, $4, $5
+ 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
+ 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) 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;
+(29 rows)
+
+-- Explain - top-level tracking with multi statement.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES ((1), (2));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+ERROR:  INSERT has more expressions than target columns
+LINE 1: ...n (costs off) INSERT INTO stats_track_tab VALUES ((1), (2));
+                                                                   ^
+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;
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+         QUERY PLAN         
+----------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+           QUERY PLAN            
+---------------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2), (3)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ 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
+ t        |     1 | explain (costs off) (SELECT $1, $2, $3, $4)
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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))
+ 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
+ t        |     1 | explain (costs off) SELECT $1, $2 UNION SELECT $3, $4
+ t        |     1 | explain (costs off) SELECT $1, $2, $3, $4, $5
+ 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
+(14 rows)
+
+-- Explain with CTE - all-level tracking
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) WITH a AS (select 4) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) 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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                              query                                                               
+----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (WITH a AS (select $1) (SELECT $2, $3))
+ t        |     1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab
+ t        |     1 | explain (costs off) WITH a AS (select $1) INSERT INTO stats_track_tab VALUES (($2))
+ t        |     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+
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                                         +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ 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+
+          |       |     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;
+(15 rows)
+
+-- Explain with CTE - top-level tracking
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) WITH a AS (select 4) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) 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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                              query                                                               
+----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (WITH a AS (select $1) (SELECT $2, $3))
+ t        |     1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab
+ t        |     1 | explain (costs off) WITH a AS (select $1) INSERT INTO stats_track_tab VALUES (($2))
+ t        |     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+
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                                         +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ 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
+(8 rows)
+
+-- Explain analyze, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                     query                                                     
+----------+-------+---------------------------------------------------------------------------------------------------------------
+ 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;
+(4 rows)
+
+-- Explain analyze, top tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                     query                                                     
+----------+-------+---------------------------------------------------------------------------------------------------------------
+ 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
+(3 rows)
+
+-- Create Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                             query                                              
+----------+-------+------------------------------------------------------------------------------------------------
+ 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)
+
+-- Create Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE MATERIALIZED VIEW pgss_materialized_view_2 AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                              query                                               
+----------+-------+--------------------------------------------------------------------------------------------------
+ t        |     1 | CREATE MATERIALIZED VIEW pgss_materialized_view_2 AS SELECT * FROM generate_series($1, $2) as id
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Refresh Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Refresh Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT $1
+ 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)
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_ctas_3 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_4 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_3 AS SELECT $1
+ t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_4 AS EXECUTE test_prepare_pgss
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Explain with Create Table As - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
+ 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)
+
+-- Explain with Create Table As - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
+ 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)
+
+-- Declare cursor, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                          query                          
+----------+-------+---------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ 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)
+
+-- Declare cursor, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                          query                          
+----------+-------+---------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ 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)
+
 -- 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 65a17147a5a..b72964aa305 100644
--- a/contrib/pg_stat_statements/sql/level_tracking.sql
+++ b/contrib/pg_stat_statements/sql/level_tracking.sql
@@ -32,7 +32,6 @@ BEGIN
 END; $$;
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
-
 -- Procedure with multiple utility statements.
 CREATE OR REPLACE PROCEDURE proc_with_utility_stmt()
 LANGUAGE SQL
@@ -55,6 +54,193 @@ CALL proc_with_utility_stmt();
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
 
+-- Explain - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1;
+explain (costs off) (SELECT 1, 2);
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) DELETE FROM stats_track_tab;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((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 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1;
+explain (costs off) (SELECT 1, 2);
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) DELETE FROM stats_track_tab;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((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 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain - all-level tracking with multi statement.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES (1), (2);
+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;
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain - top-level tracking with multi statement.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES ((1), (2));
+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;
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with CTE - all-level tracking
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) WITH a AS (select 4) SELECT 1;
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+explain (costs off) WITH a AS (select 4) 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) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with CTE - top-level tracking
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) WITH a AS (select 4) SELECT 1;
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+explain (costs off) WITH a AS (select 4) 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) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain analyze, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain analyze, top tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE MATERIALIZED VIEW pgss_materialized_view_2 AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Refresh Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Refresh Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_ctas_3 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_4 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with Create Table As - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with Create Table As - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Declare cursor, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Declare cursor, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
 -- DO block - top-level tracking without utility.
 SET pg_stat_statements.track = 'top';
 SET pg_stat_statements.track_utility = FALSE;
-- 
2.39.5 (Apple Git-154)

From cc4f2314f06ed7b17bdb87d5558c87992dc58452 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
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)

From 0a0a481cc0218bf59b290a4097fb08da6ef696cf Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Tue, 8 Oct 2024 08:45:39 +0200
Subject: Extract nested query from PrepareStmt

Previously, executing a prepared query would report the full "PREPARE x
AS ..." in pg_stat_statements. This patch extracts the
Insert/Select/Delete/Update statement nested in prepare and report only
the relevant query part.
---
 .../expected/level_tracking.out               |  2 +-
 .../pg_stat_statements/expected/planning.out  | 10 +++---
 .../pg_stat_statements/expected/select.out    |  2 +-
 .../pg_stat_statements/expected/utility.out   | 19 +++++++++--
 contrib/pg_stat_statements/sql/planning.sql   |  4 +--
 contrib/pg_stat_statements/sql/utility.sql    |  4 +++
 src/backend/commands/prepare.c                | 33 +++++++++++++++++--
 7 files changed, 61 insertions(+), 13 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 320cb93b597..d36bdaa665b 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -886,8 +886,8 @@ 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
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
- f        |     1 | PREPARE test_prepare_pgss AS select generate_series($1, $2)
  f        |     1 | SELECT $1
+ f        |     1 | select generate_series($1, $2)
 (5 rows)
 
 -- Create Table As, top-level tracking.
diff --git a/contrib/pg_stat_statements/expected/planning.out b/contrib/pg_stat_statements/expected/planning.out
index 9effd11fdc8..3ee1928cbe9 100644
--- a/contrib/pg_stat_statements/expected/planning.out
+++ b/contrib/pg_stat_statements/expected/planning.out
@@ -58,7 +58,7 @@ SELECT 42;
 (1 row)
 
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
  plans | calls | rows |                          query                           
 -------+-------+------+----------------------------------------------------------
      0 |     1 |    0 | ALTER TABLE stats_plan_test ADD COLUMN x int
@@ -72,10 +72,10 @@ SELECT plans, calls, rows, query FROM pg_stat_statements
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
- plans_ok | calls | rows |                         query                         
-----------+-------+------+-------------------------------------------------------
- t        |     4 |    4 | PREPARE prep1 AS SELECT COUNT(*) FROM stats_plan_test
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
+ plans_ok | calls | rows |                query                 
+----------+-------+------+--------------------------------------
+ t        |     4 |    4 | SELECT COUNT(*) FROM stats_plan_test
 (1 row)
 
 -- Cleanup
diff --git a/contrib/pg_stat_statements/expected/select.out b/contrib/pg_stat_statements/expected/select.out
index dd6c756f67d..217a2c0b2bc 100644
--- a/contrib/pg_stat_statements/expected/select.out
+++ b/contrib/pg_stat_statements/expected/select.out
@@ -127,7 +127,6 @@ DEALLOCATE pgss_test;
 SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
  calls | rows |                                    query                                     
 -------+------+------------------------------------------------------------------------------
-     1 |    1 | PREPARE pgss_test (int) AS SELECT $1, $2 LIMIT $3
      4 |    4 | SELECT $1                                                                   +
        |      |   -- multiline                                                              +
        |      |   AS "text"
@@ -137,6 +136,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
      2 |    2 | SELECT $1 AS "int"
      1 |    2 | SELECT $1 AS i UNION SELECT $2 ORDER BY i
      1 |    1 | SELECT $1 || $2
+     1 |    1 | SELECT $1, $2 LIMIT $3
      0 |    0 | SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C"
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
      1 |    2 | WITH t(f) AS (                                                              +
diff --git a/contrib/pg_stat_statements/expected/utility.out b/contrib/pg_stat_statements/expected/utility.out
index 060d4416dd7..6c0a4a543c2 100644
--- a/contrib/pg_stat_statements/expected/utility.out
+++ b/contrib/pg_stat_statements/expected/utility.out
@@ -532,6 +532,19 @@ EXECUTE stat_select (2);
  2
 (1 row)
 
+PREPARE stat_select_2 AS SELECT $1, $2 AS a\; PREPARE stat_select_3 AS SELECT $1, $2, $3 AS a;
+EXECUTE stat_select_2 (2, 3);
+ ?column? | a 
+----------+---
+ 2        | 3
+(1 row)
+
+EXECUTE stat_select_3 (2, 3, 4);
+ ?column? | ?column? | a 
+----------+----------+---
+ 2        | 3        | 4
+(1 row)
+
 DEALLOCATE PREPARE stat_select;
 DEALLOCATE ALL;
 DEALLOCATE PREPARE ALL;
@@ -540,10 +553,12 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
 -------+------+----------------------------------------------------
      2 |    0 | DEALLOCATE $1
      2 |    0 | DEALLOCATE ALL
-     2 |    2 | PREPARE stat_select AS SELECT $1 AS a
+     2 |    2 | SELECT $1 AS a
      1 |    1 | SELECT $1 as a
+     1 |    1 | SELECT $1, $2 AS a
+     1 |    1 | SELECT $1, $2, $3 AS a
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(5 rows)
+(7 rows)
 
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
  t 
diff --git a/contrib/pg_stat_statements/sql/planning.sql b/contrib/pg_stat_statements/sql/planning.sql
index 46f5d9b951c..9cfe206b3b0 100644
--- a/contrib/pg_stat_statements/sql/planning.sql
+++ b/contrib/pg_stat_statements/sql/planning.sql
@@ -20,11 +20,11 @@ SELECT 42;
 SELECT 42;
 SELECT 42;
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 
 -- Cleanup
 DROP TABLE stats_plan_test;
diff --git a/contrib/pg_stat_statements/sql/utility.sql b/contrib/pg_stat_statements/sql/utility.sql
index dd97203c210..7a7df695910 100644
--- a/contrib/pg_stat_statements/sql/utility.sql
+++ b/contrib/pg_stat_statements/sql/utility.sql
@@ -270,6 +270,10 @@ EXECUTE stat_select (1);
 DEALLOCATE stat_select;
 PREPARE stat_select AS SELECT $1 AS a;
 EXECUTE stat_select (2);
+PREPARE stat_select_2 AS SELECT $1, $2 AS a\; PREPARE stat_select_3 AS SELECT $1, $2, $3 AS a;
+EXECUTE stat_select_2 (2, 3);
+EXECUTE stat_select_3 (2, 3, 4);
+
 DEALLOCATE PREPARE stat_select;
 DEALLOCATE ALL;
 DEALLOCATE PREPARE ALL;
diff --git a/src/backend/commands/prepare.c b/src/backend/commands/prepare.c
index a93f970a292..39aa447cf98 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -77,8 +77,37 @@ PrepareQuery(ParseState *pstate, PrepareStmt *stmt,
 	 */
 	rawstmt = makeNode(RawStmt);
 	rawstmt->stmt = stmt->query;
-	rawstmt->stmt_location = stmt_location;
-	rawstmt->stmt_len = stmt_len;
+
+	switch (nodeTag(stmt->query))
+	{
+		case T_InsertStmt:
+			rawstmt->stmt_location = ((InsertStmt *) stmt->query)->location;
+			break;
+		case T_DeleteStmt:
+			rawstmt->stmt_location = ((DeleteStmt *) stmt->query)->location;
+			break;
+		case T_UpdateStmt:
+			rawstmt->stmt_location = ((UpdateStmt *) stmt->query)->location;
+			break;
+		case T_MergeStmt:
+			rawstmt->stmt_location = ((MergeStmt *) stmt->query)->location;
+			break;
+		case T_SelectStmt:
+			rawstmt->stmt_location = ((SelectStmt *) stmt->query)->location;
+			rawstmt->stmt_len = ((SelectStmt *) stmt->query)->stmt_len;
+			break;
+		default:
+			elog(ERROR, "unexpected node type: %d", (int) nodeTag(stmt->query));
+			break;
+	}
+
+	/*
+	 * stmt_len will be defined for SelectStmt within parentheses. If it's
+	 * defined, use it. Otherwise, we need to compute the new length based on
+	 * the new statement location and the initial location+length.
+	 */
+	if (stmt_len > 0 && rawstmt->stmt_len == 0)
+		rawstmt->stmt_len = stmt_len - (rawstmt->stmt_location - stmt_location);
 
 	/*
 	 * Create the CachedPlanSource before we do parse analysis, since it needs
-- 
2.39.5 (Apple Git-154)

From 61f8920b7cc938cd84c7663ae8e7b78533731dad Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
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 fb3458fcdfa..176020269f7 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -881,7 +881,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';
@@ -927,7 +928,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';
@@ -961,7 +963,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';
@@ -998,8 +1001,9 @@ 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
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | SELECT $1
  f        |     1 | select generate_series($1, $2)
-(4 rows)
+(5 rows)
 
 -- Create Table As, top-level tracking.
 SET pg_stat_statements.track = 'top';
@@ -1040,7 +1044,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';
@@ -1091,7 +1096,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 18a5af6b919..c81221cdbeb 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 dcbdaa03885..d2eef8097cf 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 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)

From aeabd8310cc5400c9534e2458a59240bc4d74bc3 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Thu, 3 Oct 2024 10:21:44 +0200
Subject: Use view's definition as query string on a materialized view refresh

When creating a materialized view, the first refresh will have the
"Select" part of the statement as a query string. On subsequent refresh,
the "REFRESH MATERIALIZED" utility statement will be passed as query
string. This causes pgss to track both the top query and nested query as
a refresh.

This patch changes the query string on a refresh to fetch the view
definition instead. This will allow pgss to display the correct
statement when tracking refresh's nested query.
---
 .../expected/level_tracking.out               | 10 ++++----
 src/backend/commands/matview.c                | 24 +++++++++++++------
 src/backend/utils/adt/ruleutils.c             | 13 ++++++++++
 src/include/utils/ruleutils.h                 |  2 +-
 4 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 176020269f7..cc89ba9af41 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -959,12 +959,14 @@ SELECT pg_stat_statements_reset() IS NOT NULL AS t;
 REFRESH MATERIALIZED VIEW pgss_materialized_view;
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY toplevel desc, query COLLATE "C";
- toplevel | calls |                       query                        
-----------+-------+----------------------------------------------------
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
  t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
- f        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view;
-(3 rows)
+ f        |     1 | SELECT * FROM pg_catalog.pg_rewrite WHERE ev_class = $1 AND rulename = $2
+ f        |     1 | SELECT id                                                                +
+          |       |    FROM generate_series(1, 5) id(id);
+(4 rows)
 
 -- Refresh Materialized View, top-level tracking.
 SET pg_stat_statements.track = 'top';
diff --git a/src/backend/commands/matview.c b/src/backend/commands/matview.c
index 7cc68338837..92d5dd43864 100644
--- a/src/backend/commands/matview.c
+++ b/src/backend/commands/matview.c
@@ -39,6 +39,7 @@
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
+#include "utils/ruleutils.h"
 #include "utils/snapmgr.h"
 #include "utils/syscache.h"
 
@@ -61,8 +62,7 @@ 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,
-									   ParseState *pstate, const char *queryString,
-									   bool is_create);
+									   ParseState *pstate, 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);
@@ -327,11 +327,21 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData,
 	if (!skipData)
 	{
 		DestReceiver *dest;
-		const char *queryString = pstate->p_sourcetext;
+		ParseState *refresh_pstate = pstate;
+
+		/*
+		 * On refresh, the pstate's source text will be the refresh utility
+		 * statement. We need to fetch the the view definition to get the
+		 * query executed by the refresh.
+		 */
+		if (!is_create)
+		{
+			refresh_pstate = make_parsestate(NULL);
+			refresh_pstate->p_sourcetext = pg_get_viewdef_string(matviewOid, false);
+		}
 
 		dest = CreateTransientRelDestReceiver(OIDNewHeap);
-		processed = refresh_matview_datafill(dest, dataQuery, pstate,
-											 queryString, is_create);
+		processed = refresh_matview_datafill(dest, dataQuery, refresh_pstate, is_create);
 	}
 
 	/* Make the matview match the newly generated data. */
@@ -406,8 +416,7 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData,
  */
 static uint64
 refresh_matview_datafill(DestReceiver *dest, Query *query,
-						 ParseState *pstate, const char *queryString,
-						 bool is_create)
+						 ParseState *pstate, bool is_create)
 {
 	List	   *rewritten;
 	PlannedStmt *plan;
@@ -415,6 +424,7 @@ refresh_matview_datafill(DestReceiver *dest, Query *query,
 	Query	   *copied_query;
 	uint64		processed;
 	JumbleState *jstate = NULL;
+	const char *queryString = pstate->p_sourcetext;
 
 	/* Lock and rewrite, using a copy to preserve the original query. */
 	copied_query = copyObject(query);
diff --git a/src/backend/utils/adt/ruleutils.c b/src/backend/utils/adt/ruleutils.c
index 2177d17e278..0bf85cbb759 100644
--- a/src/backend/utils/adt/ruleutils.c
+++ b/src/backend/utils/adt/ruleutils.c
@@ -776,6 +776,19 @@ pg_get_viewdef_name_ext(PG_FUNCTION_ARGS)
 	PG_RETURN_TEXT_P(string_to_text(res));
 }
 
+/*
+ * Internal version of pg_get_viewdef
+ */
+char *
+pg_get_viewdef_string(Oid viewoid, bool pretty)
+{
+	int			prettyFlags;
+
+	prettyFlags = GET_PRETTY_FLAGS(pretty);
+
+	return pg_get_viewdef_worker(viewoid, prettyFlags, WRAP_COLUMN_DEFAULT);
+}
+
 /*
  * Common code for by-OID and by-name variants of pg_get_viewdef
  */
diff --git a/src/include/utils/ruleutils.h b/src/include/utils/ruleutils.h
index 161fb5ef02c..72177b9dce3 100644
--- a/src/include/utils/ruleutils.h
+++ b/src/include/utils/ruleutils.h
@@ -29,7 +29,7 @@ extern char *pg_get_indexdef_columns(Oid indexrelid, bool pretty);
 extern char *pg_get_indexdef_columns_extended(Oid indexrelid,
 											  bits16 flags);
 extern char *pg_get_querydef(Query *query, bool pretty);
-
+extern char *pg_get_viewdef_string(Oid viewoid, bool pretty);
 extern char *pg_get_partkeydef_columns(Oid relid, bool pretty);
 extern char *pg_get_partconstrdef_string(Oid partitionId, char *aliasname);
 
-- 
2.39.5 (Apple Git-154)

From 7dcbf714c723b310be8684392b934d99f14d94a4 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Thu, 3 Oct 2024 08:52:02 +0200
Subject: Track location to extract relevant part in nested statement

Previously, Query generated through transform would have unset
stmt_location. Extensions relying on the statement location to extract
the relevant part of the statement would fallback to use the whole
statement instead, thus showing the same string in the top and
nested level which was a source of confusion.

This patch fixes the issue by keeping track of the statement locations
and propagate it to Query during transform, allowing pgss to only show
the relevant part of the query for nested query.
---
 .../expected/level_tracking.out               | 94 +++++++++----------
 .../pg_stat_statements/expected/planning.out  | 10 +-
 .../pg_stat_statements/expected/select.out    |  2 +-
 .../pg_stat_statements/expected/utility.out   |  2 +-
 contrib/pg_stat_statements/sql/planning.sql   |  4 +-
 src/backend/optimizer/util/clauses.c          |  2 +-
 src/backend/parser/analyze.c                  | 92 ++++++++++++------
 src/backend/parser/gram.y                     | 36 +++++--
 src/backend/parser/parse_merge.c              |  2 +
 src/include/nodes/parsenodes.h                |  6 ++
 src/include/parser/analyze.h                  |  2 +-
 src/include/parser/parse_node.h               |  2 +
 12 files changed, 162 insertions(+), 92 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 9463c68154b..fb3458fcdfa 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 | 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 | SELECT $1, $2
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | TABLE test_table
+ f        |     1 | UPDATE stats_track_tab SET x=$1 WHERE x=$2
+ f        |     1 | VALUES($1)
+ f        |     1 | VALUES($1, $2)
 (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) UNION SELECT $4, $5, $6
+ 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
+ f        |     1 | SELECT $1, $2, $3, $4
+ f        |     1 | SELECT $1, $2, $3, $4, $5
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | TABLE test_table
+ 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)
+ f        |     1 | VALUES($1, $2)
 (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.
@@ -1000,7 +998,7 @@ 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
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
- f        |     1 | PREPARE test_prepare_pgss AS select generate_series($1, $2)
+ f        |     1 | select generate_series($1, $2)
 (4 rows)
 
 -- Create Table As, top-level tracking.
diff --git a/contrib/pg_stat_statements/expected/planning.out b/contrib/pg_stat_statements/expected/planning.out
index 9effd11fdc8..3ee1928cbe9 100644
--- a/contrib/pg_stat_statements/expected/planning.out
+++ b/contrib/pg_stat_statements/expected/planning.out
@@ -58,7 +58,7 @@ SELECT 42;
 (1 row)
 
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
  plans | calls | rows |                          query                           
 -------+-------+------+----------------------------------------------------------
      0 |     1 |    0 | ALTER TABLE stats_plan_test ADD COLUMN x int
@@ -72,10 +72,10 @@ SELECT plans, calls, rows, query FROM pg_stat_statements
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
- plans_ok | calls | rows |                         query                         
-----------+-------+------+-------------------------------------------------------
- t        |     4 |    4 | PREPARE prep1 AS SELECT COUNT(*) FROM stats_plan_test
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
+ plans_ok | calls | rows |                query                 
+----------+-------+------+--------------------------------------
+ t        |     4 |    4 | SELECT COUNT(*) FROM stats_plan_test
 (1 row)
 
 -- Cleanup
diff --git a/contrib/pg_stat_statements/expected/select.out b/contrib/pg_stat_statements/expected/select.out
index dd6c756f67d..217a2c0b2bc 100644
--- a/contrib/pg_stat_statements/expected/select.out
+++ b/contrib/pg_stat_statements/expected/select.out
@@ -127,7 +127,6 @@ DEALLOCATE pgss_test;
 SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
  calls | rows |                                    query                                     
 -------+------+------------------------------------------------------------------------------
-     1 |    1 | PREPARE pgss_test (int) AS SELECT $1, $2 LIMIT $3
      4 |    4 | SELECT $1                                                                   +
        |      |   -- multiline                                                              +
        |      |   AS "text"
@@ -137,6 +136,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
      2 |    2 | SELECT $1 AS "int"
      1 |    2 | SELECT $1 AS i UNION SELECT $2 ORDER BY i
      1 |    1 | SELECT $1 || $2
+     1 |    1 | SELECT $1, $2 LIMIT $3
      0 |    0 | SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C"
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
      1 |    2 | WITH t(f) AS (                                                              +
diff --git a/contrib/pg_stat_statements/expected/utility.out b/contrib/pg_stat_statements/expected/utility.out
index 060d4416dd7..aa4f0f7e628 100644
--- a/contrib/pg_stat_statements/expected/utility.out
+++ b/contrib/pg_stat_statements/expected/utility.out
@@ -540,7 +540,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
 -------+------+----------------------------------------------------
      2 |    0 | DEALLOCATE $1
      2 |    0 | DEALLOCATE ALL
-     2 |    2 | PREPARE stat_select AS SELECT $1 AS a
+     2 |    2 | SELECT $1 AS a
      1 |    1 | SELECT $1 as a
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (5 rows)
diff --git a/contrib/pg_stat_statements/sql/planning.sql b/contrib/pg_stat_statements/sql/planning.sql
index 46f5d9b951c..9cfe206b3b0 100644
--- a/contrib/pg_stat_statements/sql/planning.sql
+++ b/contrib/pg_stat_statements/sql/planning.sql
@@ -20,11 +20,11 @@ SELECT 42;
 SELECT 42;
 SELECT 42;
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 
 -- Cleanup
 DROP TABLE stats_plan_test;
diff --git a/src/backend/optimizer/util/clauses.c b/src/backend/optimizer/util/clauses.c
index b4e085e9d4b..8fc4da11483 100644
--- a/src/backend/optimizer/util/clauses.c
+++ b/src/backend/optimizer/util/clauses.c
@@ -4674,7 +4674,7 @@ inline_function(Oid funcid, Oid result_type, Oid result_collid,
 		pstate->p_sourcetext = src;
 		sql_fn_parser_setup(pstate, pinfo);
 
-		querytree = transformTopLevelStmt(pstate, linitial(raw_parsetree_list));
+		querytree = transformOptionalSelectInto(pstate, ((RawStmt *) linitial(raw_parsetree_list))->stmt);
 
 		free_parsestate(pstate);
 	}
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index e901203424d..4c0bcf93a87 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -58,7 +58,6 @@
 /* Hook for plugins to get control at end of parse analysis */
 post_parse_analyze_hook_type post_parse_analyze_hook = NULL;
 
-static Query *transformOptionalSelectInto(ParseState *pstate, Node *parseTree);
 static Query *transformDeleteStmt(ParseState *pstate, DeleteStmt *stmt);
 static Query *transformInsertStmt(ParseState *pstate, InsertStmt *stmt);
 static OnConflictExpr *transformOnConflictClause(ParseState *pstate,
@@ -113,13 +112,15 @@ 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);
 
 	pstate->p_queryEnv = queryEnv;
 
-	query = transformTopLevelStmt(pstate, parseTree);
+	query = transformOptionalSelectInto(pstate, parseTree->stmt);
 
 	if (IsQueryIdEnabled())
 		jstate = JumbleQuery(query);
@@ -153,12 +154,14 @@ 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);
 
 	pstate->p_queryEnv = queryEnv;
 
-	query = transformTopLevelStmt(pstate, parseTree);
+	query = transformOptionalSelectInto(pstate, parseTree->stmt);
 
 	/* make sure all is well with parameter types */
 	check_variable_parameters(pstate, query);
@@ -195,10 +198,12 @@ 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);
 
-	query = transformTopLevelStmt(pstate, parseTree);
+	query = transformOptionalSelectInto(pstate, parseTree->stmt);
 
 	if (IsQueryIdEnabled())
 		jstate = JumbleQuery(query);
@@ -238,27 +243,6 @@ parse_sub_analyze(Node *parseTree, ParseState *parentParseState,
 	return query;
 }
 
-/*
- * transformTopLevelStmt -
- *	  transform a Parse tree into a Query tree.
- *
- * This function is just responsible for transferring statement location data
- * from the RawStmt into the finished Query.
- */
-Query *
-transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree)
-{
-	Query	   *result;
-
-	/* We're at top level, so allow SELECT INTO */
-	result = transformOptionalSelectInto(pstate, parseTree->stmt);
-
-	result->stmt_location = parseTree->stmt_location;
-	result->stmt_len = parseTree->stmt_len;
-
-	return result;
-}
-
 /*
  * transformOptionalSelectInto -
  *	  If SELECT has INTO, convert it to CREATE TABLE AS.
@@ -269,7 +253,7 @@ transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree)
  * of the parse tree, and so we only try it before entering the recursive
  * transformStmt() processing.
  */
-static Query *
+Query *
 transformOptionalSelectInto(ParseState *pstate, Node *parseTree)
 {
 	if (IsA(parseTree, SelectStmt))
@@ -417,7 +401,9 @@ transformStmt(ParseState *pstate, Node *parseTree)
 			 */
 			result = makeNode(Query);
 			result->commandType = CMD_UTILITY;
-			result->utilityStmt = (Node *) parseTree;
+			result->utilityStmt = parseTree;
+			result->stmt_location = pstate->p_stmt_location;
+			result->stmt_len = pstate->p_stmt_len;
 			break;
 	}
 
@@ -518,6 +504,16 @@ transformDeleteStmt(ParseState *pstate, DeleteStmt *stmt)
 	Node	   *qual;
 
 	qry->commandType = CMD_DELETE;
+	qry->stmt_location = stmt->location;
+
+	/*
+	 * The parser can't provide the length of individual statements. However,
+	 * we have the statement's location plus the length (p_stmt_len) and
+	 * location (p_stmt_location) of the top level RawStmt, stored in pstate.
+	 * Thus, the statement's length is the RawStmt's length minus how much
+	 * we've advanced in the RawStmt's string.
+	 */
+	qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location);
 
 	/* process the WITH clause independently of all else */
 	if (stmt->withClause)
@@ -606,6 +602,9 @@ transformInsertStmt(ParseState *pstate, InsertStmt *stmt)
 	Assert(pstate->p_ctenamespace == NIL);
 
 	qry->commandType = CMD_INSERT;
+	qry->stmt_location = stmt->location;
+	/* see comment in transformDeleteStmt */
+	qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location);
 	pstate->p_is_insert = true;
 
 	/* process the WITH clause independently of all else */
@@ -1347,6 +1346,21 @@ transformSelectStmt(ParseState *pstate, SelectStmt *stmt)
 	ListCell   *l;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->location;
+
+	/*
+	 * If the select statement has parentheses (like (SELECT 1)), stmt_len
+	 * will be set by the parser and will represent the length of the select
+	 * inside parentheses which can be used for our query length.
+	 *
+	 * Otherwise, stmt_len will be 0 and we will fallback to computing the
+	 * length from the ParseState's length and location (see comment in
+	 * transformDeleteStmt)
+	 */
+	if (stmt->stmt_len > 0)
+		qry->stmt_len = stmt->stmt_len;
+	else
+		qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location);
 
 	/* process the WITH clause independently of all else */
 	if (stmt->withClause)
@@ -1499,6 +1513,12 @@ transformValuesClause(ParseState *pstate, SelectStmt *stmt)
 	int			i;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->location;
+	/* See comment in transformSelectStmt */
+	if (stmt->stmt_len > 0)
+		qry->stmt_len = stmt->stmt_len;
+	else
+		qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location);
 
 	/* Most SELECT stuff doesn't apply in a VALUES clause */
 	Assert(stmt->distinctClause == NIL);
@@ -1730,6 +1750,9 @@ transformSetOperationStmt(ParseState *pstate, SelectStmt *stmt)
 	int			tllen;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->location;
+	/* see comment in transformDeleteStmt */
+	qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location);
 
 	/*
 	 * Find leftmost leaf SelectStmt.  We currently only need to do this in
@@ -2396,6 +2419,8 @@ transformReturnStmt(ParseState *pstate, ReturnStmt *stmt)
 	Query	   *qry = makeNode(Query);
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = pstate->p_stmt_location;
+	qry->stmt_len = pstate->p_stmt_len;
 	qry->isReturn = true;
 
 	qry->targetList = list_make1(makeTargetEntry((Expr *) transformExpr(pstate, stmt->returnval, EXPR_KIND_SELECT_TARGET),
@@ -2429,6 +2454,9 @@ transformUpdateStmt(ParseState *pstate, UpdateStmt *stmt)
 	Node	   *qual;
 
 	qry->commandType = CMD_UPDATE;
+	qry->stmt_location = stmt->location;
+	/* see comment in transformDeleteStmt */
+	qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location);
 	pstate->p_is_insert = false;
 
 	/* process the WITH clause independently of all else */
@@ -2676,6 +2704,8 @@ transformPLAssignStmt(ParseState *pstate, PLAssignStmt *stmt)
 	 * consider WITH or INTO, and we build a targetlist our own way.
 	 */
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = pstate->p_stmt_location;
+	qry->stmt_len = pstate->p_stmt_len;
 	pstate->p_is_insert = false;
 
 	/* make FOR UPDATE/FOR SHARE info available to addRangeTableEntry */
@@ -2947,6 +2977,8 @@ transformDeclareCursorStmt(ParseState *pstate, DeclareCursorStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3002,6 +3034,8 @@ transformExplainStmt(ParseState *pstate, ExplainStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3082,6 +3116,8 @@ transformCreateTableAsStmt(ParseState *pstate, CreateTableAsStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3206,6 +3242,8 @@ transformCallStmt(ParseState *pstate, CallStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 4aa8646af7b..6073c61f110 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -190,7 +190,7 @@ static void insertSelectOptions(SelectStmt *stmt,
 								SelectLimit *limitClause,
 								WithClause *withClause,
 								core_yyscan_t yyscanner);
-static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg);
+static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location);
 static Node *doNegate(Node *n, int location);
 static void doNegateFloat(Float *v);
 static Node *makeAndExpr(Node *lexpr, Node *rexpr, int location);
@@ -12170,6 +12170,9 @@ InsertStmt:
 					$5->onConflictClause = $6;
 					$5->returningList = $7;
 					$5->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					$5->location = @$;
 					$$ = (Node *) $5;
 				}
 		;
@@ -12323,6 +12326,9 @@ DeleteStmt: opt_with_clause DELETE_P FROM relation_expr_opt_alias
 					n->whereClause = $6;
 					n->returningList = $7;
 					n->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					n->location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12397,6 +12403,9 @@ UpdateStmt: opt_with_clause UPDATE relation_expr_opt_alias
 					n->whereClause = $7;
 					n->returningList = $8;
 					n->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					n->location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12474,6 +12483,9 @@ MergeStmt:
 					m->joinCondition = $8;
 					m->mergeWhenClauses = $9;
 					m->returningList = $10;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					m->location = @$;
 
 					$$ = (Node *) m;
 				}
@@ -12714,7 +12726,12 @@ SelectStmt: select_no_parens			%prec UMINUS
 		;
 
 select_with_parens:
-			'(' select_no_parens ')'				{ $$ = $2; }
+			'(' select_no_parens ')'
+				{
+					SelectStmt *n = (SelectStmt *) $2;
+					n->stmt_len = @3 - @2;
+					$$ = $2;
+				}
 			| '(' select_with_parens ')'			{ $$ = $2; }
 		;
 
@@ -12836,6 +12853,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->location = @1;
 					$$ = (Node *) n;
 				}
 			| SELECT distinct_clause target_list
@@ -12853,6 +12871,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->location = @1;
 					$$ = (Node *) n;
 				}
 			| values_clause							{ $$ = $1; }
@@ -12873,19 +12892,20 @@ simple_select:
 
 					n->targetList = list_make1(rt);
 					n->fromClause = list_make1($2);
+					n->location = @1;
 					$$ = (Node *) n;
 				}
 			| select_clause UNION set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause INTERSECT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause EXCEPT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 		;
 
@@ -13443,6 +13463,7 @@ values_clause:
 				{
 					SelectStmt *n = makeNode(SelectStmt);
 
+					n->location = @1;
 					n->valuesLists = list_make1($3);
 					$$ = (Node *) n;
 				}
@@ -18963,11 +18984,13 @@ insertSelectOptions(SelectStmt *stmt,
 					 errmsg("multiple WITH clauses not allowed"),
 					 parser_errposition(exprLocation((Node *) withClause))));
 		stmt->withClause = withClause;
+		/* Update SelectStmt's location to the start of the with clause */
+		stmt->location = withClause->location;
 	}
 }
 
 static Node *
-makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
+makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location)
 {
 	SelectStmt *n = makeNode(SelectStmt);
 
@@ -18975,6 +18998,7 @@ makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
 	n->all = all;
 	n->larg = (SelectStmt *) larg;
 	n->rarg = (SelectStmt *) rarg;
+	n->location = location;
 	return (Node *) n;
 }
 
diff --git a/src/backend/parser/parse_merge.c b/src/backend/parser/parse_merge.c
index 87df79027d7..f5f34116a61 100644
--- a/src/backend/parser/parse_merge.c
+++ b/src/backend/parser/parse_merge.c
@@ -118,6 +118,8 @@ transformMergeStmt(ParseState *pstate, MergeStmt *stmt)
 	Assert(pstate->p_ctenamespace == NIL);
 
 	qry->commandType = CMD_MERGE;
+	qry->stmt_location = stmt->location;
+	qry->stmt_len = pstate->p_stmt_len - (stmt->location - pstate->p_stmt_location);
 	qry->hasRecursive = false;
 
 	/* process the WITH clause independently of all else */
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 5b62df32733..11879b7f4fa 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -2045,6 +2045,7 @@ typedef struct InsertStmt
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
 	OverridingKind override;	/* OVERRIDING clause */
+	ParseLoc	location;		/* token location, or -1 if unknown */
 } InsertStmt;
 
 /* ----------------------
@@ -2059,6 +2060,7 @@ typedef struct DeleteStmt
 	Node	   *whereClause;	/* qualifications */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	location;		/* token location, or -1 if unknown */
 } DeleteStmt;
 
 /* ----------------------
@@ -2074,6 +2076,7 @@ typedef struct UpdateStmt
 	List	   *fromClause;		/* optional from clause for more tables */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	location;		/* token location, or -1 if unknown */
 } UpdateStmt;
 
 /* ----------------------
@@ -2089,6 +2092,7 @@ typedef struct MergeStmt
 	List	   *mergeWhenClauses;	/* list of MergeWhenClause(es) */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	location;		/* token location, or -1 if unknown */
 } MergeStmt;
 
 /* ----------------------
@@ -2158,6 +2162,8 @@ typedef struct SelectStmt
 	bool		all;			/* ALL specified? */
 	struct SelectStmt *larg;	/* left child */
 	struct SelectStmt *rarg;	/* right child */
+	ParseLoc	location;		/* token location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 	/* Eventually add fields for CORRESPONDING spec here */
 } SelectStmt;
 
diff --git a/src/include/parser/analyze.h b/src/include/parser/analyze.h
index 28b66fccb43..0225258fe2e 100644
--- a/src/include/parser/analyze.h
+++ b/src/include/parser/analyze.h
@@ -46,7 +46,7 @@ extern List *transformUpdateTargetList(ParseState *pstate,
 									   List *origTlist);
 extern List *transformReturningList(ParseState *pstate, List *returningList,
 									ParseExprKind exprKind);
-extern Query *transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree);
+extern Query *transformOptionalSelectInto(ParseState *pstate, Node *parseTree);
 extern Query *transformStmt(ParseState *pstate, Node *parseTree);
 
 extern bool stmt_requires_parse_analysis(RawStmt *parseTree);
diff --git a/src/include/parser/parse_node.h b/src/include/parser/parse_node.h
index 543df568147..ba572b3aea2 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 */
-- 
2.39.5 (Apple Git-154)

From 1b8094a8bce62c2e46576cd38c5951e51c4d16d2 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Mon, 7 Oct 2024 10:45:49 +0200
Subject: Add tests covering pgss nested queries

What pgss reports for nested statements can be confusing. Some
statements like CreateTableAs, DeclareCursor and CreateMaterializedView
don't jumble the nested query and thus won't report it in pgss. Explain
explicitely Jumble the nested query and will appear in pgss. However,
the reported query string will be the same as the top level statement.

Multi statements queries are also not handled correctly. The statement
length used in the top RawStmt is not propagated in the nested
statement. Thus, it fallbacks to the "use the whole query string"
default which is reported to pgss.

This patch creates tests to cover the current behaviour.
---
 .../expected/level_tracking.out               | 1012 +++++++++++++++++
 .../pg_stat_statements/sql/level_tracking.sql |  201 +++-
 2 files changed, 1212 insertions(+), 1 deletion(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index bb65e98ce09..9463c68154b 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -112,6 +112,1018 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (2 rows)
 
+-- Explain - all-level tracking.
+CREATE TABLE test_table (x int);
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab;
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+explain (costs off) (TABLE test_table);
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                   query                                                    
+----------+-------+------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (SELECT $1, $2)
+ t        |     1 | explain (costs off) (TABLE test_table)
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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                                                                   +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1 UNION SELECT $2
+ 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+
+          |       |     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);
+(23 rows)
+
+-- Explain - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab;
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+explain (costs off) (TABLE test_table);
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                   query                                                    
+----------+-------+------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (SELECT $1, $2)
+ t        |     1 | explain (costs off) (TABLE test_table)
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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                                                                   +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1 UNION SELECT $2
+ 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)
+(12 rows)
+
+-- Explain - all-level tracking with multi statement.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES (1), (2);
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+           QUERY PLAN            
+---------------------------------
+ Insert on stats_track_tab
+   ->  Values Scan on "*VALUES*"
+(2 rows)
+
+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;
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+         QUERY PLAN         
+----------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+           QUERY PLAN            
+---------------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2), (3)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ 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
+ t        |     1 | explain (costs off) (SELECT $1, $2, $3, $4)
+ t        |     1 | explain (costs off) (TABLE test_table)
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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                                                                                         +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1, $2
+ t        |     1 | explain (costs off) SELECT $1, $2 UNION SELECT $3, $4
+ t        |     1 | explain (costs off) SELECT $1, $2, $3, $4, $5
+ t        |     1 | explain (costs off) TABLE stats_track_tab
+ 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));
+(37 rows)
+
+-- Explain - top-level tracking with multi statement.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES ((1), (2));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+ERROR:  INSERT has more expressions than target columns
+LINE 1: ...n (costs off) INSERT INTO stats_track_tab VALUES ((1), (2));
+                                                                   ^
+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;
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+         QUERY PLAN         
+----------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+           QUERY PLAN            
+---------------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2), (3)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ 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
+ t        |     1 | explain (costs off) (SELECT $1, $2, $3, $4)
+ t        |     1 | explain (costs off) (TABLE test_table)
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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))
+ 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
+ t        |     1 | explain (costs off) SELECT $1, $2 UNION SELECT $3, $4
+ t        |     1 | explain (costs off) SELECT $1, $2, $3, $4, $5
+ t        |     1 | explain (costs off) TABLE stats_track_tab
+ 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)
+(18 rows)
+
+-- Explain with CTE - all-level tracking
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) WITH a AS (select 4) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) 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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                              query                                                               
+----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (WITH a AS (select $1) (SELECT $2, $3))
+ t        |     1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab
+ t        |     1 | explain (costs off) WITH a AS (select $1) INSERT INTO stats_track_tab VALUES (($2))
+ t        |     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+
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                                         +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ 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+
+          |       |     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;
+(15 rows)
+
+-- Explain with CTE - top-level tracking
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) WITH a AS (select 4) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) 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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                              query                                                               
+----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (WITH a AS (select $1) (SELECT $2, $3))
+ t        |     1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab
+ t        |     1 | explain (costs off) WITH a AS (select $1) INSERT INTO stats_track_tab VALUES (($2))
+ t        |     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+
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                                         +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ 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
+(8 rows)
+
+-- Explain analyze, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                     query                                                     
+----------+-------+---------------------------------------------------------------------------------------------------------------
+ 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;
+(4 rows)
+
+-- Explain analyze, top tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                     query                                                     
+----------+-------+---------------------------------------------------------------------------------------------------------------
+ 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
+(3 rows)
+
+-- Create Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                             query                                              
+----------+-------+------------------------------------------------------------------------------------------------
+ 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)
+
+-- Create Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE MATERIALIZED VIEW pgss_materialized_view_2 AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                              query                                               
+----------+-------+--------------------------------------------------------------------------------------------------
+ t        |     1 | CREATE MATERIALIZED VIEW pgss_materialized_view_2 AS SELECT * FROM generate_series($1, $2) as id
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Refresh Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Refresh Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT $1
+ 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)
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_ctas_3 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_4 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_3 AS SELECT $1
+ t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_4 AS EXECUTE test_prepare_pgss
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Explain with Create Table As - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
+ 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)
+
+-- Explain with Create Table As - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
+ 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)
+
+-- Declare cursor, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                          query                          
+----------+-------+---------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ 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)
+
+-- Declare cursor, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                          query                          
+----------+-------+---------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ 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)
+
 -- 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 65a17147a5a..313985d1b68 100644
--- a/contrib/pg_stat_statements/sql/level_tracking.sql
+++ b/contrib/pg_stat_statements/sql/level_tracking.sql
@@ -32,7 +32,6 @@ BEGIN
 END; $$;
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
-
 -- Procedure with multiple utility statements.
 CREATE OR REPLACE PROCEDURE proc_with_utility_stmt()
 LANGUAGE SQL
@@ -55,6 +54,206 @@ CALL proc_with_utility_stmt();
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
 
+-- Explain - all-level tracking.
+CREATE TABLE test_table (x int);
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1;
+explain (costs off) (SELECT 1, 2);
+explain (costs off) TABLE stats_track_tab;
+explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1);
+explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) DELETE FROM stats_track_tab;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((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 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1;
+explain (costs off) (SELECT 1, 2);
+explain (costs off) TABLE stats_track_tab;
+explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1);
+explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) DELETE FROM stats_track_tab;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((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 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain - all-level tracking with multi statement.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES (1), (2);
+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;
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain - top-level tracking with multi statement.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES ((1), (2));
+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;
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with CTE - all-level tracking
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) WITH a AS (select 4) SELECT 1;
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+explain (costs off) WITH a AS (select 4) 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) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with CTE - top-level tracking
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) WITH a AS (select 4) SELECT 1;
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+explain (costs off) WITH a AS (select 4) 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) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain analyze, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain analyze, top tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE MATERIALIZED VIEW pgss_materialized_view_2 AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Refresh Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Refresh Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_ctas_3 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_4 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with Create Table As - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with Create Table As - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Declare cursor, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Declare cursor, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
 -- DO block - top-level tracking without utility.
 SET pg_stat_statements.track = 'top';
 SET pg_stat_statements.track_utility = FALSE;
-- 
2.39.5 (Apple Git-154)

From 9593de5aedd8934ae6b12895e015cbe31ff5a946 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Thu, 3 Oct 2024 08:52:02 +0200
Subject: Track location to extract relevant part in nested statement

Previously, Query generated through transform would have unset
stmt_location. Extensions relying on the statement location to extract
the relevant part of the statement would fallback to use the whole
statement instead, thus showing the same string in the top and
nested level which was a source of confusion.

This patch fixes the issue by keeping track of the statement locations
and propagate it to Query during transform, allowing pgss to only show
the relevant part of the query for nested query.
---
 .../expected/level_tracking.out               | 118 +++++++++---------
 .../pg_stat_statements/expected/planning.out  |  10 +-
 .../pg_stat_statements/expected/select.out    |   2 +-
 .../pg_stat_statements/expected/utility.out   |   2 +-
 contrib/pg_stat_statements/sql/planning.sql   |   4 +-
 src/backend/optimizer/util/clauses.c          |   2 +-
 src/backend/parser/analyze.c                  |  96 +++++++++-----
 src/backend/parser/gram.y                     |  71 ++++++++++-
 src/backend/parser/parse_merge.c              |   2 +
 src/include/nodes/parsenodes.h                |  10 ++
 src/include/parser/analyze.h                  |   3 +-
 src/include/parser/parse_node.h               |   2 +
 12 files changed, 217 insertions(+), 105 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 8c5c1b5b55d..183a8cfc195 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 | 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 | SELECT $1, $2
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | TABLE test_table
+ f        |     1 | UPDATE stats_track_tab SET x=$1 WHERE x=$2
+ f        |     1 | VALUES($1)
+ f        |     1 | VALUES($1, $2)
 (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) UNION SELECT $4, $5, $6
+ 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
+ f        |     1 | SELECT $1, $2, $3, $4
+ f        |     1 | SELECT $1, $2, $3, $4, $5
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | TABLE test_table
+ 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)
+ f        |     1 | VALUES($1, $2)
 (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.
@@ -1000,7 +998,7 @@ 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
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
- f        |     1 | PREPARE test_prepare_pgss AS select generate_series($1, $2)
+ f        |     1 | select generate_series($1, $2)
 (4 rows)
 
 -- Create Table As, top-level tracking.
@@ -1151,25 +1149,25 @@ COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout;
 2
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY toplevel desc, query COLLATE "C";
- toplevel | calls |                                  query                                   
-----------+-------+--------------------------------------------------------------------------
+ toplevel | calls |                                  query                                  
+----------+-------+-------------------------------------------------------------------------
  t        |     1 | COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout
  t        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout
- t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id          +
-          |       |     WHEN MATCHED THEN UPDATE SET x = id                                 +
+ t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id         +
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                +
           |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout
  t        |     1 | COPY (SELECT 1 UNION SELECT 2) to stdout
  t        |     1 | COPY (SELECT 1) to stdout
  t        |     1 | COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
- f        |     1 | COPY (DELETE FROM stats_track_tab WHERE x=$1 RETURNING x) to stdout
- f        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES ($1) RETURNING x) to stdout
- f        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT $1 id) ON x = id         +
-          |       |     WHEN MATCHED THEN UPDATE SET x = id                                 +
-          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout
- f        |     1 | COPY (SELECT $1 UNION SELECT $2) to stdout
- f        |     1 | COPY (SELECT $1) to stdout
- f        |     1 | COPY (UPDATE stats_track_tab SET x=$1 WHERE x=$2 RETURNING x) to stdout
+ f        |     1 | DELETE FROM stats_track_tab WHERE x=$1 RETURNING x
+ f        |     1 | INSERT INTO stats_track_tab (x) VALUES ($1) RETURNING x
+ f        |     1 | MERGE INTO stats_track_tab USING (SELECT $1 id) ON x = id              +
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x
+ f        |     1 | SELECT $1
+ f        |     1 | SELECT $1 UNION SELECT $2
+ f        |     1 | UPDATE stats_track_tab SET x=$1 WHERE x=$2 RETURNING x
 (13 rows)
 
 -- COPY - top-level tracking.
diff --git a/contrib/pg_stat_statements/expected/planning.out b/contrib/pg_stat_statements/expected/planning.out
index 9effd11fdc8..3ee1928cbe9 100644
--- a/contrib/pg_stat_statements/expected/planning.out
+++ b/contrib/pg_stat_statements/expected/planning.out
@@ -58,7 +58,7 @@ SELECT 42;
 (1 row)
 
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
  plans | calls | rows |                          query                           
 -------+-------+------+----------------------------------------------------------
      0 |     1 |    0 | ALTER TABLE stats_plan_test ADD COLUMN x int
@@ -72,10 +72,10 @@ SELECT plans, calls, rows, query FROM pg_stat_statements
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
- plans_ok | calls | rows |                         query                         
-----------+-------+------+-------------------------------------------------------
- t        |     4 |    4 | PREPARE prep1 AS SELECT COUNT(*) FROM stats_plan_test
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
+ plans_ok | calls | rows |                query                 
+----------+-------+------+--------------------------------------
+ t        |     4 |    4 | SELECT COUNT(*) FROM stats_plan_test
 (1 row)
 
 -- Cleanup
diff --git a/contrib/pg_stat_statements/expected/select.out b/contrib/pg_stat_statements/expected/select.out
index dd6c756f67d..217a2c0b2bc 100644
--- a/contrib/pg_stat_statements/expected/select.out
+++ b/contrib/pg_stat_statements/expected/select.out
@@ -127,7 +127,6 @@ DEALLOCATE pgss_test;
 SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
  calls | rows |                                    query                                     
 -------+------+------------------------------------------------------------------------------
-     1 |    1 | PREPARE pgss_test (int) AS SELECT $1, $2 LIMIT $3
      4 |    4 | SELECT $1                                                                   +
        |      |   -- multiline                                                              +
        |      |   AS "text"
@@ -137,6 +136,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
      2 |    2 | SELECT $1 AS "int"
      1 |    2 | SELECT $1 AS i UNION SELECT $2 ORDER BY i
      1 |    1 | SELECT $1 || $2
+     1 |    1 | SELECT $1, $2 LIMIT $3
      0 |    0 | SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C"
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
      1 |    2 | WITH t(f) AS (                                                              +
diff --git a/contrib/pg_stat_statements/expected/utility.out b/contrib/pg_stat_statements/expected/utility.out
index 060d4416dd7..aa4f0f7e628 100644
--- a/contrib/pg_stat_statements/expected/utility.out
+++ b/contrib/pg_stat_statements/expected/utility.out
@@ -540,7 +540,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
 -------+------+----------------------------------------------------
      2 |    0 | DEALLOCATE $1
      2 |    0 | DEALLOCATE ALL
-     2 |    2 | PREPARE stat_select AS SELECT $1 AS a
+     2 |    2 | SELECT $1 AS a
      1 |    1 | SELECT $1 as a
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (5 rows)
diff --git a/contrib/pg_stat_statements/sql/planning.sql b/contrib/pg_stat_statements/sql/planning.sql
index 46f5d9b951c..9cfe206b3b0 100644
--- a/contrib/pg_stat_statements/sql/planning.sql
+++ b/contrib/pg_stat_statements/sql/planning.sql
@@ -20,11 +20,11 @@ SELECT 42;
 SELECT 42;
 SELECT 42;
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 
 -- Cleanup
 DROP TABLE stats_plan_test;
diff --git a/src/backend/optimizer/util/clauses.c b/src/backend/optimizer/util/clauses.c
index b4e085e9d4b..8fc4da11483 100644
--- a/src/backend/optimizer/util/clauses.c
+++ b/src/backend/optimizer/util/clauses.c
@@ -4674,7 +4674,7 @@ inline_function(Oid funcid, Oid result_type, Oid result_collid,
 		pstate->p_sourcetext = src;
 		sql_fn_parser_setup(pstate, pinfo);
 
-		querytree = transformTopLevelStmt(pstate, linitial(raw_parsetree_list));
+		querytree = transformOptionalSelectInto(pstate, ((RawStmt *) linitial(raw_parsetree_list))->stmt);
 
 		free_parsestate(pstate);
 	}
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index 2d3d8fcf769..0c7e63b783f 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -58,7 +58,6 @@
 /* Hook for plugins to get control at end of parse analysis */
 post_parse_analyze_hook_type post_parse_analyze_hook = NULL;
 
-static Query *transformOptionalSelectInto(ParseState *pstate, Node *parseTree);
 static Query *transformDeleteStmt(ParseState *pstate, DeleteStmt *stmt);
 static Query *transformInsertStmt(ParseState *pstate, InsertStmt *stmt);
 static OnConflictExpr *transformOnConflictClause(ParseState *pstate,
@@ -113,13 +112,15 @@ 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);
 
 	pstate->p_queryEnv = queryEnv;
 
-	query = transformTopLevelStmt(pstate, parseTree);
+	query = transformOptionalSelectInto(pstate, parseTree->stmt);
 
 	if (IsQueryIdEnabled())
 		jstate = JumbleQuery(query);
@@ -153,12 +154,14 @@ 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);
 
 	pstate->p_queryEnv = queryEnv;
 
-	query = transformTopLevelStmt(pstate, parseTree);
+	query = transformOptionalSelectInto(pstate, parseTree->stmt);
 
 	/* make sure all is well with parameter types */
 	check_variable_parameters(pstate, query);
@@ -195,10 +198,12 @@ 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);
 
-	query = transformTopLevelStmt(pstate, parseTree);
+	query = transformOptionalSelectInto(pstate, parseTree->stmt);
 
 	if (IsQueryIdEnabled())
 		jstate = JumbleQuery(query);
@@ -238,27 +243,6 @@ parse_sub_analyze(Node *parseTree, ParseState *parentParseState,
 	return query;
 }
 
-/*
- * transformTopLevelStmt -
- *	  transform a Parse tree into a Query tree.
- *
- * This function is just responsible for transferring statement location data
- * from the RawStmt into the finished Query.
- */
-Query *
-transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree)
-{
-	Query	   *result;
-
-	/* We're at top level, so allow SELECT INTO */
-	result = transformOptionalSelectInto(pstate, parseTree->stmt);
-
-	result->stmt_location = parseTree->stmt_location;
-	result->stmt_len = parseTree->stmt_len;
-
-	return result;
-}
-
 /*
  * transformOptionalSelectInto -
  *	  If SELECT has INTO, convert it to CREATE TABLE AS.
@@ -269,7 +253,7 @@ transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree)
  * of the parse tree, and so we only try it before entering the recursive
  * transformStmt() processing.
  */
-static Query *
+Query *
 transformOptionalSelectInto(ParseState *pstate, Node *parseTree)
 {
 	if (IsA(parseTree, SelectStmt))
@@ -417,7 +401,9 @@ transformStmt(ParseState *pstate, Node *parseTree)
 			 */
 			result = makeNode(Query);
 			result->commandType = CMD_UTILITY;
-			result->utilityStmt = (Node *) parseTree;
+			result->utilityStmt = parseTree;
+			result->stmt_location = pstate->p_stmt_location;
+			result->stmt_len = pstate->p_stmt_len;
 			break;
 	}
 
@@ -506,6 +492,37 @@ analyze_requires_snapshot(RawStmt *parseTree)
 	return stmt_requires_parse_analysis(parseTree);
 }
 
+/*
+ * setQueryStmtLen
+ *		Set stmt_len in Query.
+ *
+ * Some statements, like PreparableStmt, can be located within parentheses.
+ * For example "(SELECT 1)" or "COPY (UPDATE ...) to x;". For those, we can't
+ * use the whole string from the statement's location or the SQL string will
+ * yield "SELECT 1)". The parser will set stmt_len, reflecting the size of the
+ * statement within the parentheses. Thus, when stmt_len is available, we use it
+ * for the Query's stmt_len.
+ *
+ * For other cases, the parser can't provide the length of individual statements.
+ * However, we have the statement's location plus the length (p_stmt_len) and
+ * location (p_stmt_location) of the top level RawStmt, stored in pstate. Thus,
+ * the statement's length is the RawStmt's length minus how much we've advanced
+ * in the RawStmt's string.
+ */
+void
+setQueryStmtLen(ParseState *pstate, Query *qry, int stmt_len)
+{
+	if (stmt_len > 0)
+		/* Statement's length is known, use it */
+		qry->stmt_len = stmt_len;
+	else
+		/*
+		 * Compute the statement's length from statement's location and RawStmt's
+		 * length and location
+		 */
+		qry->stmt_len = pstate->p_stmt_len - (qry->stmt_location - pstate->p_stmt_location);
+}
+
 /*
  * transformDeleteStmt -
  *	  transforms a Delete Statement
@@ -518,6 +535,8 @@ transformDeleteStmt(ParseState *pstate, DeleteStmt *stmt)
 	Node	   *qual;
 
 	qry->commandType = CMD_DELETE;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 
 	/* process the WITH clause independently of all else */
 	if (stmt->withClause)
@@ -606,6 +625,8 @@ transformInsertStmt(ParseState *pstate, InsertStmt *stmt)
 	Assert(pstate->p_ctenamespace == NIL);
 
 	qry->commandType = CMD_INSERT;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 	pstate->p_is_insert = true;
 
 	/* process the WITH clause independently of all else */
@@ -1331,7 +1352,6 @@ count_rowexpr_columns(ParseState *pstate, Node *expr)
 	return -1;
 }
 
-
 /*
  * transformSelectStmt -
  *	  transforms a Select Statement
@@ -1347,6 +1367,8 @@ transformSelectStmt(ParseState *pstate, SelectStmt *stmt)
 	ListCell   *l;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 
 	/* process the WITH clause independently of all else */
 	if (stmt->withClause)
@@ -1499,6 +1521,8 @@ transformValuesClause(ParseState *pstate, SelectStmt *stmt)
 	int			i;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 
 	/* Most SELECT stuff doesn't apply in a VALUES clause */
 	Assert(stmt->distinctClause == NIL);
@@ -1730,6 +1754,8 @@ transformSetOperationStmt(ParseState *pstate, SelectStmt *stmt)
 	int			tllen;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 
 	/*
 	 * Find leftmost leaf SelectStmt.  We currently only need to do this in
@@ -2397,6 +2423,8 @@ transformReturnStmt(ParseState *pstate, ReturnStmt *stmt)
 	Query	   *qry = makeNode(Query);
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = pstate->p_stmt_location;
+	qry->stmt_len = pstate->p_stmt_len;
 	qry->isReturn = true;
 
 	qry->targetList = list_make1(makeTargetEntry((Expr *) transformExpr(pstate, stmt->returnval, EXPR_KIND_SELECT_TARGET),
@@ -2430,6 +2458,8 @@ transformUpdateStmt(ParseState *pstate, UpdateStmt *stmt)
 	Node	   *qual;
 
 	qry->commandType = CMD_UPDATE;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 	pstate->p_is_insert = false;
 
 	/* process the WITH clause independently of all else */
@@ -2677,6 +2707,8 @@ transformPLAssignStmt(ParseState *pstate, PLAssignStmt *stmt)
 	 * consider WITH or INTO, and we build a targetlist our own way.
 	 */
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = pstate->p_stmt_location;
+	qry->stmt_len = pstate->p_stmt_len;
 	pstate->p_is_insert = false;
 
 	/* make FOR UPDATE/FOR SHARE info available to addRangeTableEntry */
@@ -2948,6 +2980,8 @@ transformDeclareCursorStmt(ParseState *pstate, DeclareCursorStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3003,6 +3037,8 @@ transformExplainStmt(ParseState *pstate, ExplainStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3083,6 +3119,8 @@ transformCreateTableAsStmt(ParseState *pstate, CreateTableAsStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3207,6 +3245,8 @@ transformCallStmt(ParseState *pstate, CallStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 4aa8646af7b..d00f4b47c4a 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -167,6 +167,7 @@ static void base_yyerror(YYLTYPE *yylloc, core_yyscan_t yyscanner,
 						 const char *msg);
 static RawStmt *makeRawStmt(Node *stmt, int stmt_location);
 static void updateRawStmtEnd(RawStmt *rs, int end_location);
+static void updatePreparableStmtEnd(Node *n, int end_location);
 static Node *makeColumnRef(char *colname, List *indirection,
 						   int location, core_yyscan_t yyscanner);
 static Node *makeTypeCast(Node *arg, TypeName *typename, int location);
@@ -190,7 +191,7 @@ static void insertSelectOptions(SelectStmt *stmt,
 								SelectLimit *limitClause,
 								WithClause *withClause,
 								core_yyscan_t yyscanner);
-static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg);
+static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location);
 static Node *doNegate(Node *n, int location);
 static void doNegateFloat(Float *v);
 static Node *makeAndExpr(Node *lexpr, Node *rexpr, int location);
@@ -3403,6 +3404,7 @@ CopyStmt:	COPY opt_binary qualified_name opt_column_list
 				{
 					CopyStmt *n = makeNode(CopyStmt);
 
+					updatePreparableStmtEnd($3, @4);
 					n->relation = NULL;
 					n->query = $3;
 					n->attlist = NIL;
@@ -12170,6 +12172,9 @@ InsertStmt:
 					$5->onConflictClause = $6;
 					$5->returningList = $7;
 					$5->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					$5->stmt_location = @$;
 					$$ = (Node *) $5;
 				}
 		;
@@ -12323,6 +12328,9 @@ DeleteStmt: opt_with_clause DELETE_P FROM relation_expr_opt_alias
 					n->whereClause = $6;
 					n->returningList = $7;
 					n->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					n->stmt_location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12397,6 +12405,9 @@ UpdateStmt: opt_with_clause UPDATE relation_expr_opt_alias
 					n->whereClause = $7;
 					n->returningList = $8;
 					n->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					n->stmt_location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12474,6 +12485,9 @@ MergeStmt:
 					m->joinCondition = $8;
 					m->mergeWhenClauses = $9;
 					m->returningList = $10;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					m->stmt_location = @$;
 
 					$$ = (Node *) m;
 				}
@@ -12714,7 +12728,12 @@ SelectStmt: select_no_parens			%prec UMINUS
 		;
 
 select_with_parens:
-			'(' select_no_parens ')'				{ $$ = $2; }
+			'(' select_no_parens ')'
+				{
+					SelectStmt *n = (SelectStmt *) $2;
+					n->stmt_len = @3 - @2;
+					$$ = $2;
+				}
 			| '(' select_with_parens ')'			{ $$ = $2; }
 		;
 
@@ -12836,6 +12855,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| SELECT distinct_clause target_list
@@ -12853,6 +12873,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| values_clause							{ $$ = $1; }
@@ -12873,19 +12894,20 @@ simple_select:
 
 					n->targetList = list_make1(rt);
 					n->fromClause = list_make1($2);
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| select_clause UNION set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause INTERSECT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause EXCEPT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 		;
 
@@ -13443,6 +13465,7 @@ values_clause:
 				{
 					SelectStmt *n = makeNode(SelectStmt);
 
+					n->stmt_location = @1;
 					n->valuesLists = list_make1($3);
 					$$ = (Node *) n;
 				}
@@ -18585,6 +18608,39 @@ updateRawStmtEnd(RawStmt *rs, int end_location)
 	rs->stmt_len = end_location - rs->stmt_location;
 }
 
+/* Adjust a PreparableStmt to reflect that it doesn't run to the end of the string */
+static void
+updatePreparableStmtEnd(Node *n, int end_location)
+{
+	if (IsA(n, SelectStmt))
+	{
+		SelectStmt *stmt = (SelectStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, InsertStmt))
+	{
+		InsertStmt *stmt = (InsertStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, UpdateStmt))
+	{
+		UpdateStmt *stmt = (UpdateStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, DeleteStmt))
+	{
+		DeleteStmt *stmt = (DeleteStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, MergeStmt))
+	{
+		MergeStmt *stmt = (MergeStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else
+		elog(ERROR, "unexpected node type %d", (int) n->type);
+}
+
 static Node *
 makeColumnRef(char *colname, List *indirection,
 			  int location, core_yyscan_t yyscanner)
@@ -18963,11 +19019,13 @@ insertSelectOptions(SelectStmt *stmt,
 					 errmsg("multiple WITH clauses not allowed"),
 					 parser_errposition(exprLocation((Node *) withClause))));
 		stmt->withClause = withClause;
+		/* Update SelectStmt's location to the start of the with clause */
+		stmt->stmt_location = withClause->location;
 	}
 }
 
 static Node *
-makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
+makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location)
 {
 	SelectStmt *n = makeNode(SelectStmt);
 
@@ -18975,6 +19033,7 @@ makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
 	n->all = all;
 	n->larg = (SelectStmt *) larg;
 	n->rarg = (SelectStmt *) rarg;
+	n->stmt_location = location;
 	return (Node *) n;
 }
 
diff --git a/src/backend/parser/parse_merge.c b/src/backend/parser/parse_merge.c
index 87df79027d7..fdfa61868cf 100644
--- a/src/backend/parser/parse_merge.c
+++ b/src/backend/parser/parse_merge.c
@@ -118,6 +118,8 @@ transformMergeStmt(ParseState *pstate, MergeStmt *stmt)
 	Assert(pstate->p_ctenamespace == NIL);
 
 	qry->commandType = CMD_MERGE;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 	qry->hasRecursive = false;
 
 	/* process the WITH clause independently of all else */
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index c92cef3d16d..b40b661ec8a 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -2046,6 +2046,8 @@ typedef struct InsertStmt
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
 	OverridingKind override;	/* OVERRIDING clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } InsertStmt;
 
 /* ----------------------
@@ -2060,6 +2062,8 @@ typedef struct DeleteStmt
 	Node	   *whereClause;	/* qualifications */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } DeleteStmt;
 
 /* ----------------------
@@ -2075,6 +2079,8 @@ typedef struct UpdateStmt
 	List	   *fromClause;		/* optional from clause for more tables */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } UpdateStmt;
 
 /* ----------------------
@@ -2090,6 +2096,8 @@ typedef struct MergeStmt
 	List	   *mergeWhenClauses;	/* list of MergeWhenClause(es) */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } MergeStmt;
 
 /* ----------------------
@@ -2159,6 +2167,8 @@ typedef struct SelectStmt
 	bool		all;			/* ALL specified? */
 	struct SelectStmt *larg;	/* left child */
 	struct SelectStmt *rarg;	/* right child */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 	/* Eventually add fields for CORRESPONDING spec here */
 } SelectStmt;
 
diff --git a/src/include/parser/analyze.h b/src/include/parser/analyze.h
index 28b66fccb43..8ba4e050aff 100644
--- a/src/include/parser/analyze.h
+++ b/src/include/parser/analyze.h
@@ -46,8 +46,9 @@ extern List *transformUpdateTargetList(ParseState *pstate,
 									   List *origTlist);
 extern List *transformReturningList(ParseState *pstate, List *returningList,
 									ParseExprKind exprKind);
-extern Query *transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree);
+extern Query *transformOptionalSelectInto(ParseState *pstate, Node *parseTree);
 extern Query *transformStmt(ParseState *pstate, Node *parseTree);
+extern void setQueryStmtLen(ParseState *pstate, Query *qry, int stmt_len);
 
 extern bool stmt_requires_parse_analysis(RawStmt *parseTree);
 extern bool analyze_requires_snapshot(RawStmt *parseTree);
diff --git a/src/include/parser/parse_node.h b/src/include/parser/parse_node.h
index 543df568147..ba572b3aea2 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 */
-- 
2.39.5 (Apple Git-154)

From fafccd6dfffb2dfd9df979e17f077d5400633b89 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
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 183a8cfc195..85f01f8a090 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -881,7 +881,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';
@@ -927,7 +928,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';
@@ -961,7 +963,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';
@@ -998,8 +1001,9 @@ 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
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | SELECT $1
  f        |     1 | select generate_series($1, $2)
-(4 rows)
+(5 rows)
 
 -- Create Table As, top-level tracking.
 SET pg_stat_statements.track = 'top';
@@ -1040,7 +1044,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';
@@ -1091,7 +1096,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 18a5af6b919..c81221cdbeb 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 dcbdaa03885..d2eef8097cf 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 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)

From 09e6670ccb6d63eeeafa5aa8af4dc9ab3c6468fb Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Thu, 3 Oct 2024 10:21:44 +0200
Subject: Use view's definition as query string on a materialized view refresh

When creating a materialized view, the first refresh will have the
"Select" part of the statement as a query string. On subsequent refresh,
the "REFRESH MATERIALIZED" utility statement will be passed as query
string. This causes pgss to track both the top query and nested query as
a refresh.

This patch changes the query string on a refresh to fetch the view
definition instead. This will allow pgss to display the correct
statement when tracking refresh's nested query.
---
 .../expected/level_tracking.out               | 10 ++++----
 src/backend/commands/matview.c                | 24 +++++++++++++------
 src/backend/utils/adt/ruleutils.c             | 13 ++++++++++
 src/include/utils/ruleutils.h                 |  2 +-
 4 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 85f01f8a090..93afd7e51cc 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -959,12 +959,14 @@ SELECT pg_stat_statements_reset() IS NOT NULL AS t;
 REFRESH MATERIALIZED VIEW pgss_materialized_view;
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY toplevel desc, query COLLATE "C";
- toplevel | calls |                       query                        
-----------+-------+----------------------------------------------------
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
  t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
- f        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view;
-(3 rows)
+ f        |     1 | SELECT * FROM pg_catalog.pg_rewrite WHERE ev_class = $1 AND rulename = $2
+ f        |     1 | SELECT id                                                                +
+          |       |    FROM generate_series(1, 5) id(id);
+(4 rows)
 
 -- Refresh Materialized View, top-level tracking.
 SET pg_stat_statements.track = 'top';
diff --git a/src/backend/commands/matview.c b/src/backend/commands/matview.c
index 7cc68338837..92d5dd43864 100644
--- a/src/backend/commands/matview.c
+++ b/src/backend/commands/matview.c
@@ -39,6 +39,7 @@
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
+#include "utils/ruleutils.h"
 #include "utils/snapmgr.h"
 #include "utils/syscache.h"
 
@@ -61,8 +62,7 @@ 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,
-									   ParseState *pstate, const char *queryString,
-									   bool is_create);
+									   ParseState *pstate, 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);
@@ -327,11 +327,21 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData,
 	if (!skipData)
 	{
 		DestReceiver *dest;
-		const char *queryString = pstate->p_sourcetext;
+		ParseState *refresh_pstate = pstate;
+
+		/*
+		 * On refresh, the pstate's source text will be the refresh utility
+		 * statement. We need to fetch the the view definition to get the
+		 * query executed by the refresh.
+		 */
+		if (!is_create)
+		{
+			refresh_pstate = make_parsestate(NULL);
+			refresh_pstate->p_sourcetext = pg_get_viewdef_string(matviewOid, false);
+		}
 
 		dest = CreateTransientRelDestReceiver(OIDNewHeap);
-		processed = refresh_matview_datafill(dest, dataQuery, pstate,
-											 queryString, is_create);
+		processed = refresh_matview_datafill(dest, dataQuery, refresh_pstate, is_create);
 	}
 
 	/* Make the matview match the newly generated data. */
@@ -406,8 +416,7 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData,
  */
 static uint64
 refresh_matview_datafill(DestReceiver *dest, Query *query,
-						 ParseState *pstate, const char *queryString,
-						 bool is_create)
+						 ParseState *pstate, bool is_create)
 {
 	List	   *rewritten;
 	PlannedStmt *plan;
@@ -415,6 +424,7 @@ refresh_matview_datafill(DestReceiver *dest, Query *query,
 	Query	   *copied_query;
 	uint64		processed;
 	JumbleState *jstate = NULL;
+	const char *queryString = pstate->p_sourcetext;
 
 	/* Lock and rewrite, using a copy to preserve the original query. */
 	copied_query = copyObject(query);
diff --git a/src/backend/utils/adt/ruleutils.c b/src/backend/utils/adt/ruleutils.c
index 2177d17e278..0bf85cbb759 100644
--- a/src/backend/utils/adt/ruleutils.c
+++ b/src/backend/utils/adt/ruleutils.c
@@ -776,6 +776,19 @@ pg_get_viewdef_name_ext(PG_FUNCTION_ARGS)
 	PG_RETURN_TEXT_P(string_to_text(res));
 }
 
+/*
+ * Internal version of pg_get_viewdef
+ */
+char *
+pg_get_viewdef_string(Oid viewoid, bool pretty)
+{
+	int			prettyFlags;
+
+	prettyFlags = GET_PRETTY_FLAGS(pretty);
+
+	return pg_get_viewdef_worker(viewoid, prettyFlags, WRAP_COLUMN_DEFAULT);
+}
+
 /*
  * Common code for by-OID and by-name variants of pg_get_viewdef
  */
diff --git a/src/include/utils/ruleutils.h b/src/include/utils/ruleutils.h
index 161fb5ef02c..72177b9dce3 100644
--- a/src/include/utils/ruleutils.h
+++ b/src/include/utils/ruleutils.h
@@ -29,7 +29,7 @@ extern char *pg_get_indexdef_columns(Oid indexrelid, bool pretty);
 extern char *pg_get_indexdef_columns_extended(Oid indexrelid,
 											  bits16 flags);
 extern char *pg_get_querydef(Query *query, bool pretty);
-
+extern char *pg_get_viewdef_string(Oid viewoid, bool pretty);
 extern char *pg_get_partkeydef_columns(Oid relid, bool pretty);
 extern char *pg_get_partconstrdef_string(Oid partitionId, char *aliasname);
 
-- 
2.39.5 (Apple Git-154)

From 8948a8f94d45f134bf1047b496c4f52580c955db Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Mon, 7 Oct 2024 10:45:49 +0200
Subject: Add tests covering pgss nested queries

What pgss reports for nested statements can be confusing. Some
statements like CreateTableAs, DeclareCursor and CreateMaterializedView
don't jumble the nested query and thus won't report it in pgss. Explain
explicitely Jumble the nested query and will appear in pgss. However,
the reported query string will be the same as the top level statement.

Multi statements queries are also not handled correctly. The statement
length used in the top RawStmt is not propagated in the nested
statement. Thus, it fallbacks to the "use the whole query string"
default which is reported to pgss.

This patch creates tests to cover the current behaviour.
---
 .../expected/level_tracking.out               | 1100 +++++++++++++++++
 .../pg_stat_statements/sql/level_tracking.sql |  229 +++-
 2 files changed, 1328 insertions(+), 1 deletion(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index bb65e98ce09..8c5c1b5b55d 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -112,6 +112,1106 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (2 rows)
 
+-- Explain - all-level tracking.
+CREATE TABLE test_table (x int);
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab;
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+explain (costs off) (TABLE test_table);
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                   query                                                    
+----------+-------+------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (SELECT $1, $2)
+ t        |     1 | explain (costs off) (TABLE test_table)
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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                                                                   +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1 UNION SELECT $2
+ 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+
+          |       |     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);
+(23 rows)
+
+-- Explain - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab;
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+explain (costs off) (TABLE test_table);
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                   query                                                    
+----------+-------+------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (SELECT $1, $2)
+ t        |     1 | explain (costs off) (TABLE test_table)
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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                                                                   +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1 UNION SELECT $2
+ 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)
+(12 rows)
+
+-- Explain - all-level tracking with multi statement.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES (1), (2);
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+           QUERY PLAN            
+---------------------------------
+ Insert on stats_track_tab
+   ->  Values Scan on "*VALUES*"
+(2 rows)
+
+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;
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+         QUERY PLAN         
+----------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+           QUERY PLAN            
+---------------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2), (3)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ 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
+ t        |     1 | explain (costs off) (SELECT $1, $2, $3, $4)
+ t        |     1 | explain (costs off) (TABLE test_table)
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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                                                                                         +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1, $2
+ t        |     1 | explain (costs off) SELECT $1, $2 UNION SELECT $3, $4
+ t        |     1 | explain (costs off) SELECT $1, $2, $3, $4, $5
+ t        |     1 | explain (costs off) TABLE stats_track_tab
+ 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));
+(37 rows)
+
+-- Explain - top-level tracking with multi statement.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES ((1), (2));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+ERROR:  INSERT has more expressions than target columns
+LINE 1: ...n (costs off) INSERT INTO stats_track_tab VALUES ((1), (2));
+                                                                   ^
+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;
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+         QUERY PLAN         
+----------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+           QUERY PLAN            
+---------------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2), (3)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ 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
+ t        |     1 | explain (costs off) (SELECT $1, $2, $3, $4)
+ t        |     1 | explain (costs off) (TABLE test_table)
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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))
+ 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
+ t        |     1 | explain (costs off) SELECT $1, $2 UNION SELECT $3, $4
+ t        |     1 | explain (costs off) SELECT $1, $2, $3, $4, $5
+ t        |     1 | explain (costs off) TABLE stats_track_tab
+ 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)
+(18 rows)
+
+-- Explain with CTE - all-level tracking
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) WITH a AS (select 4) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) 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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                              query                                                               
+----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (WITH a AS (select $1) (SELECT $2, $3))
+ t        |     1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab
+ t        |     1 | explain (costs off) WITH a AS (select $1) INSERT INTO stats_track_tab VALUES (($2))
+ t        |     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+
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                                         +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ 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+
+          |       |     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;
+(15 rows)
+
+-- Explain with CTE - top-level tracking
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) WITH a AS (select 4) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) 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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                              query                                                               
+----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (WITH a AS (select $1) (SELECT $2, $3))
+ t        |     1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab
+ t        |     1 | explain (costs off) WITH a AS (select $1) INSERT INTO stats_track_tab VALUES (($2))
+ t        |     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+
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                                         +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ 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
+(8 rows)
+
+-- Explain analyze, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                     query                                                     
+----------+-------+---------------------------------------------------------------------------------------------------------------
+ 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;
+(4 rows)
+
+-- Explain analyze, top tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                     query                                                     
+----------+-------+---------------------------------------------------------------------------------------------------------------
+ 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
+(3 rows)
+
+-- Create Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                             query                                              
+----------+-------+------------------------------------------------------------------------------------------------
+ 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)
+
+-- Create Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE MATERIALIZED VIEW pgss_materialized_view_2 AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                              query                                               
+----------+-------+--------------------------------------------------------------------------------------------------
+ t        |     1 | CREATE MATERIALIZED VIEW pgss_materialized_view_2 AS SELECT * FROM generate_series($1, $2) as id
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Refresh Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Refresh Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT $1
+ 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)
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_ctas_3 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_4 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_3 AS SELECT $1
+ t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_4 AS EXECUTE test_prepare_pgss
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Explain with Create Table As - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
+ 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)
+
+-- Explain with Create Table As - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
+ 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)
+
+-- Declare cursor, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                          query                          
+----------+-------+---------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ 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)
+
+-- Declare cursor, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                          query                          
+----------+-------+---------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ 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)
+
+-- COPY - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+COPY (SELECT 1) to stdout;
+1
+COPY (SELECT 1 UNION SELECT 2) to stdout;
+1
+2
+COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id
+    WHEN MATCHED THEN UPDATE SET x = id
+    WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout;
+1
+COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout;
+1
+COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout;
+2
+2
+COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout;
+2
+2
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                  query                                   
+----------+-------+--------------------------------------------------------------------------
+ t        |     1 | COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout
+ t        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout
+ t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id          +
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                 +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout
+ t        |     1 | COPY (SELECT 1 UNION SELECT 2) to stdout
+ t        |     1 | COPY (SELECT 1) to stdout
+ t        |     1 | COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | COPY (DELETE FROM stats_track_tab WHERE x=$1 RETURNING x) to stdout
+ f        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES ($1) RETURNING x) to stdout
+ f        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT $1 id) ON x = id         +
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                 +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout
+ f        |     1 | COPY (SELECT $1 UNION SELECT $2) to stdout
+ f        |     1 | COPY (SELECT $1) to stdout
+ f        |     1 | COPY (UPDATE stats_track_tab SET x=$1 WHERE x=$2 RETURNING x) to stdout
+(13 rows)
+
+-- COPY - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+COPY (SELECT 1) to stdout;
+1
+COPY (SELECT 1 UNION SELECT 2) to stdout;
+1
+2
+COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id
+    WHEN MATCHED THEN UPDATE SET x = id
+    WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout;
+1
+COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout;
+1
+COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout;
+2
+2
+COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout;
+2
+2
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                  query                                  
+----------+-------+-------------------------------------------------------------------------
+ t        |     1 | COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout
+ t        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout
+ t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id         +
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout
+ t        |     1 | COPY (SELECT 1 UNION SELECT 2) to stdout
+ t        |     1 | COPY (SELECT 1) to stdout
+ t        |     1 | COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(7 rows)
+
 -- 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 65a17147a5a..87aee416d1f 100644
--- a/contrib/pg_stat_statements/sql/level_tracking.sql
+++ b/contrib/pg_stat_statements/sql/level_tracking.sql
@@ -32,7 +32,6 @@ BEGIN
 END; $$;
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
-
 -- Procedure with multiple utility statements.
 CREATE OR REPLACE PROCEDURE proc_with_utility_stmt()
 LANGUAGE SQL
@@ -55,6 +54,234 @@ CALL proc_with_utility_stmt();
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
 
+-- Explain - all-level tracking.
+CREATE TABLE test_table (x int);
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1;
+explain (costs off) (SELECT 1, 2);
+explain (costs off) TABLE stats_track_tab;
+explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1);
+explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) DELETE FROM stats_track_tab;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((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 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1;
+explain (costs off) (SELECT 1, 2);
+explain (costs off) TABLE stats_track_tab;
+explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1);
+explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) DELETE FROM stats_track_tab;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((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 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain - all-level tracking with multi statement.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES (1), (2);
+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;
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain - top-level tracking with multi statement.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES ((1), (2));
+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;
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with CTE - all-level tracking
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) WITH a AS (select 4) SELECT 1;
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+explain (costs off) WITH a AS (select 4) 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) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with CTE - top-level tracking
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) WITH a AS (select 4) SELECT 1;
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+explain (costs off) WITH a AS (select 4) 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) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain analyze, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain analyze, top tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE MATERIALIZED VIEW pgss_materialized_view_2 AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Refresh Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Refresh Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_ctas_3 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_4 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with Create Table As - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with Create Table As - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Declare cursor, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Declare cursor, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- COPY - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+COPY (SELECT 1) to stdout;
+COPY (SELECT 1 UNION SELECT 2) to stdout;
+COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id
+    WHEN MATCHED THEN UPDATE SET x = id
+    WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout;
+COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout;
+COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout;
+COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- COPY - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+COPY (SELECT 1) to stdout;
+COPY (SELECT 1 UNION SELECT 2) to stdout;
+COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id
+    WHEN MATCHED THEN UPDATE SET x = id
+    WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout;
+COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout;
+COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout;
+COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
 -- DO block - top-level tracking without utility.
 SET pg_stat_statements.track = 'top';
 SET pg_stat_statements.track_utility = FALSE;
-- 
2.39.5 (Apple Git-154)

From 8718894eecfaf03dcce44f6dd3c90e0bd7294291 Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Wed, 16 Oct 2024 20:56:28 +0800
Subject: [PATCH v9 1/3] Better error reporting from extension scripts (Was:
 Extend ALTER OPERATOR)

---
 .../pg_stat_statements/expected/select.out    |  5 +-
 contrib/pg_stat_statements/sql/select.sql     |  3 +-
 src/backend/nodes/queryjumblefuncs.c          |  6 ++
 src/backend/parser/gram.y                     | 66 +++++++------------
 4 files changed, 34 insertions(+), 46 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/select.out b/contrib/pg_stat_statements/expected/select.out
index dd6c756f67..e0e2fa265c 100644
--- a/contrib/pg_stat_statements/expected/select.out
+++ b/contrib/pg_stat_statements/expected/select.out
@@ -19,8 +19,9 @@ SELECT 1 AS "int";
    1
 (1 row)
 
+/* this comment should not appear in the output */
 SELECT 'hello'
-  -- multiline
+  -- but this one will appear
   AS "text";
  text  
 -------
@@ -129,7 +130,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
 -------+------+------------------------------------------------------------------------------
      1 |    1 | PREPARE pgss_test (int) AS SELECT $1, $2 LIMIT $3
      4 |    4 | SELECT $1                                                                   +
-       |      |   -- multiline                                                              +
+       |      |   -- but this one will appear                                               +
        |      |   AS "text"
      2 |    2 | SELECT $1 + $2
      3 |    3 | SELECT $1 + $2 + $3 AS "add"
diff --git a/contrib/pg_stat_statements/sql/select.sql b/contrib/pg_stat_statements/sql/select.sql
index eb45cb81ad..e0be58d5e2 100644
--- a/contrib/pg_stat_statements/sql/select.sql
+++ b/contrib/pg_stat_statements/sql/select.sql
@@ -12,8 +12,9 @@ SELECT pg_stat_statements_reset() IS NOT NULL AS t;
 --
 SELECT 1 AS "int";
 
+/* this comment should not appear in the output */
 SELECT 'hello'
-  -- multiline
+  -- but this one will appear
   AS "text";
 
 SELECT 'world' AS "text";
diff --git a/src/backend/nodes/queryjumblefuncs.c b/src/backend/nodes/queryjumblefuncs.c
index 5e43fd9229..e8bf95690b 100644
--- a/src/backend/nodes/queryjumblefuncs.c
+++ b/src/backend/nodes/queryjumblefuncs.c
@@ -90,6 +90,12 @@ CleanQuerytext(const char *query, int *location, int *len)
 	/*
 	 * Discard leading and trailing whitespace, too.  Use scanner_isspace()
 	 * not libc's isspace(), because we want to match the lexer's behavior.
+	 *
+	 * Note: the parser now strips leading comments and whitespace from the
+	 * reported stmt_location, so this first loop will only iterate in the
+	 * unusual case that the location didn't propagate to here.  But the
+	 * statement length will extend to the end-of-string or terminating
+	 * semicolon, so the second loop often does something useful.
 	 */
 	while (query_len > 0 && scanner_isspace(query[0]))
 		query++, query_location++, query_len--;
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 4aa8646af7..4bab2117d9 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -67,39 +67,25 @@
 
 
 /*
- * Location tracking support --- simpler than bison's default, since we only
- * want to track the start position not the end position of each nonterminal.
+ * Location tracking support.  Unlike bison's default, we only want
+ * to track the start position not the end position of each nonterminal.
+ * Nonterminals that reduce to empty receive position "-1".  Since a
+ * production's leading RHS nonterminal(s) may have reduced to empty,
+ * we have to scan to find the first one that's not -1.
  */
 #define YYLLOC_DEFAULT(Current, Rhs, N) \
 	do { \
-		if ((N) > 0) \
-			(Current) = (Rhs)[1]; \
-		else \
-			(Current) = (-1); \
+		(Current) = (-1); \
+		for (int _i = 1; _i <= (N); _i++) \
+		{ \
+			if ((Rhs)[_i] >= 0) \
+			{ \
+				(Current) = (Rhs)[_i]; \
+				break; \
+			} \
+		} \
 	} while (0)
 
-/*
- * The above macro assigns -1 (unknown) as the parse location of any
- * nonterminal that was reduced from an empty rule, or whose leftmost
- * component was reduced from an empty rule.  This is problematic
- * for nonterminals defined like
- *		OptFooList: / * EMPTY * / { ... } | OptFooList Foo { ... } ;
- * because we'll set -1 as the location during the first reduction and then
- * copy it during each subsequent reduction, leaving us with -1 for the
- * location even when the list is not empty.  To fix that, do this in the
- * action for the nonempty rule(s):
- *		if (@$ < 0) @$ = @2;
- * (Although we have many nonterminals that follow this pattern, we only
- * bother with fixing @$ like this when the nonterminal's parse location
- * is actually referenced in some rule.)
- *
- * A cleaner answer would be to make YYLLOC_DEFAULT scan all the Rhs
- * locations until it's found one that's not -1.  Then we'd get a correct
- * location for any nonterminal that isn't entirely empty.  But this way
- * would add overhead to every rule reduction, and so far there's not been
- * a compelling reason to pay that overhead.
- */
-
 /*
  * Bison doesn't allocate anything that needs to live across parser calls,
  * so we can easily have it use palloc instead of malloc.  This prevents
@@ -930,7 +916,7 @@ parse_toplevel:
 			| MODE_PLPGSQL_EXPR PLpgSQL_Expr
 			{
 				pg_yyget_extra(yyscanner)->parsetree =
-					list_make1(makeRawStmt($2, 0));
+					list_make1(makeRawStmt($2, @2));
 			}
 			| MODE_PLPGSQL_ASSIGN1 PLAssignStmt
 			{
@@ -938,7 +924,7 @@ parse_toplevel:
 
 				n->nnames = 1;
 				pg_yyget_extra(yyscanner)->parsetree =
-					list_make1(makeRawStmt((Node *) n, 0));
+					list_make1(makeRawStmt((Node *) n, @2));
 			}
 			| MODE_PLPGSQL_ASSIGN2 PLAssignStmt
 			{
@@ -946,7 +932,7 @@ parse_toplevel:
 
 				n->nnames = 2;
 				pg_yyget_extra(yyscanner)->parsetree =
-					list_make1(makeRawStmt((Node *) n, 0));
+					list_make1(makeRawStmt((Node *) n, @2));
 			}
 			| MODE_PLPGSQL_ASSIGN3 PLAssignStmt
 			{
@@ -954,19 +940,15 @@ parse_toplevel:
 
 				n->nnames = 3;
 				pg_yyget_extra(yyscanner)->parsetree =
-					list_make1(makeRawStmt((Node *) n, 0));
+					list_make1(makeRawStmt((Node *) n, @2));
 			}
 		;
 
 /*
  * At top level, we wrap each stmt with a RawStmt node carrying start location
- * and length of the stmt's text.  Notice that the start loc/len are driven
- * entirely from semicolon locations (@2).  It would seem natural to use
- * @1 or @3 to get the true start location of a stmt, but that doesn't work
- * for statements that can start with empty nonterminals (opt_with_clause is
- * the main offender here); as noted in the comments for YYLLOC_DEFAULT,
- * we'd get -1 for the location in such cases.
- * We also take care to discard empty statements entirely.
+ * and length of the stmt's text.
+ * We also take care to discard empty statements entirely (which among other
+ * things dodges the problem of assigning them a location).
  */
 stmtmulti:	stmtmulti ';' toplevel_stmt
 				{
@@ -976,14 +958,14 @@ stmtmulti:	stmtmulti ';' toplevel_stmt
 						updateRawStmtEnd(llast_node(RawStmt, $1), @2);
 					}
 					if ($3 != NULL)
-						$$ = lappend($1, makeRawStmt($3, @2 + 1));
+						$$ = lappend($1, makeRawStmt($3, @3));
 					else
 						$$ = $1;
 				}
 			| toplevel_stmt
 				{
 					if ($1 != NULL)
-						$$ = list_make1(makeRawStmt($1, 0));
+						$$ = list_make1(makeRawStmt($1, @1));
 					else
 						$$ = NIL;
 				}
@@ -1584,8 +1566,6 @@ CreateSchemaStmt:
 OptSchemaEltList:
 			OptSchemaEltList schema_stmt
 				{
-					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
-						@$ = @2;
 					$$ = lappend($1, $2);
 				}
 			| /* EMPTY */
-- 
2.34.1

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

From a5ed8307b25c932f9ef14ffc16dd893e00a2b58b Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Mon, 7 Oct 2024 10:45:49 +0200
Subject: [PATCH v9 2/3] Add tests covering pgss nested queries

What pgss reports for nested statements can be confusing. Some
statements like CreateTableAs, DeclareCursor and CreateMaterializedView
don't jumble the nested query and thus won't report it in pgss. Explain
explicitely Jumble the nested query and will appear in pgss. However,
the reported query string will be the same as the top level statement.

Multi statements queries are also not handled correctly. The statement
length used in the top RawStmt is not propagated in the nested
statement. Thus, it fallbacks to the "use the whole query string"
default which is reported to pgss.

This patch creates tests to cover the current behaviour.
---
 .../expected/level_tracking.out               | 1100 +++++++++++++++++
 .../pg_stat_statements/sql/level_tracking.sql |  229 +++-
 2 files changed, 1328 insertions(+), 1 deletion(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index bb65e98ce0..8c5c1b5b55 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -112,6 +112,1106 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (2 rows)
 
+-- Explain - all-level tracking.
+CREATE TABLE test_table (x int);
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab;
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+explain (costs off) (TABLE test_table);
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                   query                                                    
+----------+-------+------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (SELECT $1, $2)
+ t        |     1 | explain (costs off) (TABLE test_table)
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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                                                                   +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1 UNION SELECT $2
+ 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+
+          |       |     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);
+(23 rows)
+
+-- Explain - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab;
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+explain (costs off) (TABLE test_table);
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                   query                                                    
+----------+-------+------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (SELECT $1, $2)
+ t        |     1 | explain (costs off) (TABLE test_table)
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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                                                                   +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1 UNION SELECT $2
+ 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)
+(12 rows)
+
+-- Explain - all-level tracking with multi statement.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES (1), (2);
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+           QUERY PLAN            
+---------------------------------
+ Insert on stats_track_tab
+   ->  Values Scan on "*VALUES*"
+(2 rows)
+
+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;
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+         QUERY PLAN         
+----------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+           QUERY PLAN            
+---------------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2), (3)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ 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
+ t        |     1 | explain (costs off) (SELECT $1, $2, $3, $4)
+ t        |     1 | explain (costs off) (TABLE test_table)
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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                                                                                         +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1, $2
+ t        |     1 | explain (costs off) SELECT $1, $2 UNION SELECT $3, $4
+ t        |     1 | explain (costs off) SELECT $1, $2, $3, $4, $5
+ t        |     1 | explain (costs off) TABLE stats_track_tab
+ 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));
+(37 rows)
+
+-- Explain - top-level tracking with multi statement.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES ((1), (2));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+ERROR:  INSERT has more expressions than target columns
+LINE 1: ...n (costs off) INSERT INTO stats_track_tab VALUES ((1), (2));
+                                                                   ^
+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;
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+         QUERY PLAN         
+----------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+           QUERY PLAN            
+---------------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2), (3)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ 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
+ t        |     1 | explain (costs off) (SELECT $1, $2, $3, $4)
+ t        |     1 | explain (costs off) (TABLE test_table)
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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))
+ 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
+ t        |     1 | explain (costs off) SELECT $1, $2 UNION SELECT $3, $4
+ t        |     1 | explain (costs off) SELECT $1, $2, $3, $4, $5
+ t        |     1 | explain (costs off) TABLE stats_track_tab
+ 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)
+(18 rows)
+
+-- Explain with CTE - all-level tracking
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) WITH a AS (select 4) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) 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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                              query                                                               
+----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (WITH a AS (select $1) (SELECT $2, $3))
+ t        |     1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab
+ t        |     1 | explain (costs off) WITH a AS (select $1) INSERT INTO stats_track_tab VALUES (($2))
+ t        |     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+
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                                         +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ 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+
+          |       |     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;
+(15 rows)
+
+-- Explain with CTE - top-level tracking
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) WITH a AS (select 4) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) 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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                              query                                                               
+----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (WITH a AS (select $1) (SELECT $2, $3))
+ t        |     1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab
+ t        |     1 | explain (costs off) WITH a AS (select $1) INSERT INTO stats_track_tab VALUES (($2))
+ t        |     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+
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                                         +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ 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
+(8 rows)
+
+-- Explain analyze, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                     query                                                     
+----------+-------+---------------------------------------------------------------------------------------------------------------
+ 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;
+(4 rows)
+
+-- Explain analyze, top tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                                     query                                                     
+----------+-------+---------------------------------------------------------------------------------------------------------------
+ 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
+(3 rows)
+
+-- Create Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                             query                                              
+----------+-------+------------------------------------------------------------------------------------------------
+ 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)
+
+-- Create Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE MATERIALIZED VIEW pgss_materialized_view_2 AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                              query                                               
+----------+-------+--------------------------------------------------------------------------------------------------
+ t        |     1 | CREATE MATERIALIZED VIEW pgss_materialized_view_2 AS SELECT * FROM generate_series($1, $2) as id
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Refresh Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Refresh Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT $1
+ 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)
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_ctas_3 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_4 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_3 AS SELECT $1
+ t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_4 AS EXECUTE test_prepare_pgss
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Explain with Create Table As - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
+ 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)
+
+-- Explain with Create Table As - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
+ 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)
+
+-- Declare cursor, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                          query                          
+----------+-------+---------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ 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)
+
+-- Declare cursor, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                          query                          
+----------+-------+---------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ 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)
+
+-- COPY - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+COPY (SELECT 1) to stdout;
+1
+COPY (SELECT 1 UNION SELECT 2) to stdout;
+1
+2
+COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id
+    WHEN MATCHED THEN UPDATE SET x = id
+    WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout;
+1
+COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout;
+1
+COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout;
+2
+2
+COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout;
+2
+2
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                  query                                   
+----------+-------+--------------------------------------------------------------------------
+ t        |     1 | COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout
+ t        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout
+ t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id          +
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                 +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout
+ t        |     1 | COPY (SELECT 1 UNION SELECT 2) to stdout
+ t        |     1 | COPY (SELECT 1) to stdout
+ t        |     1 | COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | COPY (DELETE FROM stats_track_tab WHERE x=$1 RETURNING x) to stdout
+ f        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES ($1) RETURNING x) to stdout
+ f        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT $1 id) ON x = id         +
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                 +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout
+ f        |     1 | COPY (SELECT $1 UNION SELECT $2) to stdout
+ f        |     1 | COPY (SELECT $1) to stdout
+ f        |     1 | COPY (UPDATE stats_track_tab SET x=$1 WHERE x=$2 RETURNING x) to stdout
+(13 rows)
+
+-- COPY - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+COPY (SELECT 1) to stdout;
+1
+COPY (SELECT 1 UNION SELECT 2) to stdout;
+1
+2
+COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id
+    WHEN MATCHED THEN UPDATE SET x = id
+    WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout;
+1
+COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout;
+1
+COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout;
+2
+2
+COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout;
+2
+2
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+ toplevel | calls |                                  query                                  
+----------+-------+-------------------------------------------------------------------------
+ t        |     1 | COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout
+ t        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout
+ t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id         +
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout
+ t        |     1 | COPY (SELECT 1 UNION SELECT 2) to stdout
+ t        |     1 | COPY (SELECT 1) to stdout
+ t        |     1 | COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(7 rows)
+
 -- 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 65a17147a5..87aee416d1 100644
--- a/contrib/pg_stat_statements/sql/level_tracking.sql
+++ b/contrib/pg_stat_statements/sql/level_tracking.sql
@@ -32,7 +32,6 @@ BEGIN
 END; $$;
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
-
 -- Procedure with multiple utility statements.
 CREATE OR REPLACE PROCEDURE proc_with_utility_stmt()
 LANGUAGE SQL
@@ -55,6 +54,234 @@ CALL proc_with_utility_stmt();
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
 
+-- Explain - all-level tracking.
+CREATE TABLE test_table (x int);
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1;
+explain (costs off) (SELECT 1, 2);
+explain (costs off) TABLE stats_track_tab;
+explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1);
+explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) DELETE FROM stats_track_tab;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((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 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1;
+explain (costs off) (SELECT 1, 2);
+explain (costs off) TABLE stats_track_tab;
+explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1);
+explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) DELETE FROM stats_track_tab;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((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 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain - all-level tracking with multi statement.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES (1), (2);
+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;
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain - top-level tracking with multi statement.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES ((1), (2));
+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;
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with CTE - all-level tracking
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) WITH a AS (select 4) SELECT 1;
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+explain (costs off) WITH a AS (select 4) 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) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with CTE - top-level tracking
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) WITH a AS (select 4) SELECT 1;
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+explain (costs off) WITH a AS (select 4) 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) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain analyze, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain analyze, top tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE MATERIALIZED VIEW pgss_materialized_view_2 AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Refresh Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Refresh Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_ctas_3 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_4 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with Create Table As - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Explain with Create Table As - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Declare cursor, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- Declare cursor, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- COPY - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+COPY (SELECT 1) to stdout;
+COPY (SELECT 1 UNION SELECT 2) to stdout;
+COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id
+    WHEN MATCHED THEN UPDATE SET x = id
+    WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout;
+COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout;
+COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout;
+COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
+-- COPY - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+COPY (SELECT 1) to stdout;
+COPY (SELECT 1 UNION SELECT 2) to stdout;
+COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id
+    WHEN MATCHED THEN UPDATE SET x = id
+    WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout;
+COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout;
+COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout;
+COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY toplevel desc, query COLLATE "C";
+
 -- DO block - top-level tracking without utility.
 SET pg_stat_statements.track = 'top';
 SET pg_stat_statements.track_utility = FALSE;
-- 
2.34.1

From 361e3df50ca2f649b44fc0802a3bcea6503a92ab Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
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 dc46c08a67d..d1638a23159 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -880,8 +880,9 @@ 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
  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';
@@ -926,8 +927,9 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  toplevel | calls |                                             query                                              
 ----------+-------+------------------------------------------------------------------------------------------------
  t        |     1 | CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series($1, $2) as id
+ f        |     1 | SELECT * FROM generate_series($1, $2) as id
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(2 rows)
+(3 rows)
 
 -- Create Materialized View, top-level tracking.
 SET pg_stat_statements.track = 'top';
@@ -960,8 +962,9 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  toplevel | calls |                       query                        
 ----------+-------+----------------------------------------------------
  t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
+ f        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view;
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(2 rows)
+(3 rows)
 
 -- Refresh Materialized View, top-level tracking.
 SET pg_stat_statements.track = 'top';
@@ -997,9 +1000,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';
@@ -1038,9 +1042,10 @@ SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
  toplevel | calls |                                   query                                   
 ----------+-------+---------------------------------------------------------------------------
+ f        |     1 | SELECT $1
  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)
+(3 rows)
 
 -- Explain with Create Table As - top-level tracking.
 SET pg_stat_statements.track = 'top';
@@ -1090,8 +1095,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';
diff --git a/src/backend/commands/createas.c b/src/backend/commands/createas.c
index 68ec122dbf9..7ed42aa4419 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 18a5af6b919..c81221cdbeb 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 dcbdaa03885..d2eef8097cf 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 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)

From b1fd24689ae2af217cc8ef89f3c2ce7edfc061bf Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Thu, 3 Oct 2024 08:52:02 +0200
Subject: Track location to extract relevant part in nested statement

Previously, Query generated through transform would have unset
stmt_location. Extensions relying on the statement location to extract
the relevant part of the statement would fallback to use the whole
statement instead, thus showing the same string in the top and
nested level which was a source of confusion.

This patch fixes the issue by keeping track of the statement locations
and propagate it to Query during transform, allowing pgss to only show
the relevant part of the query for nested query.
---
 .../expected/level_tracking.out               | 122 +++++++++---------
 .../pg_stat_statements/expected/planning.out  |  10 +-
 .../pg_stat_statements/expected/select.out    |   2 +-
 .../pg_stat_statements/expected/utility.out   |   2 +-
 contrib/pg_stat_statements/sql/planning.sql   |   4 +-
 src/backend/optimizer/util/clauses.c          |   2 +-
 src/backend/parser/analyze.c                  |  96 ++++++++++----
 src/backend/parser/gram.y                     |  71 +++++++++-
 src/backend/parser/parse_merge.c              |   2 +
 src/include/nodes/parsenodes.h                |  10 ++
 src/include/parser/analyze.h                  |   3 +-
 src/include/parser/parse_node.h               |   2 +
 12 files changed, 219 insertions(+), 107 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index d926082ac10..dc46c08a67d 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -207,33 +207,33 @@ SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
  toplevel | calls |                                                   query                                                    
 ----------+-------+------------------------------------------------------------------------------------------------------------
+ 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 | SELECT $1
+ f        |     1 | SELECT $1 UNION SELECT $2
+ f        |     1 | SELECT $1, $2
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | TABLE test_table
+ f        |     1 | UPDATE stats_track_tab SET x=$1 WHERE x=$2
+ f        |     1 | VALUES($1)
+ f        |     1 | VALUES($1, $2)
  t        |     1 | explain (costs off) (SELECT $1, $2)
- f        |     1 | explain (costs off) (SELECT $1, $2);
  t        |     1 | explain (costs off) (TABLE test_table)
- f        |     1 | explain (costs off) (TABLE test_table);
  t        |     1 | explain (costs off) (VALUES($1, $2))
- f        |     1 | explain (costs off) (VALUES($1, $2));
  t        |     1 | explain (costs off) DELETE FROM stats_track_tab
- f        |     1 | explain (costs off) DELETE FROM stats_track_tab;
  t        |     1 | explain (costs off) INSERT INTO stats_track_tab VALUES (($1))
- f        |     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                                                                   +
           |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
- 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);
  t        |     1 | explain (costs off) SELECT $1
  t        |     1 | explain (costs off) SELECT $1 UNION SELECT $2
- f        |     1 | explain (costs off) SELECT $1 UNION SELECT $2;
- f        |     1 | explain (costs off) SELECT $1;
  t        |     1 | explain (costs off) TABLE stats_track_tab
- f        |     1 | explain (costs off) TABLE stats_track_tab;
  t        |     1 | explain (costs off) UPDATE stats_track_tab SET x=$1 WHERE x=$2
- f        |     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) VALUES($1);
 (23 rows)
 
 -- Explain - top-level tracking.
@@ -480,51 +480,49 @@ explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                                              query                                                               
-----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ toplevel | calls |                                                   query                                                    
+----------+-------+------------------------------------------------------------------------------------------------------------
+ f        |     1 | (SELECT $1, $2, $3) UNION SELECT $4, $5, $6
+ 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
+ f        |     1 | SELECT $1, $2, $3, $4
+ f        |     1 | SELECT $1, $2, $3, $4, $5
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | TABLE test_table
+ 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)
+ f        |     1 | VALUES($1, $2)
  t        |     1 | explain (costs off) (SELECT $1, $2, $3)
  t        |     1 | explain (costs off) (SELECT $1, $2, $3) UNION SELECT $4, $5, $6
- f        |     1 | explain (costs off) (SELECT $1, $2, $3); explain (costs off) (SELECT 1, 2, 3, 4);
  t        |     1 | 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);
  t        |     1 | explain (costs off) (TABLE test_table)
  t        |     1 | explain (costs off) (VALUES($1, $2))
  t        |     1 | explain (costs off) DELETE FROM stats_track_tab
  t        |     1 | 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) DELETE FROM stats_track_tab; 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))
- 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);
- 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)
- 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;
  t        |     1 | explain (costs off) SELECT $1
  t        |     1 | explain (costs off) SELECT $1, $2
  t        |     1 | explain (costs off) SELECT $1, $2 UNION SELECT $3, $4
- f        |     1 | explain (costs off) SELECT $1, $2 UNION SELECT $3, $4; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
  t        |     1 | explain (costs off) SELECT $1, $2, $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;
  t        |     1 | explain (costs off) TABLE stats_track_tab
- 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);
  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
- 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;
  t        |     1 | explain (costs off) VALUES($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));
 (37 rows)
 
 -- Explain - top-level tracking with multi statement.
@@ -753,24 +751,24 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  toplevel | calls |                                                              query                                                               
 ----------+-------+----------------------------------------------------------------------------------------------------------------------------------
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ 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 | 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
  t        |     1 | explain (costs off) (WITH a AS (select $1) (SELECT $2, $3))
- f        |     1 | explain (costs off) (WITH a AS (select $1) (SELECT $2, $3));
  t        |     1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab
- f        |     1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab;
  t        |     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) INSERT INTO stats_track_tab VALUES (($2));
  t        |     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+
           |       |     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) 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);
  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
- 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;
  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) UPDATE stats_track_tab SET x=$2 WHERE x=$3;
 (15 rows)
 
 -- Explain with CTE - top-level tracking
@@ -881,7 +879,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
- f        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1;
+ f        |     1 | SELECT $1
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (4 rows)
 
@@ -999,8 +997,8 @@ 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 | PREPARE test_prepare_pgss AS select generate_series($1, $2)
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | select generate_series($1, $2)
 (4 rows)
 
 -- Create Table As, top-level tracking.
@@ -1151,25 +1149,25 @@ COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout;
 2
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                  query                                   
-----------+-------+--------------------------------------------------------------------------
- f        |     1 | COPY (DELETE FROM stats_track_tab WHERE x=$1 RETURNING x) to stdout
+ toplevel | calls |                                  query                                  
+----------+-------+-------------------------------------------------------------------------
  t        |     1 | COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout
- f        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES ($1) RETURNING x) to stdout
  t        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout
- f        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT $1 id) ON x = id         +
-          |       |     WHEN MATCHED THEN UPDATE SET x = id                                 +
+ t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id         +
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                +
           |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout
- t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id          +
-          |       |     WHEN MATCHED THEN UPDATE SET x = id                                 +
-          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout
- f        |     1 | COPY (SELECT $1 UNION SELECT $2) to stdout
- f        |     1 | COPY (SELECT $1) to stdout
  t        |     1 | COPY (SELECT 1 UNION SELECT 2) to stdout
  t        |     1 | COPY (SELECT 1) to stdout
- f        |     1 | COPY (UPDATE stats_track_tab SET x=$1 WHERE x=$2 RETURNING x) to stdout
  t        |     1 | COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout
+ f        |     1 | DELETE FROM stats_track_tab WHERE x=$1 RETURNING x
+ f        |     1 | INSERT INTO stats_track_tab (x) VALUES ($1) RETURNING x
+ f        |     1 | MERGE INTO stats_track_tab USING (SELECT $1 id) ON x = id              +
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x
+ f        |     1 | SELECT $1
+ f        |     1 | SELECT $1 UNION SELECT $2
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | UPDATE stats_track_tab SET x=$1 WHERE x=$2 RETURNING x
 (13 rows)
 
 -- COPY - top-level tracking.
diff --git a/contrib/pg_stat_statements/expected/planning.out b/contrib/pg_stat_statements/expected/planning.out
index 9effd11fdc8..3ee1928cbe9 100644
--- a/contrib/pg_stat_statements/expected/planning.out
+++ b/contrib/pg_stat_statements/expected/planning.out
@@ -58,7 +58,7 @@ SELECT 42;
 (1 row)
 
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
  plans | calls | rows |                          query                           
 -------+-------+------+----------------------------------------------------------
      0 |     1 |    0 | ALTER TABLE stats_plan_test ADD COLUMN x int
@@ -72,10 +72,10 @@ SELECT plans, calls, rows, query FROM pg_stat_statements
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
- plans_ok | calls | rows |                         query                         
-----------+-------+------+-------------------------------------------------------
- t        |     4 |    4 | PREPARE prep1 AS SELECT COUNT(*) FROM stats_plan_test
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
+ plans_ok | calls | rows |                query                 
+----------+-------+------+--------------------------------------
+ t        |     4 |    4 | SELECT COUNT(*) FROM stats_plan_test
 (1 row)
 
 -- Cleanup
diff --git a/contrib/pg_stat_statements/expected/select.out b/contrib/pg_stat_statements/expected/select.out
index dd6c756f67d..217a2c0b2bc 100644
--- a/contrib/pg_stat_statements/expected/select.out
+++ b/contrib/pg_stat_statements/expected/select.out
@@ -127,7 +127,6 @@ DEALLOCATE pgss_test;
 SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
  calls | rows |                                    query                                     
 -------+------+------------------------------------------------------------------------------
-     1 |    1 | PREPARE pgss_test (int) AS SELECT $1, $2 LIMIT $3
      4 |    4 | SELECT $1                                                                   +
        |      |   -- multiline                                                              +
        |      |   AS "text"
@@ -137,6 +136,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
      2 |    2 | SELECT $1 AS "int"
      1 |    2 | SELECT $1 AS i UNION SELECT $2 ORDER BY i
      1 |    1 | SELECT $1 || $2
+     1 |    1 | SELECT $1, $2 LIMIT $3
      0 |    0 | SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C"
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
      1 |    2 | WITH t(f) AS (                                                              +
diff --git a/contrib/pg_stat_statements/expected/utility.out b/contrib/pg_stat_statements/expected/utility.out
index 060d4416dd7..aa4f0f7e628 100644
--- a/contrib/pg_stat_statements/expected/utility.out
+++ b/contrib/pg_stat_statements/expected/utility.out
@@ -540,7 +540,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
 -------+------+----------------------------------------------------
      2 |    0 | DEALLOCATE $1
      2 |    0 | DEALLOCATE ALL
-     2 |    2 | PREPARE stat_select AS SELECT $1 AS a
+     2 |    2 | SELECT $1 AS a
      1 |    1 | SELECT $1 as a
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (5 rows)
diff --git a/contrib/pg_stat_statements/sql/planning.sql b/contrib/pg_stat_statements/sql/planning.sql
index 46f5d9b951c..9cfe206b3b0 100644
--- a/contrib/pg_stat_statements/sql/planning.sql
+++ b/contrib/pg_stat_statements/sql/planning.sql
@@ -20,11 +20,11 @@ SELECT 42;
 SELECT 42;
 SELECT 42;
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 
 -- Cleanup
 DROP TABLE stats_plan_test;
diff --git a/src/backend/optimizer/util/clauses.c b/src/backend/optimizer/util/clauses.c
index 8e39795e245..2bd14335bdb 100644
--- a/src/backend/optimizer/util/clauses.c
+++ b/src/backend/optimizer/util/clauses.c
@@ -4686,7 +4686,7 @@ inline_function(Oid funcid, Oid result_type, Oid result_collid,
 		pstate->p_sourcetext = src;
 		sql_fn_parser_setup(pstate, pinfo);
 
-		querytree = transformTopLevelStmt(pstate, linitial(raw_parsetree_list));
+		querytree = transformOptionalSelectInto(pstate, ((RawStmt *) linitial(raw_parsetree_list))->stmt);
 
 		free_parsestate(pstate);
 	}
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index 8a6ba1692e8..5f4372c4a00 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -58,7 +58,6 @@
 /* Hook for plugins to get control at end of parse analysis */
 post_parse_analyze_hook_type post_parse_analyze_hook = NULL;
 
-static Query *transformOptionalSelectInto(ParseState *pstate, Node *parseTree);
 static Query *transformDeleteStmt(ParseState *pstate, DeleteStmt *stmt);
 static Query *transformInsertStmt(ParseState *pstate, InsertStmt *stmt);
 static OnConflictExpr *transformOnConflictClause(ParseState *pstate,
@@ -113,13 +112,15 @@ 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);
 
 	pstate->p_queryEnv = queryEnv;
 
-	query = transformTopLevelStmt(pstate, parseTree);
+	query = transformOptionalSelectInto(pstate, parseTree->stmt);
 
 	if (IsQueryIdEnabled())
 		jstate = JumbleQuery(query);
@@ -153,12 +154,14 @@ 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);
 
 	pstate->p_queryEnv = queryEnv;
 
-	query = transformTopLevelStmt(pstate, parseTree);
+	query = transformOptionalSelectInto(pstate, parseTree->stmt);
 
 	/* make sure all is well with parameter types */
 	check_variable_parameters(pstate, query);
@@ -195,10 +198,12 @@ 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);
 
-	query = transformTopLevelStmt(pstate, parseTree);
+	query = transformOptionalSelectInto(pstate, parseTree->stmt);
 
 	if (IsQueryIdEnabled())
 		jstate = JumbleQuery(query);
@@ -238,27 +243,6 @@ parse_sub_analyze(Node *parseTree, ParseState *parentParseState,
 	return query;
 }
 
-/*
- * transformTopLevelStmt -
- *	  transform a Parse tree into a Query tree.
- *
- * This function is just responsible for transferring statement location data
- * from the RawStmt into the finished Query.
- */
-Query *
-transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree)
-{
-	Query	   *result;
-
-	/* We're at top level, so allow SELECT INTO */
-	result = transformOptionalSelectInto(pstate, parseTree->stmt);
-
-	result->stmt_location = parseTree->stmt_location;
-	result->stmt_len = parseTree->stmt_len;
-
-	return result;
-}
-
 /*
  * transformOptionalSelectInto -
  *	  If SELECT has INTO, convert it to CREATE TABLE AS.
@@ -269,7 +253,7 @@ transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree)
  * of the parse tree, and so we only try it before entering the recursive
  * transformStmt() processing.
  */
-static Query *
+Query *
 transformOptionalSelectInto(ParseState *pstate, Node *parseTree)
 {
 	if (IsA(parseTree, SelectStmt))
@@ -417,7 +401,9 @@ transformStmt(ParseState *pstate, Node *parseTree)
 			 */
 			result = makeNode(Query);
 			result->commandType = CMD_UTILITY;
-			result->utilityStmt = (Node *) parseTree;
+			result->utilityStmt = parseTree;
+			result->stmt_location = pstate->p_stmt_location;
+			result->stmt_len = pstate->p_stmt_len;
 			break;
 	}
 
@@ -506,6 +492,37 @@ analyze_requires_snapshot(RawStmt *parseTree)
 	return stmt_requires_parse_analysis(parseTree);
 }
 
+/*
+ * setQueryStmtLen
+ *		Set stmt_len in Query.
+ *
+ * Some statements, like PreparableStmt, can be located within parentheses.
+ * For example "(SELECT 1)" or "COPY (UPDATE ...) to x;". For those, we can't
+ * use the whole string from the statement's location or the SQL string will
+ * yield "SELECT 1)". The parser will set stmt_len, reflecting the size of the
+ * statement within the parentheses. Thus, when stmt_len is available, we use it
+ * for the Query's stmt_len.
+ *
+ * For other cases, the parser can't provide the length of individual statements.
+ * However, we have the statement's location plus the length (p_stmt_len) and
+ * location (p_stmt_location) of the top level RawStmt, stored in pstate. Thus,
+ * the statement's length is the RawStmt's length minus how much we've advanced
+ * in the RawStmt's string.
+ */
+void
+setQueryStmtLen(ParseState *pstate, Query *qry, int stmt_len)
+{
+	if (stmt_len > 0)
+		/* Statement's length is known, use it */
+		qry->stmt_len = stmt_len;
+	else
+		/*
+		 * Compute the statement's length from statement's location and RawStmt's
+		 * length and location
+		 */
+		qry->stmt_len = pstate->p_stmt_len - (qry->stmt_location - pstate->p_stmt_location);
+}
+
 /*
  * transformDeleteStmt -
  *	  transforms a Delete Statement
@@ -518,6 +535,8 @@ transformDeleteStmt(ParseState *pstate, DeleteStmt *stmt)
 	Node	   *qual;
 
 	qry->commandType = CMD_DELETE;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 
 	/* process the WITH clause independently of all else */
 	if (stmt->withClause)
@@ -606,6 +625,8 @@ transformInsertStmt(ParseState *pstate, InsertStmt *stmt)
 	Assert(pstate->p_ctenamespace == NIL);
 
 	qry->commandType = CMD_INSERT;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 	pstate->p_is_insert = true;
 
 	/* process the WITH clause independently of all else */
@@ -1331,7 +1352,6 @@ count_rowexpr_columns(ParseState *pstate, Node *expr)
 	return -1;
 }
 
-
 /*
  * transformSelectStmt -
  *	  transforms a Select Statement
@@ -1347,6 +1367,8 @@ transformSelectStmt(ParseState *pstate, SelectStmt *stmt)
 	ListCell   *l;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 
 	/* process the WITH clause independently of all else */
 	if (stmt->withClause)
@@ -1499,6 +1521,8 @@ transformValuesClause(ParseState *pstate, SelectStmt *stmt)
 	int			i;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 
 	/* Most SELECT stuff doesn't apply in a VALUES clause */
 	Assert(stmt->distinctClause == NIL);
@@ -1730,6 +1754,8 @@ transformSetOperationStmt(ParseState *pstate, SelectStmt *stmt)
 	int			tllen;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 
 	/*
 	 * Find leftmost leaf SelectStmt.  We currently only need to do this in
@@ -2397,6 +2423,8 @@ transformReturnStmt(ParseState *pstate, ReturnStmt *stmt)
 	Query	   *qry = makeNode(Query);
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = pstate->p_stmt_location;
+	qry->stmt_len = pstate->p_stmt_len;
 	qry->isReturn = true;
 
 	qry->targetList = list_make1(makeTargetEntry((Expr *) transformExpr(pstate, stmt->returnval, EXPR_KIND_SELECT_TARGET),
@@ -2430,6 +2458,8 @@ transformUpdateStmt(ParseState *pstate, UpdateStmt *stmt)
 	Node	   *qual;
 
 	qry->commandType = CMD_UPDATE;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 	pstate->p_is_insert = false;
 
 	/* process the WITH clause independently of all else */
@@ -2677,6 +2707,8 @@ transformPLAssignStmt(ParseState *pstate, PLAssignStmt *stmt)
 	 * consider WITH or INTO, and we build a targetlist our own way.
 	 */
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = pstate->p_stmt_location;
+	qry->stmt_len = pstate->p_stmt_len;
 	pstate->p_is_insert = false;
 
 	/* make FOR UPDATE/FOR SHARE info available to addRangeTableEntry */
@@ -2948,6 +2980,8 @@ transformDeclareCursorStmt(ParseState *pstate, DeclareCursorStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3003,6 +3037,8 @@ transformExplainStmt(ParseState *pstate, ExplainStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3083,6 +3119,8 @@ transformCreateTableAsStmt(ParseState *pstate, CreateTableAsStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3207,6 +3245,8 @@ transformCallStmt(ParseState *pstate, CallStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 4aa8646af7b..d00f4b47c4a 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -167,6 +167,7 @@ static void base_yyerror(YYLTYPE *yylloc, core_yyscan_t yyscanner,
 						 const char *msg);
 static RawStmt *makeRawStmt(Node *stmt, int stmt_location);
 static void updateRawStmtEnd(RawStmt *rs, int end_location);
+static void updatePreparableStmtEnd(Node *n, int end_location);
 static Node *makeColumnRef(char *colname, List *indirection,
 						   int location, core_yyscan_t yyscanner);
 static Node *makeTypeCast(Node *arg, TypeName *typename, int location);
@@ -190,7 +191,7 @@ static void insertSelectOptions(SelectStmt *stmt,
 								SelectLimit *limitClause,
 								WithClause *withClause,
 								core_yyscan_t yyscanner);
-static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg);
+static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location);
 static Node *doNegate(Node *n, int location);
 static void doNegateFloat(Float *v);
 static Node *makeAndExpr(Node *lexpr, Node *rexpr, int location);
@@ -3403,6 +3404,7 @@ CopyStmt:	COPY opt_binary qualified_name opt_column_list
 				{
 					CopyStmt *n = makeNode(CopyStmt);
 
+					updatePreparableStmtEnd($3, @4);
 					n->relation = NULL;
 					n->query = $3;
 					n->attlist = NIL;
@@ -12170,6 +12172,9 @@ InsertStmt:
 					$5->onConflictClause = $6;
 					$5->returningList = $7;
 					$5->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					$5->stmt_location = @$;
 					$$ = (Node *) $5;
 				}
 		;
@@ -12323,6 +12328,9 @@ DeleteStmt: opt_with_clause DELETE_P FROM relation_expr_opt_alias
 					n->whereClause = $6;
 					n->returningList = $7;
 					n->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					n->stmt_location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12397,6 +12405,9 @@ UpdateStmt: opt_with_clause UPDATE relation_expr_opt_alias
 					n->whereClause = $7;
 					n->returningList = $8;
 					n->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					n->stmt_location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12474,6 +12485,9 @@ MergeStmt:
 					m->joinCondition = $8;
 					m->mergeWhenClauses = $9;
 					m->returningList = $10;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					m->stmt_location = @$;
 
 					$$ = (Node *) m;
 				}
@@ -12714,7 +12728,12 @@ SelectStmt: select_no_parens			%prec UMINUS
 		;
 
 select_with_parens:
-			'(' select_no_parens ')'				{ $$ = $2; }
+			'(' select_no_parens ')'
+				{
+					SelectStmt *n = (SelectStmt *) $2;
+					n->stmt_len = @3 - @2;
+					$$ = $2;
+				}
 			| '(' select_with_parens ')'			{ $$ = $2; }
 		;
 
@@ -12836,6 +12855,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| SELECT distinct_clause target_list
@@ -12853,6 +12873,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| values_clause							{ $$ = $1; }
@@ -12873,19 +12894,20 @@ simple_select:
 
 					n->targetList = list_make1(rt);
 					n->fromClause = list_make1($2);
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| select_clause UNION set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause INTERSECT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause EXCEPT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 		;
 
@@ -13443,6 +13465,7 @@ values_clause:
 				{
 					SelectStmt *n = makeNode(SelectStmt);
 
+					n->stmt_location = @1;
 					n->valuesLists = list_make1($3);
 					$$ = (Node *) n;
 				}
@@ -18585,6 +18608,39 @@ updateRawStmtEnd(RawStmt *rs, int end_location)
 	rs->stmt_len = end_location - rs->stmt_location;
 }
 
+/* Adjust a PreparableStmt to reflect that it doesn't run to the end of the string */
+static void
+updatePreparableStmtEnd(Node *n, int end_location)
+{
+	if (IsA(n, SelectStmt))
+	{
+		SelectStmt *stmt = (SelectStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, InsertStmt))
+	{
+		InsertStmt *stmt = (InsertStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, UpdateStmt))
+	{
+		UpdateStmt *stmt = (UpdateStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, DeleteStmt))
+	{
+		DeleteStmt *stmt = (DeleteStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, MergeStmt))
+	{
+		MergeStmt *stmt = (MergeStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else
+		elog(ERROR, "unexpected node type %d", (int) n->type);
+}
+
 static Node *
 makeColumnRef(char *colname, List *indirection,
 			  int location, core_yyscan_t yyscanner)
@@ -18963,11 +19019,13 @@ insertSelectOptions(SelectStmt *stmt,
 					 errmsg("multiple WITH clauses not allowed"),
 					 parser_errposition(exprLocation((Node *) withClause))));
 		stmt->withClause = withClause;
+		/* Update SelectStmt's location to the start of the with clause */
+		stmt->stmt_location = withClause->location;
 	}
 }
 
 static Node *
-makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
+makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location)
 {
 	SelectStmt *n = makeNode(SelectStmt);
 
@@ -18975,6 +19033,7 @@ makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
 	n->all = all;
 	n->larg = (SelectStmt *) larg;
 	n->rarg = (SelectStmt *) rarg;
+	n->stmt_location = location;
 	return (Node *) n;
 }
 
diff --git a/src/backend/parser/parse_merge.c b/src/backend/parser/parse_merge.c
index 87df79027d7..fdfa61868cf 100644
--- a/src/backend/parser/parse_merge.c
+++ b/src/backend/parser/parse_merge.c
@@ -118,6 +118,8 @@ transformMergeStmt(ParseState *pstate, MergeStmt *stmt)
 	Assert(pstate->p_ctenamespace == NIL);
 
 	qry->commandType = CMD_MERGE;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 	qry->hasRecursive = false;
 
 	/* process the WITH clause independently of all else */
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index c92cef3d16d..b40b661ec8a 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -2046,6 +2046,8 @@ typedef struct InsertStmt
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
 	OverridingKind override;	/* OVERRIDING clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } InsertStmt;
 
 /* ----------------------
@@ -2060,6 +2062,8 @@ typedef struct DeleteStmt
 	Node	   *whereClause;	/* qualifications */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } DeleteStmt;
 
 /* ----------------------
@@ -2075,6 +2079,8 @@ typedef struct UpdateStmt
 	List	   *fromClause;		/* optional from clause for more tables */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } UpdateStmt;
 
 /* ----------------------
@@ -2090,6 +2096,8 @@ typedef struct MergeStmt
 	List	   *mergeWhenClauses;	/* list of MergeWhenClause(es) */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } MergeStmt;
 
 /* ----------------------
@@ -2159,6 +2167,8 @@ typedef struct SelectStmt
 	bool		all;			/* ALL specified? */
 	struct SelectStmt *larg;	/* left child */
 	struct SelectStmt *rarg;	/* right child */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 	/* Eventually add fields for CORRESPONDING spec here */
 } SelectStmt;
 
diff --git a/src/include/parser/analyze.h b/src/include/parser/analyze.h
index 28b66fccb43..8ba4e050aff 100644
--- a/src/include/parser/analyze.h
+++ b/src/include/parser/analyze.h
@@ -46,8 +46,9 @@ extern List *transformUpdateTargetList(ParseState *pstate,
 									   List *origTlist);
 extern List *transformReturningList(ParseState *pstate, List *returningList,
 									ParseExprKind exprKind);
-extern Query *transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree);
+extern Query *transformOptionalSelectInto(ParseState *pstate, Node *parseTree);
 extern Query *transformStmt(ParseState *pstate, Node *parseTree);
+extern void setQueryStmtLen(ParseState *pstate, Query *qry, int stmt_len);
 
 extern bool stmt_requires_parse_analysis(RawStmt *parseTree);
 extern bool analyze_requires_snapshot(RawStmt *parseTree);
diff --git a/src/include/parser/parse_node.h b/src/include/parser/parse_node.h
index 543df568147..ba572b3aea2 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 */
-- 
2.39.5 (Apple Git-154)

From 98bb46aeda63ec84f96c9b7bd00ae615b46c7e01 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Mon, 7 Oct 2024 10:45:49 +0200
Subject: Add tests covering pgss nested queries

What pgss reports for nested statements can be confusing. Some
statements like CreateTableAs, DeclareCursor and CreateMaterializedView
don't jumble the nested query and thus won't report it in pgss. Explain
explicitely Jumble the nested query and will appear in pgss. However,
the reported query string will be the same as the top level statement.

Multi statements queries are also not handled correctly. The statement
length used in the top RawStmt is not propagated in the nested
statement. Thus, it fallbacks to the "use the whole query string"
default which is reported to pgss.

This patch creates tests to cover the current behaviour.
---
 .../expected/level_tracking.out               | 1100 +++++++++++++++++
 .../pg_stat_statements/sql/level_tracking.sql |  229 +++-
 2 files changed, 1328 insertions(+), 1 deletion(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index bb65e98ce09..d926082ac10 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -112,6 +112,1106 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (2 rows)
 
+-- Explain - all-level tracking.
+CREATE TABLE test_table (x int);
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab;
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+explain (costs off) (TABLE test_table);
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                                                   query                                                    
+----------+-------+------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (SELECT $1, $2)
+ f        |     1 | explain (costs off) (SELECT $1, $2);
+ t        |     1 | explain (costs off) (TABLE test_table)
+ f        |     1 | explain (costs off) (TABLE test_table);
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ f        |     1 | explain (costs off) (VALUES($1, $2));
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ f        |     1 | explain (costs off) DELETE FROM stats_track_tab;
+ t        |     1 | explain (costs off) INSERT INTO stats_track_tab VALUES (($1))
+ f        |     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                                                                   +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ 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);
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1 UNION SELECT $2
+ f        |     1 | explain (costs off) SELECT $1 UNION SELECT $2;
+ f        |     1 | explain (costs off) SELECT $1;
+ t        |     1 | explain (costs off) TABLE stats_track_tab
+ f        |     1 | explain (costs off) TABLE stats_track_tab;
+ t        |     1 | explain (costs off) UPDATE stats_track_tab SET x=$1 WHERE x=$2
+ f        |     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) VALUES($1);
+(23 rows)
+
+-- Explain - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab;
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+explain (costs off) (TABLE test_table);
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                                                   query                                                    
+----------+-------+------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (SELECT $1, $2)
+ t        |     1 | explain (costs off) (TABLE test_table)
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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                                                                   +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1 UNION SELECT $2
+ 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)
+(12 rows)
+
+-- Explain - all-level tracking with multi statement.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES (1), (2);
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+           QUERY PLAN            
+---------------------------------
+ Insert on stats_track_tab
+   ->  Values Scan on "*VALUES*"
+(2 rows)
+
+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;
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+         QUERY PLAN         
+----------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+           QUERY PLAN            
+---------------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2), (3)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ 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
+ f        |     1 | explain (costs off) (SELECT $1, $2, $3); explain (costs off) (SELECT 1, 2, 3, 4);
+ t        |     1 | 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);
+ t        |     1 | explain (costs off) (TABLE test_table)
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ t        |     1 | 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) DELETE FROM stats_track_tab; 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))
+ 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);
+ 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)
+ 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;
+ t        |     1 | explain (costs off) SELECT $1
+ t        |     1 | explain (costs off) SELECT $1, $2
+ t        |     1 | explain (costs off) SELECT $1, $2 UNION SELECT $3, $4
+ f        |     1 | explain (costs off) SELECT $1, $2 UNION SELECT $3, $4; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+ t        |     1 | explain (costs off) SELECT $1, $2, $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;
+ t        |     1 | explain (costs off) TABLE stats_track_tab
+ 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);
+ 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
+ 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;
+ t        |     1 | explain (costs off) VALUES($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));
+(37 rows)
+
+-- Explain - top-level tracking with multi statement.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+         QUERY PLAN          
+-----------------------------
+ Seq Scan on stats_track_tab
+(1 row)
+
+       QUERY PLAN       
+------------------------
+ Seq Scan on test_table
+(1 row)
+
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES ((1), (2));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+ERROR:  INSERT has more expressions than target columns
+LINE 1: ...n (costs off) INSERT INTO stats_track_tab VALUES ((1), (2));
+                                                                   ^
+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;
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+         QUERY PLAN         
+----------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+           QUERY PLAN            
+---------------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1), (2), (3)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ 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
+ t        |     1 | explain (costs off) (SELECT $1, $2, $3, $4)
+ t        |     1 | explain (costs off) (TABLE test_table)
+ t        |     1 | explain (costs off) (VALUES($1, $2))
+ t        |     1 | explain (costs off) DELETE FROM stats_track_tab
+ 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))
+ 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
+ t        |     1 | explain (costs off) SELECT $1, $2 UNION SELECT $3, $4
+ t        |     1 | explain (costs off) SELECT $1, $2, $3, $4, $5
+ t        |     1 | explain (costs off) TABLE stats_track_tab
+ 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)
+(18 rows)
+
+-- Explain with CTE - all-level tracking
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) WITH a AS (select 4) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) 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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                                                              query                                                               
+----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (WITH a AS (select $1) (SELECT $2, $3))
+ f        |     1 | explain (costs off) (WITH a AS (select $1) (SELECT $2, $3));
+ t        |     1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab
+ f        |     1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab;
+ t        |     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) INSERT INTO stats_track_tab VALUES (($2));
+ t        |     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+
+          |       |     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) 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);
+ 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
+ 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;
+ 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) UPDATE stats_track_tab SET x=$2 WHERE x=$3;
+(15 rows)
+
+-- Explain with CTE - top-level tracking
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) WITH a AS (select 4) SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+            QUERY PLAN             
+-----------------------------------
+ Update on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+         Filter: (x = 1)
+(3 rows)
+
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+            QUERY PLAN             
+-----------------------------------
+ Delete on stats_track_tab
+   ->  Seq Scan on stats_track_tab
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+        QUERY PLAN         
+---------------------------
+ Insert on stats_track_tab
+   ->  Result
+(2 rows)
+
+explain (costs off) WITH a AS (select 4) 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);
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Merge on stats_track_tab
+   ->  Hash Right Join
+         Hash Cond: (stats_track_tab.x = id.id)
+         ->  Seq Scan on stats_track_tab
+         ->  Hash
+               ->  Function Scan on generate_series id
+(6 rows)
+
+explain (costs off) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+        QUERY PLAN        
+--------------------------
+ Unique
+   ->  Sort
+         Sort Key: (1)
+         ->  Append
+               ->  Result
+               ->  Result
+(6 rows)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                                                              query                                                               
+----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ t        |     1 | explain (costs off) (WITH a AS (select $1) (SELECT $2, $3))
+ t        |     1 | explain (costs off) WITH a AS (select $1) DELETE FROM stats_track_tab
+ t        |     1 | explain (costs off) WITH a AS (select $1) INSERT INTO stats_track_tab VALUES (($2))
+ t        |     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+
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                                                                         +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id)
+ 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
+(8 rows)
+
+-- Explain analyze, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                                                     query                                                     
+----------+-------+---------------------------------------------------------------------------------------------------------------
+ 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
+ f        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1;
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(4 rows)
+
+-- Explain analyze, top tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+           QUERY PLAN           
+--------------------------------
+ Result (actual rows=1 loops=1)
+(1 row)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+                     QUERY PLAN                      
+-----------------------------------------------------
+ Seq Scan on stats_track_tab (actual rows=0 loops=1)
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                                                     query                                                     
+----------+-------+---------------------------------------------------------------------------------------------------------------
+ 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
+(3 rows)
+
+-- Create Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                                             query                                              
+----------+-------+------------------------------------------------------------------------------------------------
+ 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)
+
+-- Create Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE MATERIALIZED VIEW pgss_materialized_view_2 AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                                              query                                               
+----------+-------+--------------------------------------------------------------------------------------------------
+ t        |     1 | CREATE MATERIALIZED VIEW pgss_materialized_view_2 AS SELECT * FROM generate_series($1, $2) as id
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Refresh Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Refresh Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                       query                        
+----------+-------+----------------------------------------------------
+ t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(2 rows)
+
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ 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 | PREPARE test_prepare_pgss AS select generate_series($1, $2)
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(4 rows)
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+CREATE TEMPORARY TABLE pgss_ctas_3 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_4 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_3 AS SELECT $1
+ t        |     1 | CREATE TEMPORARY TABLE pgss_ctas_4 AS EXECUTE test_prepare_pgss
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(3 rows)
+
+-- Explain with Create Table As - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
+ 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)
+
+-- Explain with Create Table As - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+ QUERY PLAN 
+------------
+ Result
+(1 row)
+
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
+ 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)
+
+-- Declare cursor, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                          query                          
+----------+-------+---------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ 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)
+
+-- Declare cursor, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+ x 
+---
+(0 rows)
+
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                          query                          
+----------+-------+---------------------------------------------------------
+ t        |     1 | BEGIN
+ t        |     1 | CLOSE foocur
+ t        |     1 | COMMIT
+ 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)
+
+-- COPY - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+COPY (SELECT 1) to stdout;
+1
+COPY (SELECT 1 UNION SELECT 2) to stdout;
+1
+2
+COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id
+    WHEN MATCHED THEN UPDATE SET x = id
+    WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout;
+1
+COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout;
+1
+COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout;
+2
+2
+COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout;
+2
+2
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                                  query                                   
+----------+-------+--------------------------------------------------------------------------
+ f        |     1 | COPY (DELETE FROM stats_track_tab WHERE x=$1 RETURNING x) to stdout
+ t        |     1 | COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout
+ f        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES ($1) RETURNING x) to stdout
+ t        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout
+ f        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT $1 id) ON x = id         +
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                 +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout
+ t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id          +
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                 +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout
+ f        |     1 | COPY (SELECT $1 UNION SELECT $2) to stdout
+ f        |     1 | COPY (SELECT $1) to stdout
+ t        |     1 | COPY (SELECT 1 UNION SELECT 2) to stdout
+ t        |     1 | COPY (SELECT 1) to stdout
+ f        |     1 | COPY (UPDATE stats_track_tab SET x=$1 WHERE x=$2 RETURNING x) to stdout
+ t        |     1 | COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(13 rows)
+
+-- COPY - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+ t 
+---
+ t
+(1 row)
+
+COPY (SELECT 1) to stdout;
+1
+COPY (SELECT 1 UNION SELECT 2) to stdout;
+1
+2
+COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id
+    WHEN MATCHED THEN UPDATE SET x = id
+    WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout;
+1
+COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout;
+1
+COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout;
+2
+2
+COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout;
+2
+2
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+ toplevel | calls |                                  query                                  
+----------+-------+-------------------------------------------------------------------------
+ t        |     1 | COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout
+ t        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout
+ t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id         +
+          |       |     WHEN MATCHED THEN UPDATE SET x = id                                +
+          |       |     WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout
+ t        |     1 | COPY (SELECT 1 UNION SELECT 2) to stdout
+ t        |     1 | COPY (SELECT 1) to stdout
+ t        |     1 | COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout
+ t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+(7 rows)
+
 -- 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 65a17147a5a..8970f648820 100644
--- a/contrib/pg_stat_statements/sql/level_tracking.sql
+++ b/contrib/pg_stat_statements/sql/level_tracking.sql
@@ -32,7 +32,6 @@ BEGIN
 END; $$;
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
-
 -- Procedure with multiple utility statements.
 CREATE OR REPLACE PROCEDURE proc_with_utility_stmt()
 LANGUAGE SQL
@@ -55,6 +54,234 @@ CALL proc_with_utility_stmt();
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
 
+-- Explain - all-level tracking.
+CREATE TABLE test_table (x int);
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1;
+explain (costs off) (SELECT 1, 2);
+explain (costs off) TABLE stats_track_tab;
+explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1);
+explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) DELETE FROM stats_track_tab;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((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 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Explain - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1;
+explain (costs off) (SELECT 1, 2);
+explain (costs off) TABLE stats_track_tab;
+explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1);
+explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) DELETE FROM stats_track_tab;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((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 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Explain - all-level tracking with multi statement.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES (1), (2);
+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;
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Explain - top-level tracking with multi statement.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) SELECT 1\; explain (costs off) SELECT 1, 2;
+explain (costs off) (SELECT 1, 2, 3)\; explain (costs off) (SELECT 1, 2, 3, 4);
+explain (costs off) TABLE stats_track_tab\; explain (costs off) (TABLE test_table);
+explain (costs off) VALUES(1)\; explain (costs off) (VALUES(1, 2));
+explain (costs off) UPDATE stats_track_tab SET x=1 WHERE x=1\; explain (costs off) UPDATE stats_track_tab SET x=1;
+explain (costs off) DELETE FROM stats_track_tab\; explain (costs off) DELETE FROM stats_track_tab WHERE x=1;
+explain (costs off) INSERT INTO stats_track_tab VALUES ((1))\; explain (costs off) INSERT INTO stats_track_tab VALUES ((1), (2));
+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;
+explain (costs off) SELECT 1, 2 UNION SELECT 3, 4\; explain (costs off) (SELECT 1, 2, 3) UNION SELECT 3, 4, 5;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Explain with CTE - all-level tracking
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) WITH a AS (select 4) SELECT 1;
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+explain (costs off) WITH a AS (select 4) 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) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Explain with CTE - top-level tracking
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) WITH a AS (select 4) SELECT 1;
+explain (costs off) (WITH a AS (select 4) (SELECT 1, 2));
+explain (costs off) WITH a AS (select 4) UPDATE stats_track_tab SET x=1 WHERE x=1;
+explain (costs off) WITH a AS (select 4) DELETE FROM stats_track_tab;
+explain (costs off) WITH a AS (select 4) INSERT INTO stats_track_tab VALUES ((1));
+explain (costs off) WITH a AS (select 4) 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) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Explain analyze, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Explain analyze, top tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT 100;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) DECLARE foocur CURSOR FOR SELECT * FROM stats_track_tab;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Create Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE MATERIALIZED VIEW pgss_materialized_view AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Create Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE MATERIALIZED VIEW pgss_materialized_view_2 AS SELECT * FROM generate_series(1, 5) as id;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Refresh Materialized View, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Refresh Materialized View, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+REFRESH MATERIALIZED VIEW pgss_materialized_view;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Create Table As, all-level tracking.
+SET pg_stat_statements.track = 'all';
+PREPARE test_prepare_pgss AS select generate_series(1, 10);
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_ctas_1 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_2 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Create Table As, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+CREATE TEMPORARY TABLE pgss_ctas_3 AS SELECT 1;
+CREATE TEMPORARY TABLE pgss_ctas_4 AS EXECUTE test_prepare_pgss;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Explain with Create Table As - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Explain with Create Table As - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+explain (costs off) CREATE TEMPORARY TABLE pgss_explain_ctas AS SELECT 1;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Declare cursor, all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- Declare cursor, top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+BEGIN;
+DECLARE FOOCUR CURSOR FOR SELECT * from stats_track_tab;
+FETCH FORWARD 1 FROM foocur;
+CLOSE foocur;
+COMMIT;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- COPY - all-level tracking.
+SET pg_stat_statements.track = 'all';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+COPY (SELECT 1) to stdout;
+COPY (SELECT 1 UNION SELECT 2) to stdout;
+COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id
+    WHEN MATCHED THEN UPDATE SET x = id
+    WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout;
+COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout;
+COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout;
+COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
+-- COPY - top-level tracking.
+SET pg_stat_statements.track = 'top';
+SELECT pg_stat_statements_reset() IS NOT NULL AS t;
+COPY (SELECT 1) to stdout;
+COPY (SELECT 1 UNION SELECT 2) to stdout;
+COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id
+    WHEN MATCHED THEN UPDATE SET x = id
+    WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) to stdout;
+COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) to stdout;
+COPY (UPDATE stats_track_tab SET x=2 WHERE x=1 RETURNING x) to stdout;
+COPY (DELETE FROM stats_track_tab WHERE x=2 RETURNING x) to stdout;
+SELECT toplevel, calls, query FROM pg_stat_statements
+  ORDER BY query COLLATE "C";
+
 -- DO block - top-level tracking without utility.
 SET pg_stat_statements.track = 'top';
 SET pg_stat_statements.track_utility = FALSE;
-- 
2.39.5 (Apple Git-154)

From 4308e20f7b5e6be26d40c69f41ed52185fa86bae Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Thu, 3 Oct 2024 10:21:44 +0200
Subject: Use view's definition as query string on a materialized view refresh

When creating a materialized view, the first refresh will have the
"Select" part of the statement as a query string. On subsequent refresh,
the "REFRESH MATERIALIZED" utility statement will be passed as query
string. This causes pgss to track both the top query and nested query as
a refresh.

This patch changes the query string on a refresh to fetch the view
definition instead. This will allow pgss to display the correct
statement when tracking refresh's nested query.
---
 .../expected/level_tracking.out               | 10 ++++----
 src/backend/commands/matview.c                | 24 +++++++++++++------
 src/backend/utils/adt/ruleutils.c             | 13 ++++++++++
 src/include/utils/ruleutils.h                 |  2 +-
 4 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index d1638a23159..c85a9deeec9 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -959,12 +959,14 @@ SELECT pg_stat_statements_reset() IS NOT NULL AS t;
 REFRESH MATERIALIZED VIEW pgss_materialized_view;
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                       query                        
-----------+-------+----------------------------------------------------
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
  t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
- f        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view;
+ f        |     1 | SELECT * FROM pg_catalog.pg_rewrite WHERE ev_class = $1 AND rulename = $2
+ f        |     1 | SELECT id                                                                +
+          |       |    FROM generate_series(1, 5) id(id);
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(3 rows)
+(4 rows)
 
 -- Refresh Materialized View, top-level tracking.
 SET pg_stat_statements.track = 'top';
diff --git a/src/backend/commands/matview.c b/src/backend/commands/matview.c
index 7cc68338837..92d5dd43864 100644
--- a/src/backend/commands/matview.c
+++ b/src/backend/commands/matview.c
@@ -39,6 +39,7 @@
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
+#include "utils/ruleutils.h"
 #include "utils/snapmgr.h"
 #include "utils/syscache.h"
 
@@ -61,8 +62,7 @@ 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,
-									   ParseState *pstate, const char *queryString,
-									   bool is_create);
+									   ParseState *pstate, 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);
@@ -327,11 +327,21 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData,
 	if (!skipData)
 	{
 		DestReceiver *dest;
-		const char *queryString = pstate->p_sourcetext;
+		ParseState *refresh_pstate = pstate;
+
+		/*
+		 * On refresh, the pstate's source text will be the refresh utility
+		 * statement. We need to fetch the the view definition to get the
+		 * query executed by the refresh.
+		 */
+		if (!is_create)
+		{
+			refresh_pstate = make_parsestate(NULL);
+			refresh_pstate->p_sourcetext = pg_get_viewdef_string(matviewOid, false);
+		}
 
 		dest = CreateTransientRelDestReceiver(OIDNewHeap);
-		processed = refresh_matview_datafill(dest, dataQuery, pstate,
-											 queryString, is_create);
+		processed = refresh_matview_datafill(dest, dataQuery, refresh_pstate, is_create);
 	}
 
 	/* Make the matview match the newly generated data. */
@@ -406,8 +416,7 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData,
  */
 static uint64
 refresh_matview_datafill(DestReceiver *dest, Query *query,
-						 ParseState *pstate, const char *queryString,
-						 bool is_create)
+						 ParseState *pstate, bool is_create)
 {
 	List	   *rewritten;
 	PlannedStmt *plan;
@@ -415,6 +424,7 @@ refresh_matview_datafill(DestReceiver *dest, Query *query,
 	Query	   *copied_query;
 	uint64		processed;
 	JumbleState *jstate = NULL;
+	const char *queryString = pstate->p_sourcetext;
 
 	/* Lock and rewrite, using a copy to preserve the original query. */
 	copied_query = copyObject(query);
diff --git a/src/backend/utils/adt/ruleutils.c b/src/backend/utils/adt/ruleutils.c
index 2177d17e278..0bf85cbb759 100644
--- a/src/backend/utils/adt/ruleutils.c
+++ b/src/backend/utils/adt/ruleutils.c
@@ -776,6 +776,19 @@ pg_get_viewdef_name_ext(PG_FUNCTION_ARGS)
 	PG_RETURN_TEXT_P(string_to_text(res));
 }
 
+/*
+ * Internal version of pg_get_viewdef
+ */
+char *
+pg_get_viewdef_string(Oid viewoid, bool pretty)
+{
+	int			prettyFlags;
+
+	prettyFlags = GET_PRETTY_FLAGS(pretty);
+
+	return pg_get_viewdef_worker(viewoid, prettyFlags, WRAP_COLUMN_DEFAULT);
+}
+
 /*
  * Common code for by-OID and by-name variants of pg_get_viewdef
  */
diff --git a/src/include/utils/ruleutils.h b/src/include/utils/ruleutils.h
index 161fb5ef02c..72177b9dce3 100644
--- a/src/include/utils/ruleutils.h
+++ b/src/include/utils/ruleutils.h
@@ -29,7 +29,7 @@ extern char *pg_get_indexdef_columns(Oid indexrelid, bool pretty);
 extern char *pg_get_indexdef_columns_extended(Oid indexrelid,
 											  bits16 flags);
 extern char *pg_get_querydef(Query *query, bool pretty);
-
+extern char *pg_get_viewdef_string(Oid viewoid, bool pretty);
 extern char *pg_get_partkeydef_columns(Oid relid, bool pretty);
 extern char *pg_get_partconstrdef_string(Oid partitionId, char *aliasname);
 
-- 
2.39.5 (Apple Git-154)

From f630b9445100e1ef96e989ef064bd1e3fcb9f171 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Tue, 22 Oct 2024 13:28:28 +0900
Subject: [PATCH v10 1/3] Track location to extract relevant part in nested
 statement

Previously, Query generated through transform would have unset
stmt_location. Extensions relying on the statement location to extract
the relevant part of the statement would fallback to use the whole
statement instead, thus showing the same string in the top and
nested level which was a source of confusion.

This patch fixes the issue by keeping track of the statement locations
and propagate it to Query during transform, allowing pgss to only show
the relevant part of the query for nested query.
---
 src/include/nodes/parsenodes.h                |  10 ++
 src/include/parser/analyze.h                  |   3 +-
 src/include/parser/parse_node.h               |   2 +
 src/backend/optimizer/util/clauses.c          |   2 +-
 src/backend/parser/analyze.c                  |  96 +++++++---
 src/backend/parser/gram.y                     |  71 +++++++-
 src/backend/parser/parse_merge.c              |   2 +
 .../expected/level_tracking.out               | 165 +++++++++---------
 .../pg_stat_statements/expected/planning.out  |  10 +-
 .../pg_stat_statements/expected/select.out    |   2 +-
 .../pg_stat_statements/expected/utility.out   |   2 +-
 contrib/pg_stat_statements/sql/planning.sql   |   4 +-
 12 files changed, 240 insertions(+), 129 deletions(-)

diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index c92cef3d16..b40b661ec8 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -2046,6 +2046,8 @@ typedef struct InsertStmt
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
 	OverridingKind override;	/* OVERRIDING clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } InsertStmt;
 
 /* ----------------------
@@ -2060,6 +2062,8 @@ typedef struct DeleteStmt
 	Node	   *whereClause;	/* qualifications */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } DeleteStmt;
 
 /* ----------------------
@@ -2075,6 +2079,8 @@ typedef struct UpdateStmt
 	List	   *fromClause;		/* optional from clause for more tables */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } UpdateStmt;
 
 /* ----------------------
@@ -2090,6 +2096,8 @@ typedef struct MergeStmt
 	List	   *mergeWhenClauses;	/* list of MergeWhenClause(es) */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } MergeStmt;
 
 /* ----------------------
@@ -2159,6 +2167,8 @@ typedef struct SelectStmt
 	bool		all;			/* ALL specified? */
 	struct SelectStmt *larg;	/* left child */
 	struct SelectStmt *rarg;	/* right child */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 	/* Eventually add fields for CORRESPONDING spec here */
 } SelectStmt;
 
diff --git a/src/include/parser/analyze.h b/src/include/parser/analyze.h
index 28b66fccb4..8ba4e050af 100644
--- a/src/include/parser/analyze.h
+++ b/src/include/parser/analyze.h
@@ -46,8 +46,9 @@ extern List *transformUpdateTargetList(ParseState *pstate,
 									   List *origTlist);
 extern List *transformReturningList(ParseState *pstate, List *returningList,
 									ParseExprKind exprKind);
-extern Query *transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree);
+extern Query *transformOptionalSelectInto(ParseState *pstate, Node *parseTree);
 extern Query *transformStmt(ParseState *pstate, Node *parseTree);
+extern void setQueryStmtLen(ParseState *pstate, Query *qry, int stmt_len);
 
 extern bool stmt_requires_parse_analysis(RawStmt *parseTree);
 extern bool analyze_requires_snapshot(RawStmt *parseTree);
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/backend/optimizer/util/clauses.c b/src/backend/optimizer/util/clauses.c
index 8e39795e24..2bd14335bd 100644
--- a/src/backend/optimizer/util/clauses.c
+++ b/src/backend/optimizer/util/clauses.c
@@ -4686,7 +4686,7 @@ inline_function(Oid funcid, Oid result_type, Oid result_collid,
 		pstate->p_sourcetext = src;
 		sql_fn_parser_setup(pstate, pinfo);
 
-		querytree = transformTopLevelStmt(pstate, linitial(raw_parsetree_list));
+		querytree = transformOptionalSelectInto(pstate, ((RawStmt *) linitial(raw_parsetree_list))->stmt);
 
 		free_parsestate(pstate);
 	}
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index 8a6ba1692e..5f4372c4a0 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -58,7 +58,6 @@
 /* Hook for plugins to get control at end of parse analysis */
 post_parse_analyze_hook_type post_parse_analyze_hook = NULL;
 
-static Query *transformOptionalSelectInto(ParseState *pstate, Node *parseTree);
 static Query *transformDeleteStmt(ParseState *pstate, DeleteStmt *stmt);
 static Query *transformInsertStmt(ParseState *pstate, InsertStmt *stmt);
 static OnConflictExpr *transformOnConflictClause(ParseState *pstate,
@@ -113,13 +112,15 @@ 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);
 
 	pstate->p_queryEnv = queryEnv;
 
-	query = transformTopLevelStmt(pstate, parseTree);
+	query = transformOptionalSelectInto(pstate, parseTree->stmt);
 
 	if (IsQueryIdEnabled())
 		jstate = JumbleQuery(query);
@@ -153,12 +154,14 @@ 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);
 
 	pstate->p_queryEnv = queryEnv;
 
-	query = transformTopLevelStmt(pstate, parseTree);
+	query = transformOptionalSelectInto(pstate, parseTree->stmt);
 
 	/* make sure all is well with parameter types */
 	check_variable_parameters(pstate, query);
@@ -195,10 +198,12 @@ 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);
 
-	query = transformTopLevelStmt(pstate, parseTree);
+	query = transformOptionalSelectInto(pstate, parseTree->stmt);
 
 	if (IsQueryIdEnabled())
 		jstate = JumbleQuery(query);
@@ -238,27 +243,6 @@ parse_sub_analyze(Node *parseTree, ParseState *parentParseState,
 	return query;
 }
 
-/*
- * transformTopLevelStmt -
- *	  transform a Parse tree into a Query tree.
- *
- * This function is just responsible for transferring statement location data
- * from the RawStmt into the finished Query.
- */
-Query *
-transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree)
-{
-	Query	   *result;
-
-	/* We're at top level, so allow SELECT INTO */
-	result = transformOptionalSelectInto(pstate, parseTree->stmt);
-
-	result->stmt_location = parseTree->stmt_location;
-	result->stmt_len = parseTree->stmt_len;
-
-	return result;
-}
-
 /*
  * transformOptionalSelectInto -
  *	  If SELECT has INTO, convert it to CREATE TABLE AS.
@@ -269,7 +253,7 @@ transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree)
  * of the parse tree, and so we only try it before entering the recursive
  * transformStmt() processing.
  */
-static Query *
+Query *
 transformOptionalSelectInto(ParseState *pstate, Node *parseTree)
 {
 	if (IsA(parseTree, SelectStmt))
@@ -417,7 +401,9 @@ transformStmt(ParseState *pstate, Node *parseTree)
 			 */
 			result = makeNode(Query);
 			result->commandType = CMD_UTILITY;
-			result->utilityStmt = (Node *) parseTree;
+			result->utilityStmt = parseTree;
+			result->stmt_location = pstate->p_stmt_location;
+			result->stmt_len = pstate->p_stmt_len;
 			break;
 	}
 
@@ -506,6 +492,37 @@ analyze_requires_snapshot(RawStmt *parseTree)
 	return stmt_requires_parse_analysis(parseTree);
 }
 
+/*
+ * setQueryStmtLen
+ *		Set stmt_len in Query.
+ *
+ * Some statements, like PreparableStmt, can be located within parentheses.
+ * For example "(SELECT 1)" or "COPY (UPDATE ...) to x;". For those, we can't
+ * use the whole string from the statement's location or the SQL string will
+ * yield "SELECT 1)". The parser will set stmt_len, reflecting the size of the
+ * statement within the parentheses. Thus, when stmt_len is available, we use it
+ * for the Query's stmt_len.
+ *
+ * For other cases, the parser can't provide the length of individual statements.
+ * However, we have the statement's location plus the length (p_stmt_len) and
+ * location (p_stmt_location) of the top level RawStmt, stored in pstate. Thus,
+ * the statement's length is the RawStmt's length minus how much we've advanced
+ * in the RawStmt's string.
+ */
+void
+setQueryStmtLen(ParseState *pstate, Query *qry, int stmt_len)
+{
+	if (stmt_len > 0)
+		/* Statement's length is known, use it */
+		qry->stmt_len = stmt_len;
+	else
+		/*
+		 * Compute the statement's length from statement's location and RawStmt's
+		 * length and location
+		 */
+		qry->stmt_len = pstate->p_stmt_len - (qry->stmt_location - pstate->p_stmt_location);
+}
+
 /*
  * transformDeleteStmt -
  *	  transforms a Delete Statement
@@ -518,6 +535,8 @@ transformDeleteStmt(ParseState *pstate, DeleteStmt *stmt)
 	Node	   *qual;
 
 	qry->commandType = CMD_DELETE;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 
 	/* process the WITH clause independently of all else */
 	if (stmt->withClause)
@@ -606,6 +625,8 @@ transformInsertStmt(ParseState *pstate, InsertStmt *stmt)
 	Assert(pstate->p_ctenamespace == NIL);
 
 	qry->commandType = CMD_INSERT;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 	pstate->p_is_insert = true;
 
 	/* process the WITH clause independently of all else */
@@ -1331,7 +1352,6 @@ count_rowexpr_columns(ParseState *pstate, Node *expr)
 	return -1;
 }
 
-
 /*
  * transformSelectStmt -
  *	  transforms a Select Statement
@@ -1347,6 +1367,8 @@ transformSelectStmt(ParseState *pstate, SelectStmt *stmt)
 	ListCell   *l;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 
 	/* process the WITH clause independently of all else */
 	if (stmt->withClause)
@@ -1499,6 +1521,8 @@ transformValuesClause(ParseState *pstate, SelectStmt *stmt)
 	int			i;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 
 	/* Most SELECT stuff doesn't apply in a VALUES clause */
 	Assert(stmt->distinctClause == NIL);
@@ -1730,6 +1754,8 @@ transformSetOperationStmt(ParseState *pstate, SelectStmt *stmt)
 	int			tllen;
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 
 	/*
 	 * Find leftmost leaf SelectStmt.  We currently only need to do this in
@@ -2397,6 +2423,8 @@ transformReturnStmt(ParseState *pstate, ReturnStmt *stmt)
 	Query	   *qry = makeNode(Query);
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = pstate->p_stmt_location;
+	qry->stmt_len = pstate->p_stmt_len;
 	qry->isReturn = true;
 
 	qry->targetList = list_make1(makeTargetEntry((Expr *) transformExpr(pstate, stmt->returnval, EXPR_KIND_SELECT_TARGET),
@@ -2430,6 +2458,8 @@ transformUpdateStmt(ParseState *pstate, UpdateStmt *stmt)
 	Node	   *qual;
 
 	qry->commandType = CMD_UPDATE;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 	pstate->p_is_insert = false;
 
 	/* process the WITH clause independently of all else */
@@ -2677,6 +2707,8 @@ transformPLAssignStmt(ParseState *pstate, PLAssignStmt *stmt)
 	 * consider WITH or INTO, and we build a targetlist our own way.
 	 */
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = pstate->p_stmt_location;
+	qry->stmt_len = pstate->p_stmt_len;
 	pstate->p_is_insert = false;
 
 	/* make FOR UPDATE/FOR SHARE info available to addRangeTableEntry */
@@ -2948,6 +2980,8 @@ transformDeclareCursorStmt(ParseState *pstate, DeclareCursorStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3003,6 +3037,8 @@ transformExplainStmt(ParseState *pstate, ExplainStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3083,6 +3119,8 @@ transformCreateTableAsStmt(ParseState *pstate, CreateTableAsStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3207,6 +3245,8 @@ transformCallStmt(ParseState *pstate, CallStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 4aa8646af7..d00f4b47c4 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -167,6 +167,7 @@ static void base_yyerror(YYLTYPE *yylloc, core_yyscan_t yyscanner,
 						 const char *msg);
 static RawStmt *makeRawStmt(Node *stmt, int stmt_location);
 static void updateRawStmtEnd(RawStmt *rs, int end_location);
+static void updatePreparableStmtEnd(Node *n, int end_location);
 static Node *makeColumnRef(char *colname, List *indirection,
 						   int location, core_yyscan_t yyscanner);
 static Node *makeTypeCast(Node *arg, TypeName *typename, int location);
@@ -190,7 +191,7 @@ static void insertSelectOptions(SelectStmt *stmt,
 								SelectLimit *limitClause,
 								WithClause *withClause,
 								core_yyscan_t yyscanner);
-static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg);
+static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location);
 static Node *doNegate(Node *n, int location);
 static void doNegateFloat(Float *v);
 static Node *makeAndExpr(Node *lexpr, Node *rexpr, int location);
@@ -3403,6 +3404,7 @@ CopyStmt:	COPY opt_binary qualified_name opt_column_list
 				{
 					CopyStmt *n = makeNode(CopyStmt);
 
+					updatePreparableStmtEnd($3, @4);
 					n->relation = NULL;
 					n->query = $3;
 					n->attlist = NIL;
@@ -12170,6 +12172,9 @@ InsertStmt:
 					$5->onConflictClause = $6;
 					$5->returningList = $7;
 					$5->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					$5->stmt_location = @$;
 					$$ = (Node *) $5;
 				}
 		;
@@ -12323,6 +12328,9 @@ DeleteStmt: opt_with_clause DELETE_P FROM relation_expr_opt_alias
 					n->whereClause = $6;
 					n->returningList = $7;
 					n->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					n->stmt_location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12397,6 +12405,9 @@ UpdateStmt: opt_with_clause UPDATE relation_expr_opt_alias
 					n->whereClause = $7;
 					n->returningList = $8;
 					n->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					n->stmt_location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12474,6 +12485,9 @@ MergeStmt:
 					m->joinCondition = $8;
 					m->mergeWhenClauses = $9;
 					m->returningList = $10;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					m->stmt_location = @$;
 
 					$$ = (Node *) m;
 				}
@@ -12714,7 +12728,12 @@ SelectStmt: select_no_parens			%prec UMINUS
 		;
 
 select_with_parens:
-			'(' select_no_parens ')'				{ $$ = $2; }
+			'(' select_no_parens ')'
+				{
+					SelectStmt *n = (SelectStmt *) $2;
+					n->stmt_len = @3 - @2;
+					$$ = $2;
+				}
 			| '(' select_with_parens ')'			{ $$ = $2; }
 		;
 
@@ -12836,6 +12855,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| SELECT distinct_clause target_list
@@ -12853,6 +12873,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| values_clause							{ $$ = $1; }
@@ -12873,19 +12894,20 @@ simple_select:
 
 					n->targetList = list_make1(rt);
 					n->fromClause = list_make1($2);
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| select_clause UNION set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause INTERSECT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause EXCEPT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 		;
 
@@ -13443,6 +13465,7 @@ values_clause:
 				{
 					SelectStmt *n = makeNode(SelectStmt);
 
+					n->stmt_location = @1;
 					n->valuesLists = list_make1($3);
 					$$ = (Node *) n;
 				}
@@ -18585,6 +18608,39 @@ updateRawStmtEnd(RawStmt *rs, int end_location)
 	rs->stmt_len = end_location - rs->stmt_location;
 }
 
+/* Adjust a PreparableStmt to reflect that it doesn't run to the end of the string */
+static void
+updatePreparableStmtEnd(Node *n, int end_location)
+{
+	if (IsA(n, SelectStmt))
+	{
+		SelectStmt *stmt = (SelectStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, InsertStmt))
+	{
+		InsertStmt *stmt = (InsertStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, UpdateStmt))
+	{
+		UpdateStmt *stmt = (UpdateStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, DeleteStmt))
+	{
+		DeleteStmt *stmt = (DeleteStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, MergeStmt))
+	{
+		MergeStmt *stmt = (MergeStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else
+		elog(ERROR, "unexpected node type %d", (int) n->type);
+}
+
 static Node *
 makeColumnRef(char *colname, List *indirection,
 			  int location, core_yyscan_t yyscanner)
@@ -18963,11 +19019,13 @@ insertSelectOptions(SelectStmt *stmt,
 					 errmsg("multiple WITH clauses not allowed"),
 					 parser_errposition(exprLocation((Node *) withClause))));
 		stmt->withClause = withClause;
+		/* Update SelectStmt's location to the start of the with clause */
+		stmt->stmt_location = withClause->location;
 	}
 }
 
 static Node *
-makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
+makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location)
 {
 	SelectStmt *n = makeNode(SelectStmt);
 
@@ -18975,6 +19033,7 @@ makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
 	n->all = all;
 	n->larg = (SelectStmt *) larg;
 	n->rarg = (SelectStmt *) rarg;
+	n->stmt_location = location;
 	return (Node *) n;
 }
 
diff --git a/src/backend/parser/parse_merge.c b/src/backend/parser/parse_merge.c
index 87df79027d..fdfa61868c 100644
--- a/src/backend/parser/parse_merge.c
+++ b/src/backend/parser/parse_merge.c
@@ -118,6 +118,8 @@ transformMergeStmt(ParseState *pstate, MergeStmt *stmt)
 	Assert(pstate->p_ctenamespace == NIL);
 
 	qry->commandType = CMD_MERGE;
+	qry->stmt_location = stmt->stmt_location;
+	setQueryStmtLen(pstate, qry, stmt->stmt_len);
 	qry->hasRecursive = false;
 
 	/* process the WITH clause independently of all else */
diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 8f008f8bfd..489dc7143f 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -206,37 +206,37 @@ EXPLAIN (COSTS OFF) SELECT 1 UNION SELECT 2;
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                query                                
-----------+-------+---------------------------------------------------------------------
+ toplevel | calls |                               query                                
+----------+-------+--------------------------------------------------------------------
+ f        |     1 | DELETE FROM stats_track_tab
  t        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2)
- f        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2);
  t        |     1 | EXPLAIN (COSTS OFF) (TABLE test_table)
- f        |     1 | EXPLAIN (COSTS OFF) (TABLE test_table);
  t        |     1 | EXPLAIN (COSTS OFF) (VALUES ($1, $2))
- f        |     1 | EXPLAIN (COSTS OFF) (VALUES ($1, $2));
  t        |     1 | EXPLAIN (COSTS OFF) DELETE FROM stats_track_tab
- f        |     1 | EXPLAIN (COSTS OFF) DELETE FROM stats_track_tab;
  t        |     1 | EXPLAIN (COSTS OFF) INSERT INTO stats_track_tab VALUES (($1))
- f        |     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)
- 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);
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1 UNION SELECT $2
- f        |     1 | EXPLAIN (COSTS OFF) SELECT $1 UNION SELECT $2;
- f        |     1 | EXPLAIN (COSTS OFF) SELECT $1;
  t        |     1 | EXPLAIN (COSTS OFF) TABLE stats_track_tab
- f        |     1 | EXPLAIN (COSTS OFF) TABLE stats_track_tab;
  t        |     1 | EXPLAIN (COSTS OFF) UPDATE stats_track_tab SET x = $1 WHERE x = $2
- f        |     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) VALUES ($1);
+ 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 UNION SELECT $2
+ f        |     1 | SELECT $1, $2
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | TABLE test_table
+ f        |     1 | UPDATE stats_track_tab SET x = $1 WHERE x = $2
+ f        |     1 | VALUES ($1)
+ f        |     1 | VALUES ($1, $2)
 (23 rows)
 
 -- EXPLAIN - top-level tracking.
@@ -405,20 +405,20 @@ EXPLAIN (COSTS OFF) SELECT 1, 2 UNION SELECT 3, 4\; EXPLAIN (COSTS OFF) (SELECT
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                                        query                                                        
-----------+-------+---------------------------------------------------------------------------------------------------------------------
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ f        |     1 | (SELECT $1, $2, $3) UNION SELECT $4, $5, $6
  t        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2, $3)
  t        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2, $3) UNION SELECT $4, $5, $6
- f        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2, $3); EXPLAIN (COSTS OFF) (SELECT 1, 2, 3, 4);
  t        |     1 | 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);
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1, $2
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1, $2 UNION SELECT $3, $4
- 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 | SELECT $1
+ f        |     1 | SELECT $1, $2
+ f        |     1 | SELECT $1, $2 UNION SELECT $3, $4
+ f        |     1 | SELECT $1, $2, $3
+ f        |     1 | SELECT $1, $2, $3, $4
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (13 rows)
 
@@ -494,29 +494,29 @@ EXPLAIN (COSTS OFF) INSERT INTO stats_track_tab VALUES ((1))\; EXPLAIN (COSTS OF
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                                              query                                                               
-----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ toplevel | calls |                               query                                
+----------+-------+--------------------------------------------------------------------
+ f        |     1 | DELETE FROM stats_track_tab
+ f        |     1 | DELETE FROM stats_track_tab WHERE x = $1
  t        |     1 | EXPLAIN (COSTS OFF) (TABLE test_table)
  t        |     1 | EXPLAIN (COSTS OFF) (VALUES ($1, $2))
  t        |     1 | EXPLAIN (COSTS OFF) DELETE FROM stats_track_tab
  t        |     1 | 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) DELETE FROM stats_track_tab; 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))
- 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);
  t        |     1 | EXPLAIN (COSTS OFF) TABLE stats_track_tab
- 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);
  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
- 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;
  t        |     1 | EXPLAIN (COSTS OFF) VALUES ($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 | INSERT INTO stats_track_tab VALUES ($1), ($2)
+ f        |     1 | INSERT INTO stats_track_tab VALUES (($1))
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | TABLE test_table
+ 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)
+ f        |     1 | VALUES ($1, $2)
 (21 rows)
 
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
@@ -547,21 +547,18 @@ EXPLAIN (COSTS OFF) MERGE INTO stats_track_tab
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                             query                                              
-----------+-------+------------------------------------------------------------------------------------------------
- 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                                                         +
+ toplevel | calls |                             query                             
+----------+-------+---------------------------------------------------------------
+ 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)
- 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;
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1, $2, $3, $4, $5
+ 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, $2, $3, $4, $5
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (5 rows)
 
@@ -789,29 +786,29 @@ EXPLAIN (COSTS OFF) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                           query                                           
-----------+-------+-------------------------------------------------------------------------------------------
+ toplevel | calls |                                          query                                           
+----------+-------+------------------------------------------------------------------------------------------
  t        |     1 | EXPLAIN (COSTS OFF) (WITH a AS (SELECT $1) (SELECT $2, $3))
- f        |     1 | EXPLAIN (COSTS OFF) (WITH a AS (SELECT $1) (SELECT $2, $3));
  t        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) DELETE FROM stats_track_tab
- f        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) DELETE FROM stats_track_tab;
  t        |     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) INSERT INTO stats_track_tab VALUES (($2));
- t        |     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                            +
-          |       |   WHEN MATCHED THEN UPDATE SET x = id                                                    +
+ t        |     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                           +
+          |       |   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) 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);
  t        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) SELECT $2
- f        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) SELECT $2;
  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) UPDATE stats_track_tab SET x = $2 WHERE x = $3;
  t        |     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 UNION SELECT $3;
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ 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 | WITH a AS (SELECT $1) SELECT $2
+ f        |     1 | WITH a AS (SELECT $1) UPDATE stats_track_tab SET x = $2 WHERE x = $3
+ f        |     1 | WITH a AS (select $1) SELECT $2 UNION SELECT $3
 (15 rows)
 
 -- EXPLAIN with CTEs - top-level tracking
@@ -921,12 +918,12 @@ EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                              query                               
-----------+-------+------------------------------------------------------------------
- t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)           +
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ 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
- f        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1;
+ f        |     1 | SELECT $1
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (4 rows)
 
@@ -1050,8 +1047,8 @@ 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 | PREPARE test_prepare_pgss AS select generate_series($1, $2)
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | select generate_series($1, $2)
 (4 rows)
 
 -- CREATE TABLE AS, top-level tracking.
@@ -1202,25 +1199,25 @@ COPY (DELETE FROM stats_track_tab WHERE x = 2 RETURNING x) TO stdout;
 2
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                    query                                    
-----------+-------+-----------------------------------------------------------------------------
- f        |     1 | COPY (DELETE FROM stats_track_tab WHERE x = $1 RETURNING x) TO stdout
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
  t        |     1 | COPY (DELETE FROM stats_track_tab WHERE x = 2 RETURNING x) TO stdout
- f        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES ($1) RETURNING x) TO stdout
  t        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) TO stdout
- f        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT $1 id) ON x = id            +
-          |       |   WHEN MATCHED THEN UPDATE SET x = id                                      +
+ t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id           +
+          |       |   WHEN MATCHED THEN UPDATE SET x = id                                    +
           |       |   WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) TO stdout
- t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id             +
-          |       |   WHEN MATCHED THEN UPDATE SET x = id                                      +
-          |       |   WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) TO stdout
- f        |     1 | COPY (SELECT $1 UNION SELECT $2) TO stdout
- f        |     1 | COPY (SELECT $1) TO stdout
  t        |     1 | COPY (SELECT 1 UNION SELECT 2) TO stdout
  t        |     1 | COPY (SELECT 1) TO stdout
- f        |     1 | COPY (UPDATE stats_track_tab SET x = $1 WHERE x = $2 RETURNING x) TO stdout
  t        |     1 | COPY (UPDATE stats_track_tab SET x = 2 WHERE x = 1 RETURNING x) TO stdout
+ f        |     1 | DELETE FROM stats_track_tab WHERE x = $1 RETURNING x
+ f        |     1 | INSERT INTO stats_track_tab (x) VALUES ($1) RETURNING x
+ f        |     1 | MERGE INTO stats_track_tab USING (SELECT $1 id) ON x = id                +
+          |       |   WHEN MATCHED THEN UPDATE SET x = id                                    +
+          |       |   WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x
+ f        |     1 | SELECT $1
+ f        |     1 | SELECT $1 UNION SELECT $2
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | UPDATE stats_track_tab SET x = $1 WHERE x = $2 RETURNING x
 (13 rows)
 
 -- COPY - top-level tracking.
diff --git a/contrib/pg_stat_statements/expected/planning.out b/contrib/pg_stat_statements/expected/planning.out
index 9effd11fdc..3ee1928cbe 100644
--- a/contrib/pg_stat_statements/expected/planning.out
+++ b/contrib/pg_stat_statements/expected/planning.out
@@ -58,7 +58,7 @@ SELECT 42;
 (1 row)
 
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
  plans | calls | rows |                          query                           
 -------+-------+------+----------------------------------------------------------
      0 |     1 |    0 | ALTER TABLE stats_plan_test ADD COLUMN x int
@@ -72,10 +72,10 @@ SELECT plans, calls, rows, query FROM pg_stat_statements
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
- plans_ok | calls | rows |                         query                         
-----------+-------+------+-------------------------------------------------------
- t        |     4 |    4 | PREPARE prep1 AS SELECT COUNT(*) FROM stats_plan_test
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
+ plans_ok | calls | rows |                query                 
+----------+-------+------+--------------------------------------
+ t        |     4 |    4 | SELECT COUNT(*) FROM stats_plan_test
 (1 row)
 
 -- Cleanup
diff --git a/contrib/pg_stat_statements/expected/select.out b/contrib/pg_stat_statements/expected/select.out
index dd6c756f67..217a2c0b2b 100644
--- a/contrib/pg_stat_statements/expected/select.out
+++ b/contrib/pg_stat_statements/expected/select.out
@@ -127,7 +127,6 @@ DEALLOCATE pgss_test;
 SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
  calls | rows |                                    query                                     
 -------+------+------------------------------------------------------------------------------
-     1 |    1 | PREPARE pgss_test (int) AS SELECT $1, $2 LIMIT $3
      4 |    4 | SELECT $1                                                                   +
        |      |   -- multiline                                                              +
        |      |   AS "text"
@@ -137,6 +136,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
      2 |    2 | SELECT $1 AS "int"
      1 |    2 | SELECT $1 AS i UNION SELECT $2 ORDER BY i
      1 |    1 | SELECT $1 || $2
+     1 |    1 | SELECT $1, $2 LIMIT $3
      0 |    0 | SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C"
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
      1 |    2 | WITH t(f) AS (                                                              +
diff --git a/contrib/pg_stat_statements/expected/utility.out b/contrib/pg_stat_statements/expected/utility.out
index 060d4416dd..aa4f0f7e62 100644
--- a/contrib/pg_stat_statements/expected/utility.out
+++ b/contrib/pg_stat_statements/expected/utility.out
@@ -540,7 +540,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
 -------+------+----------------------------------------------------
      2 |    0 | DEALLOCATE $1
      2 |    0 | DEALLOCATE ALL
-     2 |    2 | PREPARE stat_select AS SELECT $1 AS a
+     2 |    2 | SELECT $1 AS a
      1 |    1 | SELECT $1 as a
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (5 rows)
diff --git a/contrib/pg_stat_statements/sql/planning.sql b/contrib/pg_stat_statements/sql/planning.sql
index 46f5d9b951..9cfe206b3b 100644
--- a/contrib/pg_stat_statements/sql/planning.sql
+++ b/contrib/pg_stat_statements/sql/planning.sql
@@ -20,11 +20,11 @@ SELECT 42;
 SELECT 42;
 SELECT 42;
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 
 -- Cleanup
 DROP TABLE stats_plan_test;
-- 
2.45.2

From 66e34b70813b548045cc180b745a8b72869e4fff Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Tue, 22 Oct 2024 13:30:14 +0900
Subject: [PATCH v10 2/3] 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.
---
 src/include/commands/explain.h                |  4 +-
 src/include/commands/matview.h                |  5 ++-
 src/include/commands/prepare.h                |  4 +-
 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/test/regress/expected/explain.out         | 17 ++++++++
 src/test/regress/sql/explain.sql              |  4 ++
 .../expected/level_tracking.out               | 18 +++++---
 12 files changed, 116 insertions(+), 48 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/matview.h b/src/include/commands/matview.h
index c8811e8fc7..6602640b40 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 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..7ed42aa441 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 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/matview.c b/src/backend/commands/matview.c
index 010097873d..7cc6833883 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 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/backend/tcop/utility.c b/src/backend/tcop/utility.c
index b2ea8125c9..4768b4f746 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/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..5300bc3baf 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';
@@ -974,8 +975,9 @@ 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
+ f        |     1 | SELECT * FROM generate_series($1, $2) as id
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(2 rows)
+(3 rows)
 
 -- CREATE MATERIALIZED VIEW, top-level tracking.
 SET pg_stat_statements.track = 'top';
@@ -1010,8 +1012,9 @@ SELECT toplevel, calls, query FROM pg_stat_statements
  toplevel | calls |                       query                        
 ----------+-------+----------------------------------------------------
  t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
+ f        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view;
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(2 rows)
+(3 rows)
 
 -- REFRESH MATERIALIZED VIEW, top-level tracking.
 SET pg_stat_statements.track = 'top';
@@ -1047,9 +1050,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 +1093,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 +1145,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

From 71821c2a07b7414828b308d9f33ed9e6151ec475 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Tue, 22 Oct 2024 13:32:05 +0900
Subject: [PATCH v10 3/3] Use view's definition as query string on a
 materialized view refresh

When creating a materialized view, the first refresh will have the
"Select" part of the statement as a query string. On subsequent refresh,
the "REFRESH MATERIALIZED" utility statement will be passed as query
string. This causes pgss to track both the top query and nested query as
a refresh.

This patch changes the query string on a refresh to fetch the view
definition instead. This will allow pgss to display the correct
statement when tracking refresh's nested query.
---
 src/include/utils/ruleutils.h                 |  2 +-
 src/backend/commands/matview.c                | 24 +++++++++++++------
 src/backend/utils/adt/ruleutils.c             | 13 ++++++++++
 .../expected/level_tracking.out               | 10 ++++----
 4 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/src/include/utils/ruleutils.h b/src/include/utils/ruleutils.h
index 161fb5ef02..72177b9dce 100644
--- a/src/include/utils/ruleutils.h
+++ b/src/include/utils/ruleutils.h
@@ -29,7 +29,7 @@ extern char *pg_get_indexdef_columns(Oid indexrelid, bool pretty);
 extern char *pg_get_indexdef_columns_extended(Oid indexrelid,
 											  bits16 flags);
 extern char *pg_get_querydef(Query *query, bool pretty);
-
+extern char *pg_get_viewdef_string(Oid viewoid, bool pretty);
 extern char *pg_get_partkeydef_columns(Oid relid, bool pretty);
 extern char *pg_get_partconstrdef_string(Oid partitionId, char *aliasname);
 
diff --git a/src/backend/commands/matview.c b/src/backend/commands/matview.c
index 7cc6833883..92d5dd4386 100644
--- a/src/backend/commands/matview.c
+++ b/src/backend/commands/matview.c
@@ -39,6 +39,7 @@
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
+#include "utils/ruleutils.h"
 #include "utils/snapmgr.h"
 #include "utils/syscache.h"
 
@@ -61,8 +62,7 @@ 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,
-									   ParseState *pstate, const char *queryString,
-									   bool is_create);
+									   ParseState *pstate, 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);
@@ -327,11 +327,21 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData,
 	if (!skipData)
 	{
 		DestReceiver *dest;
-		const char *queryString = pstate->p_sourcetext;
+		ParseState *refresh_pstate = pstate;
+
+		/*
+		 * On refresh, the pstate's source text will be the refresh utility
+		 * statement. We need to fetch the the view definition to get the
+		 * query executed by the refresh.
+		 */
+		if (!is_create)
+		{
+			refresh_pstate = make_parsestate(NULL);
+			refresh_pstate->p_sourcetext = pg_get_viewdef_string(matviewOid, false);
+		}
 
 		dest = CreateTransientRelDestReceiver(OIDNewHeap);
-		processed = refresh_matview_datafill(dest, dataQuery, pstate,
-											 queryString, is_create);
+		processed = refresh_matview_datafill(dest, dataQuery, refresh_pstate, is_create);
 	}
 
 	/* Make the matview match the newly generated data. */
@@ -406,8 +416,7 @@ RefreshMatViewByOid(Oid matviewOid, bool is_create, bool skipData,
  */
 static uint64
 refresh_matview_datafill(DestReceiver *dest, Query *query,
-						 ParseState *pstate, const char *queryString,
-						 bool is_create)
+						 ParseState *pstate, bool is_create)
 {
 	List	   *rewritten;
 	PlannedStmt *plan;
@@ -415,6 +424,7 @@ refresh_matview_datafill(DestReceiver *dest, Query *query,
 	Query	   *copied_query;
 	uint64		processed;
 	JumbleState *jstate = NULL;
+	const char *queryString = pstate->p_sourcetext;
 
 	/* Lock and rewrite, using a copy to preserve the original query. */
 	copied_query = copyObject(query);
diff --git a/src/backend/utils/adt/ruleutils.c b/src/backend/utils/adt/ruleutils.c
index 2177d17e27..0bf85cbb75 100644
--- a/src/backend/utils/adt/ruleutils.c
+++ b/src/backend/utils/adt/ruleutils.c
@@ -776,6 +776,19 @@ pg_get_viewdef_name_ext(PG_FUNCTION_ARGS)
 	PG_RETURN_TEXT_P(string_to_text(res));
 }
 
+/*
+ * Internal version of pg_get_viewdef
+ */
+char *
+pg_get_viewdef_string(Oid viewoid, bool pretty)
+{
+	int			prettyFlags;
+
+	prettyFlags = GET_PRETTY_FLAGS(pretty);
+
+	return pg_get_viewdef_worker(viewoid, prettyFlags, WRAP_COLUMN_DEFAULT);
+}
+
 /*
  * Common code for by-OID and by-name variants of pg_get_viewdef
  */
diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 5300bc3baf..1f7a31b60b 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -1009,12 +1009,14 @@ SELECT pg_stat_statements_reset() IS NOT NULL AS t;
 REFRESH MATERIALIZED VIEW pgss_materialized_view;
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                       query                        
-----------+-------+----------------------------------------------------
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
  t        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view
- f        |     1 | REFRESH MATERIALIZED VIEW pgss_materialized_view;
+ f        |     1 | SELECT * FROM pg_catalog.pg_rewrite WHERE ev_class = $1 AND rulename = $2
+ f        |     1 | SELECT id                                                                +
+          |       |    FROM generate_series(1, 5) id(id);
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
-(3 rows)
+(4 rows)
 
 -- REFRESH MATERIALIZED VIEW, top-level tracking.
 SET pg_stat_statements.track = 'top';
-- 
2.45.2

From b52830c448a9e4625ba88e0f4ea868088faff640 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Thu, 3 Oct 2024 08:52:02 +0200
Subject: Track location to extract relevant part in nested statement

Previously, Query generated through transform would have unset
stmt_location. Extensions relying on the statement location to extract
the relevant part of the statement would fallback to use the whole
statement instead, thus showing the same string at the top and
nested level, which was a source of confusion.

This patch fixes the issue by keeping track of the statement locations
and propagate it to Query during transform, allowing pgss to only show
the relevant part of the query for nested query.
---
 .../expected/level_tracking.out               | 165 +++++++++---------
 .../pg_stat_statements/expected/planning.out  |  10 +-
 .../pg_stat_statements/expected/select.out    |   2 +-
 .../pg_stat_statements/expected/utility.out   |   2 +-
 contrib/pg_stat_statements/sql/planning.sql   |   4 +-
 src/backend/parser/analyze.c                  |  90 +++++++++-
 src/backend/parser/gram.y                     |  79 ++++++++-
 src/include/nodes/parsenodes.h                |  10 ++
 src/include/parser/parse_node.h               |   2 +
 9 files changed, 261 insertions(+), 103 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 8f008f8bfd1..489dc7143f7 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -206,37 +206,37 @@ EXPLAIN (COSTS OFF) SELECT 1 UNION SELECT 2;
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                query                                
-----------+-------+---------------------------------------------------------------------
+ toplevel | calls |                               query                                
+----------+-------+--------------------------------------------------------------------
+ f        |     1 | DELETE FROM stats_track_tab
  t        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2)
- f        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2);
  t        |     1 | EXPLAIN (COSTS OFF) (TABLE test_table)
- f        |     1 | EXPLAIN (COSTS OFF) (TABLE test_table);
  t        |     1 | EXPLAIN (COSTS OFF) (VALUES ($1, $2))
- f        |     1 | EXPLAIN (COSTS OFF) (VALUES ($1, $2));
  t        |     1 | EXPLAIN (COSTS OFF) DELETE FROM stats_track_tab
- f        |     1 | EXPLAIN (COSTS OFF) DELETE FROM stats_track_tab;
  t        |     1 | EXPLAIN (COSTS OFF) INSERT INTO stats_track_tab VALUES (($1))
- f        |     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)
- 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);
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1 UNION SELECT $2
- f        |     1 | EXPLAIN (COSTS OFF) SELECT $1 UNION SELECT $2;
- f        |     1 | EXPLAIN (COSTS OFF) SELECT $1;
  t        |     1 | EXPLAIN (COSTS OFF) TABLE stats_track_tab
- f        |     1 | EXPLAIN (COSTS OFF) TABLE stats_track_tab;
  t        |     1 | EXPLAIN (COSTS OFF) UPDATE stats_track_tab SET x = $1 WHERE x = $2
- f        |     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) VALUES ($1);
+ 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 UNION SELECT $2
+ f        |     1 | SELECT $1, $2
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | TABLE test_table
+ f        |     1 | UPDATE stats_track_tab SET x = $1 WHERE x = $2
+ f        |     1 | VALUES ($1)
+ f        |     1 | VALUES ($1, $2)
 (23 rows)
 
 -- EXPLAIN - top-level tracking.
@@ -405,20 +405,20 @@ EXPLAIN (COSTS OFF) SELECT 1, 2 UNION SELECT 3, 4\; EXPLAIN (COSTS OFF) (SELECT
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                                        query                                                        
-----------+-------+---------------------------------------------------------------------------------------------------------------------
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ f        |     1 | (SELECT $1, $2, $3) UNION SELECT $4, $5, $6
  t        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2, $3)
  t        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2, $3) UNION SELECT $4, $5, $6
- f        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2, $3); EXPLAIN (COSTS OFF) (SELECT 1, 2, 3, 4);
  t        |     1 | 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);
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1, $2
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1, $2 UNION SELECT $3, $4
- 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 | SELECT $1
+ f        |     1 | SELECT $1, $2
+ f        |     1 | SELECT $1, $2 UNION SELECT $3, $4
+ f        |     1 | SELECT $1, $2, $3
+ f        |     1 | SELECT $1, $2, $3, $4
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (13 rows)
 
@@ -494,29 +494,29 @@ EXPLAIN (COSTS OFF) INSERT INTO stats_track_tab VALUES ((1))\; EXPLAIN (COSTS OF
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                                              query                                                               
-----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ toplevel | calls |                               query                                
+----------+-------+--------------------------------------------------------------------
+ f        |     1 | DELETE FROM stats_track_tab
+ f        |     1 | DELETE FROM stats_track_tab WHERE x = $1
  t        |     1 | EXPLAIN (COSTS OFF) (TABLE test_table)
  t        |     1 | EXPLAIN (COSTS OFF) (VALUES ($1, $2))
  t        |     1 | EXPLAIN (COSTS OFF) DELETE FROM stats_track_tab
  t        |     1 | 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) DELETE FROM stats_track_tab; 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))
- 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);
  t        |     1 | EXPLAIN (COSTS OFF) TABLE stats_track_tab
- 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);
  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
- 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;
  t        |     1 | EXPLAIN (COSTS OFF) VALUES ($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 | INSERT INTO stats_track_tab VALUES ($1), ($2)
+ f        |     1 | INSERT INTO stats_track_tab VALUES (($1))
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | TABLE test_table
+ 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)
+ f        |     1 | VALUES ($1, $2)
 (21 rows)
 
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
@@ -547,21 +547,18 @@ EXPLAIN (COSTS OFF) MERGE INTO stats_track_tab
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                             query                                              
-----------+-------+------------------------------------------------------------------------------------------------
- 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                                                         +
+ toplevel | calls |                             query                             
+----------+-------+---------------------------------------------------------------
+ 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)
- 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;
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1, $2, $3, $4, $5
+ 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, $2, $3, $4, $5
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (5 rows)
 
@@ -789,29 +786,29 @@ EXPLAIN (COSTS OFF) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                           query                                           
-----------+-------+-------------------------------------------------------------------------------------------
+ toplevel | calls |                                          query                                           
+----------+-------+------------------------------------------------------------------------------------------
  t        |     1 | EXPLAIN (COSTS OFF) (WITH a AS (SELECT $1) (SELECT $2, $3))
- f        |     1 | EXPLAIN (COSTS OFF) (WITH a AS (SELECT $1) (SELECT $2, $3));
  t        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) DELETE FROM stats_track_tab
- f        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) DELETE FROM stats_track_tab;
  t        |     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) INSERT INTO stats_track_tab VALUES (($2));
- t        |     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                            +
-          |       |   WHEN MATCHED THEN UPDATE SET x = id                                                    +
+ t        |     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                           +
+          |       |   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) 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);
  t        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) SELECT $2
- f        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) SELECT $2;
  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) UPDATE stats_track_tab SET x = $2 WHERE x = $3;
  t        |     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 UNION SELECT $3;
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ 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 | WITH a AS (SELECT $1) SELECT $2
+ f        |     1 | WITH a AS (SELECT $1) UPDATE stats_track_tab SET x = $2 WHERE x = $3
+ f        |     1 | WITH a AS (select $1) SELECT $2 UNION SELECT $3
 (15 rows)
 
 -- EXPLAIN with CTEs - top-level tracking
@@ -921,12 +918,12 @@ EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                              query                               
-----------+-------+------------------------------------------------------------------
- t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)           +
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ 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
- f        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1;
+ f        |     1 | SELECT $1
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (4 rows)
 
@@ -1050,8 +1047,8 @@ 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 | PREPARE test_prepare_pgss AS select generate_series($1, $2)
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | select generate_series($1, $2)
 (4 rows)
 
 -- CREATE TABLE AS, top-level tracking.
@@ -1202,25 +1199,25 @@ COPY (DELETE FROM stats_track_tab WHERE x = 2 RETURNING x) TO stdout;
 2
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                    query                                    
-----------+-------+-----------------------------------------------------------------------------
- f        |     1 | COPY (DELETE FROM stats_track_tab WHERE x = $1 RETURNING x) TO stdout
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
  t        |     1 | COPY (DELETE FROM stats_track_tab WHERE x = 2 RETURNING x) TO stdout
- f        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES ($1) RETURNING x) TO stdout
  t        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) TO stdout
- f        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT $1 id) ON x = id            +
-          |       |   WHEN MATCHED THEN UPDATE SET x = id                                      +
+ t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id           +
+          |       |   WHEN MATCHED THEN UPDATE SET x = id                                    +
           |       |   WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) TO stdout
- t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id             +
-          |       |   WHEN MATCHED THEN UPDATE SET x = id                                      +
-          |       |   WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) TO stdout
- f        |     1 | COPY (SELECT $1 UNION SELECT $2) TO stdout
- f        |     1 | COPY (SELECT $1) TO stdout
  t        |     1 | COPY (SELECT 1 UNION SELECT 2) TO stdout
  t        |     1 | COPY (SELECT 1) TO stdout
- f        |     1 | COPY (UPDATE stats_track_tab SET x = $1 WHERE x = $2 RETURNING x) TO stdout
  t        |     1 | COPY (UPDATE stats_track_tab SET x = 2 WHERE x = 1 RETURNING x) TO stdout
+ f        |     1 | DELETE FROM stats_track_tab WHERE x = $1 RETURNING x
+ f        |     1 | INSERT INTO stats_track_tab (x) VALUES ($1) RETURNING x
+ f        |     1 | MERGE INTO stats_track_tab USING (SELECT $1 id) ON x = id                +
+          |       |   WHEN MATCHED THEN UPDATE SET x = id                                    +
+          |       |   WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x
+ f        |     1 | SELECT $1
+ f        |     1 | SELECT $1 UNION SELECT $2
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | UPDATE stats_track_tab SET x = $1 WHERE x = $2 RETURNING x
 (13 rows)
 
 -- COPY - top-level tracking.
diff --git a/contrib/pg_stat_statements/expected/planning.out b/contrib/pg_stat_statements/expected/planning.out
index 9effd11fdc8..3ee1928cbe9 100644
--- a/contrib/pg_stat_statements/expected/planning.out
+++ b/contrib/pg_stat_statements/expected/planning.out
@@ -58,7 +58,7 @@ SELECT 42;
 (1 row)
 
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
  plans | calls | rows |                          query                           
 -------+-------+------+----------------------------------------------------------
      0 |     1 |    0 | ALTER TABLE stats_plan_test ADD COLUMN x int
@@ -72,10 +72,10 @@ SELECT plans, calls, rows, query FROM pg_stat_statements
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
- plans_ok | calls | rows |                         query                         
-----------+-------+------+-------------------------------------------------------
- t        |     4 |    4 | PREPARE prep1 AS SELECT COUNT(*) FROM stats_plan_test
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
+ plans_ok | calls | rows |                query                 
+----------+-------+------+--------------------------------------
+ t        |     4 |    4 | SELECT COUNT(*) FROM stats_plan_test
 (1 row)
 
 -- Cleanup
diff --git a/contrib/pg_stat_statements/expected/select.out b/contrib/pg_stat_statements/expected/select.out
index dd6c756f67d..217a2c0b2bc 100644
--- a/contrib/pg_stat_statements/expected/select.out
+++ b/contrib/pg_stat_statements/expected/select.out
@@ -127,7 +127,6 @@ DEALLOCATE pgss_test;
 SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
  calls | rows |                                    query                                     
 -------+------+------------------------------------------------------------------------------
-     1 |    1 | PREPARE pgss_test (int) AS SELECT $1, $2 LIMIT $3
      4 |    4 | SELECT $1                                                                   +
        |      |   -- multiline                                                              +
        |      |   AS "text"
@@ -137,6 +136,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
      2 |    2 | SELECT $1 AS "int"
      1 |    2 | SELECT $1 AS i UNION SELECT $2 ORDER BY i
      1 |    1 | SELECT $1 || $2
+     1 |    1 | SELECT $1, $2 LIMIT $3
      0 |    0 | SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C"
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
      1 |    2 | WITH t(f) AS (                                                              +
diff --git a/contrib/pg_stat_statements/expected/utility.out b/contrib/pg_stat_statements/expected/utility.out
index 060d4416dd7..aa4f0f7e628 100644
--- a/contrib/pg_stat_statements/expected/utility.out
+++ b/contrib/pg_stat_statements/expected/utility.out
@@ -540,7 +540,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
 -------+------+----------------------------------------------------
      2 |    0 | DEALLOCATE $1
      2 |    0 | DEALLOCATE ALL
-     2 |    2 | PREPARE stat_select AS SELECT $1 AS a
+     2 |    2 | SELECT $1 AS a
      1 |    1 | SELECT $1 as a
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (5 rows)
diff --git a/contrib/pg_stat_statements/sql/planning.sql b/contrib/pg_stat_statements/sql/planning.sql
index 46f5d9b951c..9cfe206b3b0 100644
--- a/contrib/pg_stat_statements/sql/planning.sql
+++ b/contrib/pg_stat_statements/sql/planning.sql
@@ -20,11 +20,11 @@ SELECT 42;
 SELECT 42;
 SELECT 42;
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 
 -- Cleanup
 DROP TABLE stats_plan_test;
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index 8a6ba1692e8..dc716936195 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -238,6 +238,74 @@ parse_sub_analyze(Node *parseTree, ParseState *parentParseState,
 	return query;
 }
 
+/*
+ * setQueryLocationAndLength
+ * 		Set query's location and length from statement and ParseState
+ *
+ * Some statements, like PreparableStmt, can be located within parentheses.
+ * For example "(SELECT 1)" or "COPY (UPDATE ...) to x;". For those, we can't
+ * use the whole string from the statement's location or the SQL string will
+ * yield "SELECT 1)". The parser will set stmt_len, reflecting the size of the
+ * statement within the parentheses. Thus, when stmt_len is available, we need
+ * to use it for the Query's stmt_len.
+ *
+ * For other cases, the parser can't provide the length of individual statements.
+ * However, we have the statement's location plus the length (p_stmt_len) and
+ * location (p_stmt_location) of the top level RawStmt, stored in pstate. Thus,
+ * the statement's length is the RawStmt's length minus how much we've advanced
+ * in the RawStmt's string.
+ */
+static void
+setQueryLocationAndLength(ParseState *pstate, Query *qry, Node *parseTree)
+{
+	ParseLoc	stmt_len = 0;
+
+	switch (nodeTag(parseTree))
+	{
+		case T_InsertStmt:
+			qry->stmt_location = ((InsertStmt *) parseTree)->stmt_location;
+			stmt_len = ((InsertStmt *) parseTree)->stmt_len;
+			break;
+		case T_DeleteStmt:
+			qry->stmt_location = ((DeleteStmt *) parseTree)->stmt_location;
+			stmt_len = ((DeleteStmt *) parseTree)->stmt_len;
+			break;
+
+		case T_UpdateStmt:
+			qry->stmt_location = ((UpdateStmt *) parseTree)->stmt_location;
+			stmt_len = ((UpdateStmt *) parseTree)->stmt_len;
+			break;
+
+		case T_MergeStmt:
+			qry->stmt_location = ((MergeStmt *) parseTree)->stmt_location;
+			stmt_len = ((MergeStmt *) parseTree)->stmt_len;
+			break;
+
+		case T_SelectStmt:
+			qry->stmt_location = ((SelectStmt *) parseTree)->stmt_location;
+			stmt_len = ((SelectStmt *) parseTree)->stmt_len;
+			break;
+
+		case T_PLAssignStmt:
+			qry->stmt_location = ((PLAssignStmt *) parseTree)->location;
+			break;
+
+		default:
+			qry->stmt_location = pstate->p_stmt_location;
+			break;
+	}
+	if (stmt_len > 0)
+		/* Statement's length is known, use it */
+		qry->stmt_len = stmt_len;
+	else
+
+		/*
+		 * Compute the statement's length from statement's location and
+		 * RawStmt's length and location
+		 */
+		qry->stmt_len = pstate->p_stmt_len - (qry->stmt_location - pstate->p_stmt_location);
+}
+
 /*
  * transformTopLevelStmt -
  *	  transform a Parse tree into a Query tree.
@@ -250,12 +318,13 @@ transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree)
 {
 	Query	   *result;
 
+	/* Store RawStmt's length and location in pstate */
+	pstate->p_stmt_len = parseTree->stmt_len;
+	pstate->p_stmt_location = parseTree->stmt_location;
+
 	/* We're at top level, so allow SELECT INTO */
 	result = transformOptionalSelectInto(pstate, parseTree->stmt);
 
-	result->stmt_location = parseTree->stmt_location;
-	result->stmt_len = parseTree->stmt_len;
-
 	return result;
 }
 
@@ -417,13 +486,14 @@ transformStmt(ParseState *pstate, Node *parseTree)
 			 */
 			result = makeNode(Query);
 			result->commandType = CMD_UTILITY;
-			result->utilityStmt = (Node *) parseTree;
+			result->utilityStmt = parseTree;
 			break;
 	}
 
 	/* Mark as original query until we learn differently */
 	result->querySource = QSRC_ORIGINAL;
 	result->canSetTag = true;
+	setQueryLocationAndLength(pstate, result, parseTree);
 
 	return result;
 }
@@ -2397,6 +2467,8 @@ transformReturnStmt(ParseState *pstate, ReturnStmt *stmt)
 	Query	   *qry = makeNode(Query);
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = pstate->p_stmt_location;
+	qry->stmt_len = pstate->p_stmt_len;
 	qry->isReturn = true;
 
 	qry->targetList = list_make1(makeTargetEntry((Expr *) transformExpr(pstate, stmt->returnval, EXPR_KIND_SELECT_TARGET),
@@ -2677,6 +2749,8 @@ transformPLAssignStmt(ParseState *pstate, PLAssignStmt *stmt)
 	 * consider WITH or INTO, and we build a targetlist our own way.
 	 */
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = pstate->p_stmt_location;
+	qry->stmt_len = pstate->p_stmt_len;
 	pstate->p_is_insert = false;
 
 	/* make FOR UPDATE/FOR SHARE info available to addRangeTableEntry */
@@ -2948,6 +3022,8 @@ transformDeclareCursorStmt(ParseState *pstate, DeclareCursorStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3003,6 +3079,8 @@ transformExplainStmt(ParseState *pstate, ExplainStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3083,6 +3161,8 @@ transformCreateTableAsStmt(ParseState *pstate, CreateTableAsStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3207,6 +3287,8 @@ transformCallStmt(ParseState *pstate, CallStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 4aa8646af7b..d4bae0d8a36 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -167,6 +167,7 @@ static void base_yyerror(YYLTYPE *yylloc, core_yyscan_t yyscanner,
 						 const char *msg);
 static RawStmt *makeRawStmt(Node *stmt, int stmt_location);
 static void updateRawStmtEnd(RawStmt *rs, int end_location);
+static void updatePreparableStmtEnd(Node *n, int end_location);
 static Node *makeColumnRef(char *colname, List *indirection,
 						   int location, core_yyscan_t yyscanner);
 static Node *makeTypeCast(Node *arg, TypeName *typename, int location);
@@ -190,7 +191,7 @@ static void insertSelectOptions(SelectStmt *stmt,
 								SelectLimit *limitClause,
 								WithClause *withClause,
 								core_yyscan_t yyscanner);
-static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg);
+static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location);
 static Node *doNegate(Node *n, int location);
 static void doNegateFloat(Float *v);
 static Node *makeAndExpr(Node *lexpr, Node *rexpr, int location);
@@ -3403,6 +3404,7 @@ CopyStmt:	COPY opt_binary qualified_name opt_column_list
 				{
 					CopyStmt *n = makeNode(CopyStmt);
 
+					updatePreparableStmtEnd($3, @4);
 					n->relation = NULL;
 					n->query = $3;
 					n->attlist = NIL;
@@ -12170,6 +12172,9 @@ InsertStmt:
 					$5->onConflictClause = $6;
 					$5->returningList = $7;
 					$5->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					$5->stmt_location = @$;
 					$$ = (Node *) $5;
 				}
 		;
@@ -12323,6 +12328,9 @@ DeleteStmt: opt_with_clause DELETE_P FROM relation_expr_opt_alias
 					n->whereClause = $6;
 					n->returningList = $7;
 					n->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					n->stmt_location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12397,6 +12405,9 @@ UpdateStmt: opt_with_clause UPDATE relation_expr_opt_alias
 					n->whereClause = $7;
 					n->returningList = $8;
 					n->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					n->stmt_location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12474,6 +12485,9 @@ MergeStmt:
 					m->joinCondition = $8;
 					m->mergeWhenClauses = $9;
 					m->returningList = $10;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					m->stmt_location = @$;
 
 					$$ = (Node *) m;
 				}
@@ -12714,7 +12728,20 @@ SelectStmt: select_no_parens			%prec UMINUS
 		;
 
 select_with_parens:
-			'(' select_no_parens ')'				{ $$ = $2; }
+			'(' select_no_parens ')'
+				{
+					SelectStmt *n = (SelectStmt *) $2;
+					/*
+					 * As SelectStmt's location starts at the SELECT keyword,
+					 * we need to track the length of the SelectStmt within
+					 * parentheses to be able to extract the relevant part
+					 * of the query. Without this, the RawStmt's length will
+					 * be used and will include the closing parenthesis,
+					 * "SELECT 1)".
+					 */
+					n->stmt_len = @3 - @2;
+					$$ = $2;
+				}
 			| '(' select_with_parens ')'			{ $$ = $2; }
 		;
 
@@ -12836,6 +12863,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| SELECT distinct_clause target_list
@@ -12853,6 +12881,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| values_clause							{ $$ = $1; }
@@ -12873,19 +12902,20 @@ simple_select:
 
 					n->targetList = list_make1(rt);
 					n->fromClause = list_make1($2);
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| select_clause UNION set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause INTERSECT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause EXCEPT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 		;
 
@@ -13443,6 +13473,7 @@ values_clause:
 				{
 					SelectStmt *n = makeNode(SelectStmt);
 
+					n->stmt_location = @1;
 					n->valuesLists = list_make1($3);
 					$$ = (Node *) n;
 				}
@@ -18585,6 +18616,39 @@ updateRawStmtEnd(RawStmt *rs, int end_location)
 	rs->stmt_len = end_location - rs->stmt_location;
 }
 
+/* Adjust a PreparableStmt to reflect that it doesn't run to the end of the string */
+static void
+updatePreparableStmtEnd(Node *n, int end_location)
+{
+	if (IsA(n, SelectStmt))
+	{
+		SelectStmt *stmt = (SelectStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, InsertStmt))
+	{
+		InsertStmt *stmt = (InsertStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, UpdateStmt))
+	{
+		UpdateStmt *stmt = (UpdateStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, DeleteStmt))
+	{
+		DeleteStmt *stmt = (DeleteStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, MergeStmt))
+	{
+		MergeStmt *stmt = (MergeStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else
+		elog(ERROR, "unexpected node type %d", (int) n->type);
+}
+
 static Node *
 makeColumnRef(char *colname, List *indirection,
 			  int location, core_yyscan_t yyscanner)
@@ -18963,11 +19027,13 @@ insertSelectOptions(SelectStmt *stmt,
 					 errmsg("multiple WITH clauses not allowed"),
 					 parser_errposition(exprLocation((Node *) withClause))));
 		stmt->withClause = withClause;
+		/* Update SelectStmt's location to the start of the with clause */
+		stmt->stmt_location = withClause->location;
 	}
 }
 
 static Node *
-makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
+makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location)
 {
 	SelectStmt *n = makeNode(SelectStmt);
 
@@ -18975,6 +19041,7 @@ makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
 	n->all = all;
 	n->larg = (SelectStmt *) larg;
 	n->rarg = (SelectStmt *) rarg;
+	n->stmt_location = location;
 	return (Node *) n;
 }
 
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index c92cef3d16d..b40b661ec8a 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -2046,6 +2046,8 @@ typedef struct InsertStmt
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
 	OverridingKind override;	/* OVERRIDING clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } InsertStmt;
 
 /* ----------------------
@@ -2060,6 +2062,8 @@ typedef struct DeleteStmt
 	Node	   *whereClause;	/* qualifications */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } DeleteStmt;
 
 /* ----------------------
@@ -2075,6 +2079,8 @@ typedef struct UpdateStmt
 	List	   *fromClause;		/* optional from clause for more tables */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } UpdateStmt;
 
 /* ----------------------
@@ -2090,6 +2096,8 @@ typedef struct MergeStmt
 	List	   *mergeWhenClauses;	/* list of MergeWhenClause(es) */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } MergeStmt;
 
 /* ----------------------
@@ -2159,6 +2167,8 @@ typedef struct SelectStmt
 	bool		all;			/* ALL specified? */
 	struct SelectStmt *larg;	/* left child */
 	struct SelectStmt *rarg;	/* right child */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 	/* Eventually add fields for CORRESPONDING spec here */
 } SelectStmt;
 
diff --git a/src/include/parser/parse_node.h b/src/include/parser/parse_node.h
index 543df568147..ba572b3aea2 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 */
-- 
2.39.5 (Apple Git-154)

From f090f3dc2c87bce2f78504fc925bcf1d69a45a1b Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
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,
and DeclareCursor before it is executed.
Additionally, explain's nested query can itself be a CreateTableAs,
DeclareCursor which also needs to be jumbled. This is now done in
ExplainOneUtility.
---
 .../expected/level_tracking.out               | 12 ++++--
 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/include/commands/explain.h                |  4 +-
 src/include/commands/prepare.h                |  4 +-
 src/test/regress/expected/explain.out         | 17 ++++++++
 src/test/regress/sql/explain.sql              |  4 ++
 9 files changed, 89 insertions(+), 35 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 489dc7143f7..9aee9f5010e 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';
diff --git a/src/backend/commands/createas.c b/src/backend/commands/createas.c
index 68ec122dbf9..e4a627ad91d 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 18a5af6b919..c81221cdbeb 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 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/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/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 dcbdaa03885..d2eef8097cf 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 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)

From 2dc0834a6d8cee9a96a120743619f7defc8f24c9 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Wed, 23 Oct 2024 13:07:52 +0900
Subject: [PATCH v12 1/2] Track location to extract relevant part in nested
 statement

Previously, Query generated through transform would have unset
stmt_location. Extensions relying on the statement location to extract
the relevant part of the statement would fallback to use the whole
statement instead, thus showing the same string at the top and
nested level, which was a source of confusion.

This patch fixes the issue by keeping track of the statement locations
and propagate it to Query during transform, allowing pgss to only show
the relevant part of the query for nested query.
---
 src/include/nodes/parsenodes.h                |  10 ++
 src/include/parser/parse_node.h               |   2 +
 src/backend/parser/analyze.c                  |  90 +++++++++-
 src/backend/parser/gram.y                     |  79 ++++++++-
 .../expected/level_tracking.out               | 165 +++++++++---------
 .../pg_stat_statements/expected/planning.out  |  10 +-
 .../pg_stat_statements/expected/select.out    |   2 +-
 .../pg_stat_statements/expected/utility.out   |   2 +-
 contrib/pg_stat_statements/sql/planning.sql   |   4 +-
 9 files changed, 261 insertions(+), 103 deletions(-)

diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index c92cef3d16..b40b661ec8 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -2046,6 +2046,8 @@ typedef struct InsertStmt
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
 	OverridingKind override;	/* OVERRIDING clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } InsertStmt;
 
 /* ----------------------
@@ -2060,6 +2062,8 @@ typedef struct DeleteStmt
 	Node	   *whereClause;	/* qualifications */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } DeleteStmt;
 
 /* ----------------------
@@ -2075,6 +2079,8 @@ typedef struct UpdateStmt
 	List	   *fromClause;		/* optional from clause for more tables */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } UpdateStmt;
 
 /* ----------------------
@@ -2090,6 +2096,8 @@ typedef struct MergeStmt
 	List	   *mergeWhenClauses;	/* list of MergeWhenClause(es) */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } MergeStmt;
 
 /* ----------------------
@@ -2159,6 +2167,8 @@ typedef struct SelectStmt
 	bool		all;			/* ALL specified? */
 	struct SelectStmt *larg;	/* left child */
 	struct SelectStmt *rarg;	/* right child */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 	/* Eventually add fields for CORRESPONDING spec here */
 } SelectStmt;
 
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/backend/parser/analyze.c b/src/backend/parser/analyze.c
index 8a6ba1692e..dc71693619 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -238,6 +238,74 @@ parse_sub_analyze(Node *parseTree, ParseState *parentParseState,
 	return query;
 }
 
+/*
+ * setQueryLocationAndLength
+ * 		Set query's location and length from statement and ParseState
+ *
+ * Some statements, like PreparableStmt, can be located within parentheses.
+ * For example "(SELECT 1)" or "COPY (UPDATE ...) to x;". For those, we can't
+ * use the whole string from the statement's location or the SQL string will
+ * yield "SELECT 1)". The parser will set stmt_len, reflecting the size of the
+ * statement within the parentheses. Thus, when stmt_len is available, we need
+ * to use it for the Query's stmt_len.
+ *
+ * For other cases, the parser can't provide the length of individual statements.
+ * However, we have the statement's location plus the length (p_stmt_len) and
+ * location (p_stmt_location) of the top level RawStmt, stored in pstate. Thus,
+ * the statement's length is the RawStmt's length minus how much we've advanced
+ * in the RawStmt's string.
+ */
+static void
+setQueryLocationAndLength(ParseState *pstate, Query *qry, Node *parseTree)
+{
+	ParseLoc	stmt_len = 0;
+
+	switch (nodeTag(parseTree))
+	{
+		case T_InsertStmt:
+			qry->stmt_location = ((InsertStmt *) parseTree)->stmt_location;
+			stmt_len = ((InsertStmt *) parseTree)->stmt_len;
+			break;
+		case T_DeleteStmt:
+			qry->stmt_location = ((DeleteStmt *) parseTree)->stmt_location;
+			stmt_len = ((DeleteStmt *) parseTree)->stmt_len;
+			break;
+
+		case T_UpdateStmt:
+			qry->stmt_location = ((UpdateStmt *) parseTree)->stmt_location;
+			stmt_len = ((UpdateStmt *) parseTree)->stmt_len;
+			break;
+
+		case T_MergeStmt:
+			qry->stmt_location = ((MergeStmt *) parseTree)->stmt_location;
+			stmt_len = ((MergeStmt *) parseTree)->stmt_len;
+			break;
+
+		case T_SelectStmt:
+			qry->stmt_location = ((SelectStmt *) parseTree)->stmt_location;
+			stmt_len = ((SelectStmt *) parseTree)->stmt_len;
+			break;
+
+		case T_PLAssignStmt:
+			qry->stmt_location = ((PLAssignStmt *) parseTree)->location;
+			break;
+
+		default:
+			qry->stmt_location = pstate->p_stmt_location;
+			break;
+	}
+	if (stmt_len > 0)
+		/* Statement's length is known, use it */
+		qry->stmt_len = stmt_len;
+	else
+
+		/*
+		 * Compute the statement's length from statement's location and
+		 * RawStmt's length and location
+		 */
+		qry->stmt_len = pstate->p_stmt_len - (qry->stmt_location - pstate->p_stmt_location);
+}
+
 /*
  * transformTopLevelStmt -
  *	  transform a Parse tree into a Query tree.
@@ -250,12 +318,13 @@ transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree)
 {
 	Query	   *result;
 
+	/* Store RawStmt's length and location in pstate */
+	pstate->p_stmt_len = parseTree->stmt_len;
+	pstate->p_stmt_location = parseTree->stmt_location;
+
 	/* We're at top level, so allow SELECT INTO */
 	result = transformOptionalSelectInto(pstate, parseTree->stmt);
 
-	result->stmt_location = parseTree->stmt_location;
-	result->stmt_len = parseTree->stmt_len;
-
 	return result;
 }
 
@@ -417,13 +486,14 @@ transformStmt(ParseState *pstate, Node *parseTree)
 			 */
 			result = makeNode(Query);
 			result->commandType = CMD_UTILITY;
-			result->utilityStmt = (Node *) parseTree;
+			result->utilityStmt = parseTree;
 			break;
 	}
 
 	/* Mark as original query until we learn differently */
 	result->querySource = QSRC_ORIGINAL;
 	result->canSetTag = true;
+	setQueryLocationAndLength(pstate, result, parseTree);
 
 	return result;
 }
@@ -2397,6 +2467,8 @@ transformReturnStmt(ParseState *pstate, ReturnStmt *stmt)
 	Query	   *qry = makeNode(Query);
 
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = pstate->p_stmt_location;
+	qry->stmt_len = pstate->p_stmt_len;
 	qry->isReturn = true;
 
 	qry->targetList = list_make1(makeTargetEntry((Expr *) transformExpr(pstate, stmt->returnval, EXPR_KIND_SELECT_TARGET),
@@ -2677,6 +2749,8 @@ transformPLAssignStmt(ParseState *pstate, PLAssignStmt *stmt)
 	 * consider WITH or INTO, and we build a targetlist our own way.
 	 */
 	qry->commandType = CMD_SELECT;
+	qry->stmt_location = pstate->p_stmt_location;
+	qry->stmt_len = pstate->p_stmt_len;
 	pstate->p_is_insert = false;
 
 	/* make FOR UPDATE/FOR SHARE info available to addRangeTableEntry */
@@ -2948,6 +3022,8 @@ transformDeclareCursorStmt(ParseState *pstate, DeclareCursorStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3003,6 +3079,8 @@ transformExplainStmt(ParseState *pstate, ExplainStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3083,6 +3161,8 @@ transformCreateTableAsStmt(ParseState *pstate, CreateTableAsStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
@@ -3207,6 +3287,8 @@ transformCallStmt(ParseState *pstate, CallStmt *stmt)
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
+	result->stmt_location = pstate->p_stmt_location;
+	result->stmt_len = pstate->p_stmt_len;
 	result->utilityStmt = (Node *) stmt;
 
 	return result;
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 4bab2117d9..172a26e10e 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -153,6 +153,7 @@ static void base_yyerror(YYLTYPE *yylloc, core_yyscan_t yyscanner,
 						 const char *msg);
 static RawStmt *makeRawStmt(Node *stmt, int stmt_location);
 static void updateRawStmtEnd(RawStmt *rs, int end_location);
+static void updatePreparableStmtEnd(Node *n, int end_location);
 static Node *makeColumnRef(char *colname, List *indirection,
 						   int location, core_yyscan_t yyscanner);
 static Node *makeTypeCast(Node *arg, TypeName *typename, int location);
@@ -176,7 +177,7 @@ static void insertSelectOptions(SelectStmt *stmt,
 								SelectLimit *limitClause,
 								WithClause *withClause,
 								core_yyscan_t yyscanner);
-static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg);
+static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location);
 static Node *doNegate(Node *n, int location);
 static void doNegateFloat(Float *v);
 static Node *makeAndExpr(Node *lexpr, Node *rexpr, int location);
@@ -3383,6 +3384,7 @@ CopyStmt:	COPY opt_binary qualified_name opt_column_list
 				{
 					CopyStmt *n = makeNode(CopyStmt);
 
+					updatePreparableStmtEnd($3, @4);
 					n->relation = NULL;
 					n->query = $3;
 					n->attlist = NIL;
@@ -12150,6 +12152,9 @@ InsertStmt:
 					$5->onConflictClause = $6;
 					$5->returningList = $7;
 					$5->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					$5->stmt_location = @$;
 					$$ = (Node *) $5;
 				}
 		;
@@ -12303,6 +12308,9 @@ DeleteStmt: opt_with_clause DELETE_P FROM relation_expr_opt_alias
 					n->whereClause = $6;
 					n->returningList = $7;
 					n->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					n->stmt_location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12377,6 +12385,9 @@ UpdateStmt: opt_with_clause UPDATE relation_expr_opt_alias
 					n->whereClause = $7;
 					n->returningList = $8;
 					n->withClause = $1;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					n->stmt_location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12454,6 +12465,9 @@ MergeStmt:
 					m->joinCondition = $8;
 					m->mergeWhenClauses = $9;
 					m->returningList = $10;
+					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
+						@$ = @2;
+					m->stmt_location = @$;
 
 					$$ = (Node *) m;
 				}
@@ -12694,7 +12708,20 @@ SelectStmt: select_no_parens			%prec UMINUS
 		;
 
 select_with_parens:
-			'(' select_no_parens ')'				{ $$ = $2; }
+			'(' select_no_parens ')'
+				{
+					SelectStmt *n = (SelectStmt *) $2;
+					/*
+					 * As SelectStmt's location starts at the SELECT keyword,
+					 * we need to track the length of the SelectStmt within
+					 * parentheses to be able to extract the relevant part
+					 * of the query. Without this, the RawStmt's length will
+					 * be used and will include the closing parenthesis,
+					 * "SELECT 1)".
+					 */
+					n->stmt_len = @3 - @2;
+					$$ = $2;
+				}
 			| '(' select_with_parens ')'			{ $$ = $2; }
 		;
 
@@ -12816,6 +12843,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| SELECT distinct_clause target_list
@@ -12833,6 +12861,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| values_clause							{ $$ = $1; }
@@ -12853,19 +12882,20 @@ simple_select:
 
 					n->targetList = list_make1(rt);
 					n->fromClause = list_make1($2);
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| select_clause UNION set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause INTERSECT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause EXCEPT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 		;
 
@@ -13423,6 +13453,7 @@ values_clause:
 				{
 					SelectStmt *n = makeNode(SelectStmt);
 
+					n->stmt_location = @1;
 					n->valuesLists = list_make1($3);
 					$$ = (Node *) n;
 				}
@@ -18565,6 +18596,39 @@ updateRawStmtEnd(RawStmt *rs, int end_location)
 	rs->stmt_len = end_location - rs->stmt_location;
 }
 
+/* Adjust a PreparableStmt to reflect that it doesn't run to the end of the string */
+static void
+updatePreparableStmtEnd(Node *n, int end_location)
+{
+	if (IsA(n, SelectStmt))
+	{
+		SelectStmt *stmt = (SelectStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, InsertStmt))
+	{
+		InsertStmt *stmt = (InsertStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, UpdateStmt))
+	{
+		UpdateStmt *stmt = (UpdateStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, DeleteStmt))
+	{
+		DeleteStmt *stmt = (DeleteStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, MergeStmt))
+	{
+		MergeStmt *stmt = (MergeStmt *)n;
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else
+		elog(ERROR, "unexpected node type %d", (int) n->type);
+}
+
 static Node *
 makeColumnRef(char *colname, List *indirection,
 			  int location, core_yyscan_t yyscanner)
@@ -18943,11 +19007,13 @@ insertSelectOptions(SelectStmt *stmt,
 					 errmsg("multiple WITH clauses not allowed"),
 					 parser_errposition(exprLocation((Node *) withClause))));
 		stmt->withClause = withClause;
+		/* Update SelectStmt's location to the start of the with clause */
+		stmt->stmt_location = withClause->location;
 	}
 }
 
 static Node *
-makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
+makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location)
 {
 	SelectStmt *n = makeNode(SelectStmt);
 
@@ -18955,6 +19021,7 @@ makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
 	n->all = all;
 	n->larg = (SelectStmt *) larg;
 	n->rarg = (SelectStmt *) rarg;
+	n->stmt_location = location;
 	return (Node *) n;
 }
 
diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 8f008f8bfd..489dc7143f 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -206,37 +206,37 @@ EXPLAIN (COSTS OFF) SELECT 1 UNION SELECT 2;
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                query                                
-----------+-------+---------------------------------------------------------------------
+ toplevel | calls |                               query                                
+----------+-------+--------------------------------------------------------------------
+ f        |     1 | DELETE FROM stats_track_tab
  t        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2)
- f        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2);
  t        |     1 | EXPLAIN (COSTS OFF) (TABLE test_table)
- f        |     1 | EXPLAIN (COSTS OFF) (TABLE test_table);
  t        |     1 | EXPLAIN (COSTS OFF) (VALUES ($1, $2))
- f        |     1 | EXPLAIN (COSTS OFF) (VALUES ($1, $2));
  t        |     1 | EXPLAIN (COSTS OFF) DELETE FROM stats_track_tab
- f        |     1 | EXPLAIN (COSTS OFF) DELETE FROM stats_track_tab;
  t        |     1 | EXPLAIN (COSTS OFF) INSERT INTO stats_track_tab VALUES (($1))
- f        |     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)
- 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);
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1 UNION SELECT $2
- f        |     1 | EXPLAIN (COSTS OFF) SELECT $1 UNION SELECT $2;
- f        |     1 | EXPLAIN (COSTS OFF) SELECT $1;
  t        |     1 | EXPLAIN (COSTS OFF) TABLE stats_track_tab
- f        |     1 | EXPLAIN (COSTS OFF) TABLE stats_track_tab;
  t        |     1 | EXPLAIN (COSTS OFF) UPDATE stats_track_tab SET x = $1 WHERE x = $2
- f        |     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) VALUES ($1);
+ 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 UNION SELECT $2
+ f        |     1 | SELECT $1, $2
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | TABLE test_table
+ f        |     1 | UPDATE stats_track_tab SET x = $1 WHERE x = $2
+ f        |     1 | VALUES ($1)
+ f        |     1 | VALUES ($1, $2)
 (23 rows)
 
 -- EXPLAIN - top-level tracking.
@@ -405,20 +405,20 @@ EXPLAIN (COSTS OFF) SELECT 1, 2 UNION SELECT 3, 4\; EXPLAIN (COSTS OFF) (SELECT
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                                        query                                                        
-----------+-------+---------------------------------------------------------------------------------------------------------------------
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ f        |     1 | (SELECT $1, $2, $3) UNION SELECT $4, $5, $6
  t        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2, $3)
  t        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2, $3) UNION SELECT $4, $5, $6
- f        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2, $3); EXPLAIN (COSTS OFF) (SELECT 1, 2, 3, 4);
  t        |     1 | 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);
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1, $2
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1, $2 UNION SELECT $3, $4
- 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 | SELECT $1
+ f        |     1 | SELECT $1, $2
+ f        |     1 | SELECT $1, $2 UNION SELECT $3, $4
+ f        |     1 | SELECT $1, $2, $3
+ f        |     1 | SELECT $1, $2, $3, $4
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (13 rows)
 
@@ -494,29 +494,29 @@ EXPLAIN (COSTS OFF) INSERT INTO stats_track_tab VALUES ((1))\; EXPLAIN (COSTS OF
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                                              query                                                               
-----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ toplevel | calls |                               query                                
+----------+-------+--------------------------------------------------------------------
+ f        |     1 | DELETE FROM stats_track_tab
+ f        |     1 | DELETE FROM stats_track_tab WHERE x = $1
  t        |     1 | EXPLAIN (COSTS OFF) (TABLE test_table)
  t        |     1 | EXPLAIN (COSTS OFF) (VALUES ($1, $2))
  t        |     1 | EXPLAIN (COSTS OFF) DELETE FROM stats_track_tab
  t        |     1 | 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) DELETE FROM stats_track_tab; 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))
- 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);
  t        |     1 | EXPLAIN (COSTS OFF) TABLE stats_track_tab
- 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);
  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
- 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;
  t        |     1 | EXPLAIN (COSTS OFF) VALUES ($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 | INSERT INTO stats_track_tab VALUES ($1), ($2)
+ f        |     1 | INSERT INTO stats_track_tab VALUES (($1))
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | TABLE test_table
+ 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)
+ f        |     1 | VALUES ($1, $2)
 (21 rows)
 
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
@@ -547,21 +547,18 @@ EXPLAIN (COSTS OFF) MERGE INTO stats_track_tab
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                             query                                              
-----------+-------+------------------------------------------------------------------------------------------------
- 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                                                         +
+ toplevel | calls |                             query                             
+----------+-------+---------------------------------------------------------------
+ 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)
- 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;
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1, $2, $3, $4, $5
+ 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, $2, $3, $4, $5
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (5 rows)
 
@@ -789,29 +786,29 @@ EXPLAIN (COSTS OFF) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                           query                                           
-----------+-------+-------------------------------------------------------------------------------------------
+ toplevel | calls |                                          query                                           
+----------+-------+------------------------------------------------------------------------------------------
  t        |     1 | EXPLAIN (COSTS OFF) (WITH a AS (SELECT $1) (SELECT $2, $3))
- f        |     1 | EXPLAIN (COSTS OFF) (WITH a AS (SELECT $1) (SELECT $2, $3));
  t        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) DELETE FROM stats_track_tab
- f        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) DELETE FROM stats_track_tab;
  t        |     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) INSERT INTO stats_track_tab VALUES (($2));
- t        |     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                            +
-          |       |   WHEN MATCHED THEN UPDATE SET x = id                                                    +
+ t        |     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                           +
+          |       |   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) 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);
  t        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) SELECT $2
- f        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) SELECT $2;
  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) UPDATE stats_track_tab SET x = $2 WHERE x = $3;
  t        |     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 UNION SELECT $3;
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ 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 | WITH a AS (SELECT $1) SELECT $2
+ f        |     1 | WITH a AS (SELECT $1) UPDATE stats_track_tab SET x = $2 WHERE x = $3
+ f        |     1 | WITH a AS (select $1) SELECT $2 UNION SELECT $3
 (15 rows)
 
 -- EXPLAIN with CTEs - top-level tracking
@@ -921,12 +918,12 @@ EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                              query                               
-----------+-------+------------------------------------------------------------------
- t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)           +
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ 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
- f        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1;
+ f        |     1 | SELECT $1
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (4 rows)
 
@@ -1050,8 +1047,8 @@ 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 | PREPARE test_prepare_pgss AS select generate_series($1, $2)
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | select generate_series($1, $2)
 (4 rows)
 
 -- CREATE TABLE AS, top-level tracking.
@@ -1202,25 +1199,25 @@ COPY (DELETE FROM stats_track_tab WHERE x = 2 RETURNING x) TO stdout;
 2
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                    query                                    
-----------+-------+-----------------------------------------------------------------------------
- f        |     1 | COPY (DELETE FROM stats_track_tab WHERE x = $1 RETURNING x) TO stdout
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
  t        |     1 | COPY (DELETE FROM stats_track_tab WHERE x = 2 RETURNING x) TO stdout
- f        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES ($1) RETURNING x) TO stdout
  t        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) TO stdout
- f        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT $1 id) ON x = id            +
-          |       |   WHEN MATCHED THEN UPDATE SET x = id                                      +
+ t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id           +
+          |       |   WHEN MATCHED THEN UPDATE SET x = id                                    +
           |       |   WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) TO stdout
- t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id             +
-          |       |   WHEN MATCHED THEN UPDATE SET x = id                                      +
-          |       |   WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) TO stdout
- f        |     1 | COPY (SELECT $1 UNION SELECT $2) TO stdout
- f        |     1 | COPY (SELECT $1) TO stdout
  t        |     1 | COPY (SELECT 1 UNION SELECT 2) TO stdout
  t        |     1 | COPY (SELECT 1) TO stdout
- f        |     1 | COPY (UPDATE stats_track_tab SET x = $1 WHERE x = $2 RETURNING x) TO stdout
  t        |     1 | COPY (UPDATE stats_track_tab SET x = 2 WHERE x = 1 RETURNING x) TO stdout
+ f        |     1 | DELETE FROM stats_track_tab WHERE x = $1 RETURNING x
+ f        |     1 | INSERT INTO stats_track_tab (x) VALUES ($1) RETURNING x
+ f        |     1 | MERGE INTO stats_track_tab USING (SELECT $1 id) ON x = id                +
+          |       |   WHEN MATCHED THEN UPDATE SET x = id                                    +
+          |       |   WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x
+ f        |     1 | SELECT $1
+ f        |     1 | SELECT $1 UNION SELECT $2
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | UPDATE stats_track_tab SET x = $1 WHERE x = $2 RETURNING x
 (13 rows)
 
 -- COPY - top-level tracking.
diff --git a/contrib/pg_stat_statements/expected/planning.out b/contrib/pg_stat_statements/expected/planning.out
index 9effd11fdc..3ee1928cbe 100644
--- a/contrib/pg_stat_statements/expected/planning.out
+++ b/contrib/pg_stat_statements/expected/planning.out
@@ -58,7 +58,7 @@ SELECT 42;
 (1 row)
 
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
  plans | calls | rows |                          query                           
 -------+-------+------+----------------------------------------------------------
      0 |     1 |    0 | ALTER TABLE stats_plan_test ADD COLUMN x int
@@ -72,10 +72,10 @@ SELECT plans, calls, rows, query FROM pg_stat_statements
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
- plans_ok | calls | rows |                         query                         
-----------+-------+------+-------------------------------------------------------
- t        |     4 |    4 | PREPARE prep1 AS SELECT COUNT(*) FROM stats_plan_test
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
+ plans_ok | calls | rows |                query                 
+----------+-------+------+--------------------------------------
+ t        |     4 |    4 | SELECT COUNT(*) FROM stats_plan_test
 (1 row)
 
 -- Cleanup
diff --git a/contrib/pg_stat_statements/expected/select.out b/contrib/pg_stat_statements/expected/select.out
index e0e2fa265c..37a30af034 100644
--- a/contrib/pg_stat_statements/expected/select.out
+++ b/contrib/pg_stat_statements/expected/select.out
@@ -128,7 +128,6 @@ DEALLOCATE pgss_test;
 SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
  calls | rows |                                    query                                     
 -------+------+------------------------------------------------------------------------------
-     1 |    1 | PREPARE pgss_test (int) AS SELECT $1, $2 LIMIT $3
      4 |    4 | SELECT $1                                                                   +
        |      |   -- but this one will appear                                               +
        |      |   AS "text"
@@ -138,6 +137,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
      2 |    2 | SELECT $1 AS "int"
      1 |    2 | SELECT $1 AS i UNION SELECT $2 ORDER BY i
      1 |    1 | SELECT $1 || $2
+     1 |    1 | SELECT $1, $2 LIMIT $3
      0 |    0 | SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C"
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
      1 |    2 | WITH t(f) AS (                                                              +
diff --git a/contrib/pg_stat_statements/expected/utility.out b/contrib/pg_stat_statements/expected/utility.out
index 060d4416dd..aa4f0f7e62 100644
--- a/contrib/pg_stat_statements/expected/utility.out
+++ b/contrib/pg_stat_statements/expected/utility.out
@@ -540,7 +540,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
 -------+------+----------------------------------------------------
      2 |    0 | DEALLOCATE $1
      2 |    0 | DEALLOCATE ALL
-     2 |    2 | PREPARE stat_select AS SELECT $1 AS a
+     2 |    2 | SELECT $1 AS a
      1 |    1 | SELECT $1 as a
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (5 rows)
diff --git a/contrib/pg_stat_statements/sql/planning.sql b/contrib/pg_stat_statements/sql/planning.sql
index 46f5d9b951..9cfe206b3b 100644
--- a/contrib/pg_stat_statements/sql/planning.sql
+++ b/contrib/pg_stat_statements/sql/planning.sql
@@ -20,11 +20,11 @@ SELECT 42;
 SELECT 42;
 SELECT 42;
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 
 -- Cleanup
 DROP TABLE stats_plan_test;
-- 
2.45.2

From 181f55dfe0525b98234c02a4f76df1a8c2e3a61b Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Wed, 23 Oct 2024 15:00:07 +0900
Subject: [PATCH v12 2/2] Some edits from me..

---
 src/backend/parser/analyze.c | 33 +++++++++++++++++++--------------
 src/backend/parser/gram.y    | 19 ++++++++-----------
 2 files changed, 27 insertions(+), 25 deletions(-)

diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index dc71693619..7709fb5e36 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -243,17 +243,17 @@ parse_sub_analyze(Node *parseTree, ParseState *parentParseState,
  * 		Set query's location and length from statement and ParseState
  *
  * Some statements, like PreparableStmt, can be located within parentheses.
- * For example "(SELECT 1)" or "COPY (UPDATE ...) to x;". For those, we can't
- * use the whole string from the statement's location or the SQL string will
- * yield "SELECT 1)". The parser will set stmt_len, reflecting the size of the
- * statement within the parentheses. Thus, when stmt_len is available, we need
- * to use it for the Query's stmt_len.
+ * For example "(SELECT 1)" or "COPY (UPDATE ...) to x;".  For those, we
+ * cannot use the whole string from the statement's location or the SQL
+ * string would yield incorrectly.  The parser will set stmt_len, reflecting
+ * the size of the statement within the parentheses.  Thus, when stmt_len is
+ * available, we need to use it for the Query's stmt_len.
  *
- * For other cases, the parser can't provide the length of individual statements.
- * However, we have the statement's location plus the length (p_stmt_len) and
- * location (p_stmt_location) of the top level RawStmt, stored in pstate. Thus,
- * the statement's length is the RawStmt's length minus how much we've advanced
- * in the RawStmt's string.
+ * For other cases, the parser can't provide the length of individual
+ * statements.  However, we have the statement's location plus its length
+ * (p_stmt_len) and location (p_stmt_location) of the top level RawStmt,
+ * stored in pstate.  Thus, the statement's length is the RawStmt's length
+ * minus how much we've advanced in the RawStmt's string.
  */
 static void
 setQueryLocationAndLength(ParseState *pstate, Query *qry, Node *parseTree)
@@ -266,6 +266,7 @@ setQueryLocationAndLength(ParseState *pstate, Query *qry, Node *parseTree)
 			qry->stmt_location = ((InsertStmt *) parseTree)->stmt_location;
 			stmt_len = ((InsertStmt *) parseTree)->stmt_len;
 			break;
+
 		case T_DeleteStmt:
 			qry->stmt_location = ((DeleteStmt *) parseTree)->stmt_location;
 			stmt_len = ((DeleteStmt *) parseTree)->stmt_len;
@@ -294,16 +295,20 @@ setQueryLocationAndLength(ParseState *pstate, Query *qry, Node *parseTree)
 			qry->stmt_location = pstate->p_stmt_location;
 			break;
 	}
+
 	if (stmt_len > 0)
+	{
 		/* Statement's length is known, use it */
 		qry->stmt_len = stmt_len;
+	}
 	else
-
+	{
 		/*
-		 * Compute the statement's length from statement's location and
-		 * RawStmt's length and location
+		 * Compute the statement's length from the statement's location and
+		 * the RawStmt's length and location.
 		 */
 		qry->stmt_len = pstate->p_stmt_len - (qry->stmt_location - pstate->p_stmt_location);
+	}
 }
 
 /*
@@ -486,7 +491,7 @@ transformStmt(ParseState *pstate, Node *parseTree)
 			 */
 			result = makeNode(Query);
 			result->commandType = CMD_UTILITY;
-			result->utilityStmt = parseTree;
+			result->utilityStmt = (Node *) parseTree;
 			break;
 	}
 
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 172a26e10e..45d02e942a 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -12152,8 +12152,6 @@ InsertStmt:
 					$5->onConflictClause = $6;
 					$5->returningList = $7;
 					$5->withClause = $1;
-					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
-						@$ = @2;
 					$5->stmt_location = @$;
 					$$ = (Node *) $5;
 				}
@@ -12308,8 +12306,6 @@ DeleteStmt: opt_with_clause DELETE_P FROM relation_expr_opt_alias
 					n->whereClause = $6;
 					n->returningList = $7;
 					n->withClause = $1;
-					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
-						@$ = @2;
 					n->stmt_location = @$;
 					$$ = (Node *) n;
 				}
@@ -12385,8 +12381,6 @@ UpdateStmt: opt_with_clause UPDATE relation_expr_opt_alias
 					n->whereClause = $7;
 					n->returningList = $8;
 					n->withClause = $1;
-					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
-						@$ = @2;
 					n->stmt_location = @$;
 					$$ = (Node *) n;
 				}
@@ -12465,8 +12459,6 @@ MergeStmt:
 					m->joinCondition = $8;
 					m->mergeWhenClauses = $9;
 					m->returningList = $10;
-					if (@$ < 0)			/* see comments for YYLLOC_DEFAULT */
-						@$ = @2;
 					m->stmt_location = @$;
 
 					$$ = (Node *) m;
@@ -12711,13 +12703,13 @@ select_with_parens:
 			'(' select_no_parens ')'
 				{
 					SelectStmt *n = (SelectStmt *) $2;
+
 					/*
 					 * As SelectStmt's location starts at the SELECT keyword,
 					 * we need to track the length of the SelectStmt within
 					 * parentheses to be able to extract the relevant part
-					 * of the query. Without this, the RawStmt's length will
-					 * be used and will include the closing parenthesis,
-					 * "SELECT 1)".
+					 * of the query.  Without this, the RawStmt's length would
+					 * be used and would include the closing parenthesis.
 					 */
 					n->stmt_len = @3 - @2;
 					$$ = $2;
@@ -18603,26 +18595,31 @@ updatePreparableStmtEnd(Node *n, int end_location)
 	if (IsA(n, SelectStmt))
 	{
 		SelectStmt *stmt = (SelectStmt *)n;
+
 		stmt->stmt_len = end_location - stmt->stmt_location;
 	}
 	else if (IsA(n, InsertStmt))
 	{
 		InsertStmt *stmt = (InsertStmt *)n;
+
 		stmt->stmt_len = end_location - stmt->stmt_location;
 	}
 	else if (IsA(n, UpdateStmt))
 	{
 		UpdateStmt *stmt = (UpdateStmt *)n;
+
 		stmt->stmt_len = end_location - stmt->stmt_location;
 	}
 	else if (IsA(n, DeleteStmt))
 	{
 		DeleteStmt *stmt = (DeleteStmt *)n;
+
 		stmt->stmt_len = end_location - stmt->stmt_location;
 	}
 	else if (IsA(n, MergeStmt))
 	{
 		MergeStmt *stmt = (MergeStmt *)n;
+
 		stmt->stmt_len = end_location - stmt->stmt_location;
 	}
 	else
-- 
2.45.2

From 58467357f53dbe72527d37a22d0ccb69633bb8b9 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Thu, 3 Oct 2024 08:52:02 +0200
Subject: Track location to extract relevant part in nested statement

Previously, Query generated through transform would have unset
stmt_location. Extensions relying on the statement location to extract
the relevant part of the statement would fallback to use the whole
statement instead, thus showing the same string at the top and
nested level, which was a source of confusion.

This patch fixes the issue by keeping track of the statement locations
and propagate it to Query during transform, allowing pgss to only show
the relevant part of the query for nested query.
---
 .../expected/level_tracking.out               | 165 +++++++++---------
 .../pg_stat_statements/expected/planning.out  |  10 +-
 .../pg_stat_statements/expected/select.out    |   2 +-
 .../pg_stat_statements/expected/utility.out   |   2 +-
 contrib/pg_stat_statements/sql/planning.sql   |   4 +-
 src/backend/parser/analyze.c                  |  93 +++++++++-
 src/backend/parser/gram.y                     |  76 +++++++-
 src/include/nodes/parsenodes.h                |  10 ++
 src/include/parser/parse_node.h               |  15 ++
 9 files changed, 273 insertions(+), 104 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index 8f008f8bfd1..489dc7143f7 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -206,37 +206,37 @@ EXPLAIN (COSTS OFF) SELECT 1 UNION SELECT 2;
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                query                                
-----------+-------+---------------------------------------------------------------------
+ toplevel | calls |                               query                                
+----------+-------+--------------------------------------------------------------------
+ f        |     1 | DELETE FROM stats_track_tab
  t        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2)
- f        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2);
  t        |     1 | EXPLAIN (COSTS OFF) (TABLE test_table)
- f        |     1 | EXPLAIN (COSTS OFF) (TABLE test_table);
  t        |     1 | EXPLAIN (COSTS OFF) (VALUES ($1, $2))
- f        |     1 | EXPLAIN (COSTS OFF) (VALUES ($1, $2));
  t        |     1 | EXPLAIN (COSTS OFF) DELETE FROM stats_track_tab
- f        |     1 | EXPLAIN (COSTS OFF) DELETE FROM stats_track_tab;
  t        |     1 | EXPLAIN (COSTS OFF) INSERT INTO stats_track_tab VALUES (($1))
- f        |     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)
- 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);
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1 UNION SELECT $2
- f        |     1 | EXPLAIN (COSTS OFF) SELECT $1 UNION SELECT $2;
- f        |     1 | EXPLAIN (COSTS OFF) SELECT $1;
  t        |     1 | EXPLAIN (COSTS OFF) TABLE stats_track_tab
- f        |     1 | EXPLAIN (COSTS OFF) TABLE stats_track_tab;
  t        |     1 | EXPLAIN (COSTS OFF) UPDATE stats_track_tab SET x = $1 WHERE x = $2
- f        |     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) VALUES ($1);
+ 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 UNION SELECT $2
+ f        |     1 | SELECT $1, $2
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | TABLE test_table
+ f        |     1 | UPDATE stats_track_tab SET x = $1 WHERE x = $2
+ f        |     1 | VALUES ($1)
+ f        |     1 | VALUES ($1, $2)
 (23 rows)
 
 -- EXPLAIN - top-level tracking.
@@ -405,20 +405,20 @@ EXPLAIN (COSTS OFF) SELECT 1, 2 UNION SELECT 3, 4\; EXPLAIN (COSTS OFF) (SELECT
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                                        query                                                        
-----------+-------+---------------------------------------------------------------------------------------------------------------------
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ f        |     1 | (SELECT $1, $2, $3) UNION SELECT $4, $5, $6
  t        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2, $3)
  t        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2, $3) UNION SELECT $4, $5, $6
- f        |     1 | EXPLAIN (COSTS OFF) (SELECT $1, $2, $3); EXPLAIN (COSTS OFF) (SELECT 1, 2, 3, 4);
  t        |     1 | 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);
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1, $2
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1, $2 UNION SELECT $3, $4
- 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 | SELECT $1
+ f        |     1 | SELECT $1, $2
+ f        |     1 | SELECT $1, $2 UNION SELECT $3, $4
+ f        |     1 | SELECT $1, $2, $3
+ f        |     1 | SELECT $1, $2, $3, $4
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (13 rows)
 
@@ -494,29 +494,29 @@ EXPLAIN (COSTS OFF) INSERT INTO stats_track_tab VALUES ((1))\; EXPLAIN (COSTS OF
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                                              query                                                               
-----------+-------+----------------------------------------------------------------------------------------------------------------------------------
+ toplevel | calls |                               query                                
+----------+-------+--------------------------------------------------------------------
+ f        |     1 | DELETE FROM stats_track_tab
+ f        |     1 | DELETE FROM stats_track_tab WHERE x = $1
  t        |     1 | EXPLAIN (COSTS OFF) (TABLE test_table)
  t        |     1 | EXPLAIN (COSTS OFF) (VALUES ($1, $2))
  t        |     1 | EXPLAIN (COSTS OFF) DELETE FROM stats_track_tab
  t        |     1 | 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) DELETE FROM stats_track_tab; 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))
- 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);
  t        |     1 | EXPLAIN (COSTS OFF) TABLE stats_track_tab
- 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);
  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
- 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;
  t        |     1 | EXPLAIN (COSTS OFF) VALUES ($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 | INSERT INTO stats_track_tab VALUES ($1), ($2)
+ f        |     1 | INSERT INTO stats_track_tab VALUES (($1))
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | TABLE stats_track_tab
+ f        |     1 | TABLE test_table
+ 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)
+ f        |     1 | VALUES ($1, $2)
 (21 rows)
 
 SELECT pg_stat_statements_reset() IS NOT NULL AS t;
@@ -547,21 +547,18 @@ EXPLAIN (COSTS OFF) MERGE INTO stats_track_tab
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                             query                                              
-----------+-------+------------------------------------------------------------------------------------------------
- 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                                                         +
+ toplevel | calls |                             query                             
+----------+-------+---------------------------------------------------------------
+ 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)
- 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;
  t        |     1 | EXPLAIN (COSTS OFF) SELECT $1, $2, $3, $4, $5
+ 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, $2, $3, $4, $5
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (5 rows)
 
@@ -789,29 +786,29 @@ EXPLAIN (COSTS OFF) WITH a AS (select 4) SELECT 1 UNION SELECT 2;
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                           query                                           
-----------+-------+-------------------------------------------------------------------------------------------
+ toplevel | calls |                                          query                                           
+----------+-------+------------------------------------------------------------------------------------------
  t        |     1 | EXPLAIN (COSTS OFF) (WITH a AS (SELECT $1) (SELECT $2, $3))
- f        |     1 | EXPLAIN (COSTS OFF) (WITH a AS (SELECT $1) (SELECT $2, $3));
  t        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) DELETE FROM stats_track_tab
- f        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) DELETE FROM stats_track_tab;
  t        |     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) INSERT INTO stats_track_tab VALUES (($2));
- t        |     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                            +
-          |       |   WHEN MATCHED THEN UPDATE SET x = id                                                    +
+ t        |     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                           +
+          |       |   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) 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);
  t        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) SELECT $2
- f        |     1 | EXPLAIN (COSTS OFF) WITH a AS (SELECT $1) SELECT $2;
  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) UPDATE stats_track_tab SET x = $2 WHERE x = $3;
  t        |     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 UNION SELECT $3;
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ 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 | WITH a AS (SELECT $1) SELECT $2
+ f        |     1 | WITH a AS (SELECT $1) UPDATE stats_track_tab SET x = $2 WHERE x = $3
+ f        |     1 | WITH a AS (select $1) SELECT $2 UNION SELECT $3
 (15 rows)
 
 -- EXPLAIN with CTEs - top-level tracking
@@ -921,12 +918,12 @@ EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
 
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                              query                               
-----------+-------+------------------------------------------------------------------
- t        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)           +
+ toplevel | calls |                              query                              
+----------+-------+-----------------------------------------------------------------
+ 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
- f        |     1 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) SELECT $1;
+ f        |     1 | SELECT $1
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (4 rows)
 
@@ -1050,8 +1047,8 @@ 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 | PREPARE test_prepare_pgss AS select generate_series($1, $2)
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | select generate_series($1, $2)
 (4 rows)
 
 -- CREATE TABLE AS, top-level tracking.
@@ -1202,25 +1199,25 @@ COPY (DELETE FROM stats_track_tab WHERE x = 2 RETURNING x) TO stdout;
 2
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C";
- toplevel | calls |                                    query                                    
-----------+-------+-----------------------------------------------------------------------------
- f        |     1 | COPY (DELETE FROM stats_track_tab WHERE x = $1 RETURNING x) TO stdout
+ toplevel | calls |                                   query                                   
+----------+-------+---------------------------------------------------------------------------
  t        |     1 | COPY (DELETE FROM stats_track_tab WHERE x = 2 RETURNING x) TO stdout
- f        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES ($1) RETURNING x) TO stdout
  t        |     1 | COPY (INSERT INTO stats_track_tab (x) VALUES (1) RETURNING x) TO stdout
- f        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT $1 id) ON x = id            +
-          |       |   WHEN MATCHED THEN UPDATE SET x = id                                      +
+ t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id           +
+          |       |   WHEN MATCHED THEN UPDATE SET x = id                                    +
           |       |   WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) TO stdout
- t        |     1 | COPY (MERGE INTO stats_track_tab USING (SELECT 1 id) ON x = id             +
-          |       |   WHEN MATCHED THEN UPDATE SET x = id                                      +
-          |       |   WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x) TO stdout
- f        |     1 | COPY (SELECT $1 UNION SELECT $2) TO stdout
- f        |     1 | COPY (SELECT $1) TO stdout
  t        |     1 | COPY (SELECT 1 UNION SELECT 2) TO stdout
  t        |     1 | COPY (SELECT 1) TO stdout
- f        |     1 | COPY (UPDATE stats_track_tab SET x = $1 WHERE x = $2 RETURNING x) TO stdout
  t        |     1 | COPY (UPDATE stats_track_tab SET x = 2 WHERE x = 1 RETURNING x) TO stdout
+ f        |     1 | DELETE FROM stats_track_tab WHERE x = $1 RETURNING x
+ f        |     1 | INSERT INTO stats_track_tab (x) VALUES ($1) RETURNING x
+ f        |     1 | MERGE INTO stats_track_tab USING (SELECT $1 id) ON x = id                +
+          |       |   WHEN MATCHED THEN UPDATE SET x = id                                    +
+          |       |   WHEN NOT MATCHED THEN INSERT (x) VALUES (id) RETURNING x
+ f        |     1 | SELECT $1
+ f        |     1 | SELECT $1 UNION SELECT $2
  t        |     1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
+ f        |     1 | UPDATE stats_track_tab SET x = $1 WHERE x = $2 RETURNING x
 (13 rows)
 
 -- COPY - top-level tracking.
diff --git a/contrib/pg_stat_statements/expected/planning.out b/contrib/pg_stat_statements/expected/planning.out
index 9effd11fdc8..3ee1928cbe9 100644
--- a/contrib/pg_stat_statements/expected/planning.out
+++ b/contrib/pg_stat_statements/expected/planning.out
@@ -58,7 +58,7 @@ SELECT 42;
 (1 row)
 
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
  plans | calls | rows |                          query                           
 -------+-------+------+----------------------------------------------------------
      0 |     1 |    0 | ALTER TABLE stats_plan_test ADD COLUMN x int
@@ -72,10 +72,10 @@ SELECT plans, calls, rows, query FROM pg_stat_statements
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
- plans_ok | calls | rows |                         query                         
-----------+-------+------+-------------------------------------------------------
- t        |     4 |    4 | PREPARE prep1 AS SELECT COUNT(*) FROM stats_plan_test
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
+ plans_ok | calls | rows |                query                 
+----------+-------+------+--------------------------------------
+ t        |     4 |    4 | SELECT COUNT(*) FROM stats_plan_test
 (1 row)
 
 -- Cleanup
diff --git a/contrib/pg_stat_statements/expected/select.out b/contrib/pg_stat_statements/expected/select.out
index e0e2fa265c9..37a30af034a 100644
--- a/contrib/pg_stat_statements/expected/select.out
+++ b/contrib/pg_stat_statements/expected/select.out
@@ -128,7 +128,6 @@ DEALLOCATE pgss_test;
 SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
  calls | rows |                                    query                                     
 -------+------+------------------------------------------------------------------------------
-     1 |    1 | PREPARE pgss_test (int) AS SELECT $1, $2 LIMIT $3
      4 |    4 | SELECT $1                                                                   +
        |      |   -- but this one will appear                                               +
        |      |   AS "text"
@@ -138,6 +137,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
      2 |    2 | SELECT $1 AS "int"
      1 |    2 | SELECT $1 AS i UNION SELECT $2 ORDER BY i
      1 |    1 | SELECT $1 || $2
+     1 |    1 | SELECT $1, $2 LIMIT $3
      0 |    0 | SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C"
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
      1 |    2 | WITH t(f) AS (                                                              +
diff --git a/contrib/pg_stat_statements/expected/utility.out b/contrib/pg_stat_statements/expected/utility.out
index 060d4416dd7..aa4f0f7e628 100644
--- a/contrib/pg_stat_statements/expected/utility.out
+++ b/contrib/pg_stat_statements/expected/utility.out
@@ -540,7 +540,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
 -------+------+----------------------------------------------------
      2 |    0 | DEALLOCATE $1
      2 |    0 | DEALLOCATE ALL
-     2 |    2 | PREPARE stat_select AS SELECT $1 AS a
+     2 |    2 | SELECT $1 AS a
      1 |    1 | SELECT $1 as a
      1 |    1 | SELECT pg_stat_statements_reset() IS NOT NULL AS t
 (5 rows)
diff --git a/contrib/pg_stat_statements/sql/planning.sql b/contrib/pg_stat_statements/sql/planning.sql
index 46f5d9b951c..9cfe206b3b0 100644
--- a/contrib/pg_stat_statements/sql/planning.sql
+++ b/contrib/pg_stat_statements/sql/planning.sql
@@ -20,11 +20,11 @@ SELECT 42;
 SELECT 42;
 SELECT 42;
 SELECT plans, calls, rows, query FROM pg_stat_statements
-  WHERE query NOT LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query NOT LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 -- for the prepared statement we expect at least one replan, but cache
 -- invalidations could force more
 SELECT plans >= 2 AND plans <= calls AS plans_ok, calls, rows, query FROM pg_stat_statements
-  WHERE query LIKE 'PREPARE%' ORDER BY query COLLATE "C";
+  WHERE query LIKE 'SELECT COUNT%' ORDER BY query COLLATE "C";
 
 -- Cleanup
 DROP TABLE stats_plan_test;
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index 8a6ba1692e8..7afb065e975 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -238,24 +238,106 @@ parse_sub_analyze(Node *parseTree, ParseState *parentParseState,
 	return query;
 }
 
+/*
+ * setQueryLocationAndLength
+ * 		Set query's location and length from statement and ParseState
+ *
+ * Some statements, like PreparableStmt, can be located within parentheses.
+ * For example "(SELECT 1)" or "COPY (UPDATE ...) to x;".  For those, we
+ * cannot use the whole string from the statement's location or the SQL
+ * string would yield incorrectly.  The parser will set stmt_len, reflecting
+ * the size of the statement within the parentheses.  Thus, when stmt_len is
+ * available, we need to use it for the Query's stmt_len.
+ *
+ * For other cases, the parser can't provide the length of individual
+ * statements.  However, we have the statement's location plus the length
+ * (p_stmt_len) and location (p_stmt_location) of the top level RawStmt,
+ * stored in pstate.  Thus, the statement's length is the RawStmt's length
+ * minus how much we've advanced in the RawStmt's string.
+ */
+static void
+setQueryLocationAndLength(ParseState *pstate, Query *qry, Node *parseTree)
+{
+	ParseLoc	stmt_len = 0;
+
+	/*
+	 * We don't have information about the top RawStmt's length, leave 0 to use
+	 * the whole string
+	 */
+	if (pstate->p_stmt_len == 0)
+		return;
+
+	switch (nodeTag(parseTree))
+	{
+		case T_InsertStmt:
+			qry->stmt_location = ((InsertStmt *) parseTree)->stmt_location;
+			stmt_len = ((InsertStmt *) parseTree)->stmt_len;
+			break;
+
+		case T_DeleteStmt:
+			qry->stmt_location = ((DeleteStmt *) parseTree)->stmt_location;
+			stmt_len = ((DeleteStmt *) parseTree)->stmt_len;
+			break;
+
+		case T_UpdateStmt:
+			qry->stmt_location = ((UpdateStmt *) parseTree)->stmt_location;
+			stmt_len = ((UpdateStmt *) parseTree)->stmt_len;
+			break;
+
+		case T_MergeStmt:
+			qry->stmt_location = ((MergeStmt *) parseTree)->stmt_location;
+			stmt_len = ((MergeStmt *) parseTree)->stmt_len;
+			break;
+
+		case T_SelectStmt:
+			qry->stmt_location = ((SelectStmt *) parseTree)->stmt_location;
+			stmt_len = ((SelectStmt *) parseTree)->stmt_len;
+			break;
+
+		case T_PLAssignStmt:
+			qry->stmt_location = ((PLAssignStmt *) parseTree)->location;
+			break;
+
+		default:
+			qry->stmt_location = pstate->p_stmt_location;
+			break;
+	}
+
+	if (stmt_len > 0)
+	{
+		/* Statement's length is known, use it */
+		qry->stmt_len = stmt_len;
+	}
+	else
+	{
+		/*
+		 * Compute the statement's length from the statement's location and
+		 * the RawStmt's length and location.
+		 */
+		qry->stmt_len = pstate->p_stmt_len - (qry->stmt_location - pstate->p_stmt_location);
+	}
+	Assert(qry->stmt_len >= 0);
+}
+
 /*
  * transformTopLevelStmt -
  *	  transform a Parse tree into a Query tree.
  *
- * This function is just responsible for transferring statement location data
- * from the RawStmt into the finished Query.
+ * This function is just responsible for storing location data
+ * from the RawStmt into the ParseState.
  */
 Query *
 transformTopLevelStmt(ParseState *pstate, RawStmt *parseTree)
 {
 	Query	   *result;
 
+	/* Store RawStmt's length and location in pstate */
+	pstate->p_stmt_len = parseTree->stmt_len;
+	pstate->p_stmt_location = parseTree->stmt_location;
+
 	/* We're at top level, so allow SELECT INTO */
 	result = transformOptionalSelectInto(pstate, parseTree->stmt);
 
-	result->stmt_location = parseTree->stmt_location;
-	result->stmt_len = parseTree->stmt_len;
-
 	return result;
 }
 
@@ -424,6 +506,7 @@ transformStmt(ParseState *pstate, Node *parseTree)
 	/* Mark as original query until we learn differently */
 	result->querySource = QSRC_ORIGINAL;
 	result->canSetTag = true;
+	setQueryLocationAndLength(pstate, result, parseTree);
 
 	return result;
 }
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 4bab2117d96..45d02e942aa 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -153,6 +153,7 @@ static void base_yyerror(YYLTYPE *yylloc, core_yyscan_t yyscanner,
 						 const char *msg);
 static RawStmt *makeRawStmt(Node *stmt, int stmt_location);
 static void updateRawStmtEnd(RawStmt *rs, int end_location);
+static void updatePreparableStmtEnd(Node *n, int end_location);
 static Node *makeColumnRef(char *colname, List *indirection,
 						   int location, core_yyscan_t yyscanner);
 static Node *makeTypeCast(Node *arg, TypeName *typename, int location);
@@ -176,7 +177,7 @@ static void insertSelectOptions(SelectStmt *stmt,
 								SelectLimit *limitClause,
 								WithClause *withClause,
 								core_yyscan_t yyscanner);
-static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg);
+static Node *makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location);
 static Node *doNegate(Node *n, int location);
 static void doNegateFloat(Float *v);
 static Node *makeAndExpr(Node *lexpr, Node *rexpr, int location);
@@ -3383,6 +3384,7 @@ CopyStmt:	COPY opt_binary qualified_name opt_column_list
 				{
 					CopyStmt *n = makeNode(CopyStmt);
 
+					updatePreparableStmtEnd($3, @4);
 					n->relation = NULL;
 					n->query = $3;
 					n->attlist = NIL;
@@ -12150,6 +12152,7 @@ InsertStmt:
 					$5->onConflictClause = $6;
 					$5->returningList = $7;
 					$5->withClause = $1;
+					$5->stmt_location = @$;
 					$$ = (Node *) $5;
 				}
 		;
@@ -12303,6 +12306,7 @@ DeleteStmt: opt_with_clause DELETE_P FROM relation_expr_opt_alias
 					n->whereClause = $6;
 					n->returningList = $7;
 					n->withClause = $1;
+					n->stmt_location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12377,6 +12381,7 @@ UpdateStmt: opt_with_clause UPDATE relation_expr_opt_alias
 					n->whereClause = $7;
 					n->returningList = $8;
 					n->withClause = $1;
+					n->stmt_location = @$;
 					$$ = (Node *) n;
 				}
 		;
@@ -12454,6 +12459,7 @@ MergeStmt:
 					m->joinCondition = $8;
 					m->mergeWhenClauses = $9;
 					m->returningList = $10;
+					m->stmt_location = @$;
 
 					$$ = (Node *) m;
 				}
@@ -12694,7 +12700,20 @@ SelectStmt: select_no_parens			%prec UMINUS
 		;
 
 select_with_parens:
-			'(' select_no_parens ')'				{ $$ = $2; }
+			'(' select_no_parens ')'
+				{
+					SelectStmt *n = (SelectStmt *) $2;
+
+					/*
+					 * As SelectStmt's location starts at the SELECT keyword,
+					 * we need to track the length of the SelectStmt within
+					 * parentheses to be able to extract the relevant part
+					 * of the query.  Without this, the RawStmt's length would
+					 * be used and would include the closing parenthesis.
+					 */
+					n->stmt_len = @3 - @2;
+					$$ = $2;
+				}
 			| '(' select_with_parens ')'			{ $$ = $2; }
 		;
 
@@ -12816,6 +12835,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| SELECT distinct_clause target_list
@@ -12833,6 +12853,7 @@ simple_select:
 					n->groupDistinct = ($7)->distinct;
 					n->havingClause = $8;
 					n->windowClause = $9;
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| values_clause							{ $$ = $1; }
@@ -12853,19 +12874,20 @@ simple_select:
 
 					n->targetList = list_make1(rt);
 					n->fromClause = list_make1($2);
+					n->stmt_location = @1;
 					$$ = (Node *) n;
 				}
 			| select_clause UNION set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause INTERSECT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_INTERSECT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 			| select_clause EXCEPT set_quantifier select_clause
 				{
-					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_EXCEPT, $3 == SET_QUANTIFIER_ALL, $1, $4, @1);
 				}
 		;
 
@@ -13423,6 +13445,7 @@ values_clause:
 				{
 					SelectStmt *n = makeNode(SelectStmt);
 
+					n->stmt_location = @1;
 					n->valuesLists = list_make1($3);
 					$$ = (Node *) n;
 				}
@@ -18565,6 +18588,44 @@ updateRawStmtEnd(RawStmt *rs, int end_location)
 	rs->stmt_len = end_location - rs->stmt_location;
 }
 
+/* Adjust a PreparableStmt to reflect that it doesn't run to the end of the string */
+static void
+updatePreparableStmtEnd(Node *n, int end_location)
+{
+	if (IsA(n, SelectStmt))
+	{
+		SelectStmt *stmt = (SelectStmt *)n;
+
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, InsertStmt))
+	{
+		InsertStmt *stmt = (InsertStmt *)n;
+
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, UpdateStmt))
+	{
+		UpdateStmt *stmt = (UpdateStmt *)n;
+
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, DeleteStmt))
+	{
+		DeleteStmt *stmt = (DeleteStmt *)n;
+
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else if (IsA(n, MergeStmt))
+	{
+		MergeStmt *stmt = (MergeStmt *)n;
+
+		stmt->stmt_len = end_location - stmt->stmt_location;
+	}
+	else
+		elog(ERROR, "unexpected node type %d", (int) n->type);
+}
+
 static Node *
 makeColumnRef(char *colname, List *indirection,
 			  int location, core_yyscan_t yyscanner)
@@ -18943,11 +19004,13 @@ insertSelectOptions(SelectStmt *stmt,
 					 errmsg("multiple WITH clauses not allowed"),
 					 parser_errposition(exprLocation((Node *) withClause))));
 		stmt->withClause = withClause;
+		/* Update SelectStmt's location to the start of the with clause */
+		stmt->stmt_location = withClause->location;
 	}
 }
 
 static Node *
-makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
+makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg, int location)
 {
 	SelectStmt *n = makeNode(SelectStmt);
 
@@ -18955,6 +19018,7 @@ makeSetOp(SetOperation op, bool all, Node *larg, Node *rarg)
 	n->all = all;
 	n->larg = (SelectStmt *) larg;
 	n->rarg = (SelectStmt *) rarg;
+	n->stmt_location = location;
 	return (Node *) n;
 }
 
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index c92cef3d16d..b40b661ec8a 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -2046,6 +2046,8 @@ typedef struct InsertStmt
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
 	OverridingKind override;	/* OVERRIDING clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } InsertStmt;
 
 /* ----------------------
@@ -2060,6 +2062,8 @@ typedef struct DeleteStmt
 	Node	   *whereClause;	/* qualifications */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } DeleteStmt;
 
 /* ----------------------
@@ -2075,6 +2079,8 @@ typedef struct UpdateStmt
 	List	   *fromClause;		/* optional from clause for more tables */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } UpdateStmt;
 
 /* ----------------------
@@ -2090,6 +2096,8 @@ typedef struct MergeStmt
 	List	   *mergeWhenClauses;	/* list of MergeWhenClause(es) */
 	List	   *returningList;	/* list of expressions to return */
 	WithClause *withClause;		/* WITH clause */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 } MergeStmt;
 
 /* ----------------------
@@ -2159,6 +2167,8 @@ typedef struct SelectStmt
 	bool		all;			/* ALL specified? */
 	struct SelectStmt *larg;	/* left child */
 	struct SelectStmt *rarg;	/* right child */
+	ParseLoc	stmt_location;	/* start location, or -1 if unknown */
+	ParseLoc	stmt_len;		/* length in bytes; 0 means "rest of string" */
 	/* Eventually add fields for CORRESPONDING spec here */
 } SelectStmt;
 
diff --git a/src/include/parser/parse_node.h b/src/include/parser/parse_node.h
index 543df568147..fe1fdc4e184 100644
--- a/src/include/parser/parse_node.h
+++ b/src/include/parser/parse_node.h
@@ -108,6 +108,19 @@ typedef Node *(*CoerceParamHook) (ParseState *pstate, Param *param,
  * byte-wise locations in parse structures to character-wise cursor
  * positions.)
  *
+ * p_stmt_location: location of the top level RawStmt's start. During
+ * transformation, Query's location will be set to the statement's
+ * location if available. Otherwise, RawStmt's start will be used.
+ * Propagating the location through ParseState is needed for the Query
+ * length calculation (see p_stmt_len below).
+ *
+ * p_stmt_len: length of the top level RawStmt. Most of the time, the
+ * statement's length won't be provided by the parser, with the exception of
+ * SelectStmt within parentheses and PreparableStmt in Copy. If the statement's
+ * location is provided by the parser, the top level location and length is
+ * needed to accurately compute the Query's length. If the statement's location
+ * is not provided, RawStmt's length can be used directly.
+ *
  * p_rtable: list of RTEs that will become the rangetable of the query.
  * Note that neither relname nor refname of these entries are necessarily
  * unique; searching the rtable by name is a bad idea.
@@ -193,6 +206,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 */
-- 
2.39.5 (Apple Git-154)

From aa70b2ddc36ea33af39b113c361cdc71d49df26d Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@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