v20250301-0009-EXPLAIN-LET-support.patch

From 18127f4da2f42ebd3c35435a77e4efbf32d7c046 Mon Sep 17 00:00:00 2001
From: "okbob@github.com" <okbob@github.com>
Date: Tue, 21 May 2024 18:28:07 +0200
Subject: [PATCH 09/21] EXPLAIN LET support

Enhancing ExplainOnePlan is necessary to be EXPLAIN ANALYZE LET fully workable.
In this case we want to be result of query or expression written to target variable.
---
 doc/src/sgml/ref/explain.sgml                 |  3 +-
 src/backend/commands/explain.c                | 28 +++++++++++--
 src/backend/commands/prepare.c                |  5 ++-
 src/backend/parser/gram.y                     |  3 +-
 src/include/commands/explain.h                |  2 +-
 .../regress/expected/session_variables.out    | 40 +++++++++++++++++++
 src/test/regress/sql/session_variables.sql    | 21 ++++++++++
 7 files changed, 94 insertions(+), 8 deletions(-)

diff --git a/doc/src/sgml/ref/explain.sgml b/doc/src/sgml/ref/explain.sgml
index 7daddf03ef0..079ccdf3cec 100644
--- a/doc/src/sgml/ref/explain.sgml
+++ b/doc/src/sgml/ref/explain.sgml
@@ -98,7 +98,8 @@ EXPLAIN [ ( <replaceable class="parameter">option</replaceable> [, ...] ) ] <rep
     <command>INSERT</command>, <command>UPDATE</command>,
     <command>DELETE</command>, <command>MERGE</command>,
     <command>CREATE TABLE AS</command>,
-    or <command>EXECUTE</command> statement
+    <command>EXECUTE</command>,
+    or <command>LET</command> statement
     without letting the command affect your data, use this approach:
 <programlisting>
 BEGIN;
diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
index d8a7232cedb..b303561b8de 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -20,6 +20,7 @@
 #include "commands/explain_dr.h"
 #include "commands/explain_format.h"
 #include "commands/prepare.h"
+#include "executor/svariableReceiver.h"
 #include "foreign/fdwapi.h"
 #include "jit/jit.h"
 #include "libpq/pqformat.h"
@@ -493,8 +494,8 @@ standard_ExplainOneQuery(Query *query, int cursorOptions,
 	}
 
 	/* run it (if needed) and produce output */
-	ExplainOnePlan(plan, NULL, NULL, -1, into, es, queryString, params,
-				   queryEnv,
+	ExplainOnePlan(plan, NULL, NULL, -1, into, query->resultVariable,
+				   es, queryString, params, queryEnv,
 				   &planduration, (es->buffers ? &bufusage : NULL),
 				   es->memory ? &mem_counters : NULL);
 }
@@ -593,6 +594,25 @@ ExplainOneUtility(Node *utilityStmt, IntoClause *into, ExplainState *es,
 		else
 			ExplainDummyGroup("Notify", NULL, es);
 	}
+	else if (IsA(utilityStmt, LetStmt))
+	{
+		LetStmt    *letstmt = (LetStmt *) utilityStmt;
+		List	   *rewritten;
+		Query	   *query;
+
+		if (es->format == EXPLAIN_FORMAT_TEXT)
+			appendStringInfoString(es->str, "SET SESSION VARIABLE\n");
+		else
+			ExplainDummyGroup("Set Session Variable", NULL, es);
+
+		rewritten = QueryRewrite(castNode(Query, copyObject(letstmt->query)));
+
+		Assert(list_length(rewritten) == 1);
+		query = linitial_node(Query, rewritten);
+		ExplainOneQuery(query,
+						CURSOR_OPT_PARALLEL_OK, NULL, es,
+						pstate, params);
+	}
 	else
 	{
 		if (es->format == EXPLAIN_FORMAT_TEXT)
@@ -618,7 +638,7 @@ ExplainOneUtility(Node *utilityStmt, IntoClause *into, ExplainState *es,
 void
 ExplainOnePlan(PlannedStmt *plannedstmt, CachedPlan *cplan,
 			   CachedPlanSource *plansource, int query_index,
-			   IntoClause *into, ExplainState *es,
+			   IntoClause *into, Oid targetvar, ExplainState *es,
 			   const char *queryString, ParamListInfo params,
 			   QueryEnvironment *queryEnv, const instr_time *planduration,
 			   const BufferUsage *bufusage,
@@ -668,6 +688,8 @@ ExplainOnePlan(PlannedStmt *plannedstmt, CachedPlan *cplan,
 	 */
 	if (into)
 		dest = CreateIntoRelDestReceiver(into);
+	else if (OidIsValid(targetvar))
+		dest = CreateVariableDestReceiver(targetvar);
 	else if (es->serialize != EXPLAIN_SERIALIZE_NONE)
 		dest = CreateExplainSerializeDestReceiver(es);
 	else
diff --git a/src/backend/commands/prepare.c b/src/backend/commands/prepare.c
index 2d04c71a148..13d198f1625 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -668,8 +668,9 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 
 		if (pstmt->commandType != CMD_UTILITY)
 			ExplainOnePlan(pstmt, cplan, entry->plansource, query_index,
-						   into, es, query_string, paramLI, pstate->p_queryEnv,
-						   &planduration, (es->buffers ? &bufusage : NULL),
+						   into, InvalidOid, es, query_string, paramLI,
+						   pstate->p_queryEnv, &planduration,
+						   (es->buffers ? &bufusage : NULL),
 						   es->memory ? &mem_counters : NULL);
 		else
 			ExplainOneUtility(pstmt->utilityStmt, into, es, pstate, paramLI);
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 4fffb28dd4d..73e2ecb662c 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -12167,7 +12167,8 @@ ExplainableStmt:
 			| CreateAsStmt
 			| CreateMatViewStmt
 			| RefreshMatViewStmt
-			| ExecuteStmt					/* by default all are $$=$1 */
+			| ExecuteStmt
+			| LetStmt						/* by default all are $$=$1 */
 		;
 
 /*****************************************************************************
diff --git a/src/include/commands/explain.h b/src/include/commands/explain.h
index 64547bd9b9c..18932a9fe14 100644
--- a/src/include/commands/explain.h
+++ b/src/include/commands/explain.h
@@ -105,7 +105,7 @@ extern void ExplainOneUtility(Node *utilityStmt, IntoClause *into,
 
 extern void ExplainOnePlan(PlannedStmt *plannedstmt, CachedPlan *cplan,
 						   CachedPlanSource *plansource, int plan_index,
-						   IntoClause *into, ExplainState *es,
+						   IntoClause *into, Oid targetvar, ExplainState *es,
 						   const char *queryString,
 						   ParamListInfo params, QueryEnvironment *queryEnv,
 						   const instr_time *planduration,
diff --git a/src/test/regress/expected/session_variables.out b/src/test/regress/expected/session_variables.out
index 9ac08a14e0b..22cdd356641 100644
--- a/src/test/regress/expected/session_variables.out
+++ b/src/test/regress/expected/session_variables.out
@@ -2028,3 +2028,43 @@ DROP SCHEMA testvar;
 SET session_variables_ambiguity_warning TO DEFAULT;
 SET session_variables_use_fence_warning_guard TO DEFAULT;
 SET search_path TO DEFAULT;
+CREATE VARIABLE var1 bigint;
+CREATE TABLE var_tab_test_table(a int);
+INSERT INTO var_tab_test_table SELECT * FROM generate_series(1,10);
+VACUUM ANALYZE var_tab_test_table;
+EXPLAIN (COSTS OFF) LET var1 = (SELECT count(*) FROM var_tab_test_table);
+                  QUERY PLAN                  
+----------------------------------------------
+ SET SESSION VARIABLE
+ Result
+   InitPlan 1
+     ->  Aggregate
+           ->  Seq Scan on var_tab_test_table
+(5 rows)
+
+-- should be NULL
+SELECT var1;
+ var1 
+------
+     
+(1 row)
+
+EXPLAIN (COSTS OFF, TIMING OFF, ANALYZE, SUMMARY OFF, BUFFERS OFF) LET var1 = (SELECT count(*) FROM var_tab_test_table);
+                                QUERY PLAN                                
+--------------------------------------------------------------------------
+ SET SESSION VARIABLE
+ Result (actual rows=1.00 loops=1)
+   InitPlan 1
+     ->  Aggregate (actual rows=1.00 loops=1)
+           ->  Seq Scan on var_tab_test_table (actual rows=10.00 loops=1)
+(5 rows)
+
+-- should be 10
+SELECT var1;
+ var1 
+------
+   10
+(1 row)
+
+DROP VARIABLE var1;
+DROP TABLE var_tab_test_table;
diff --git a/src/test/regress/sql/session_variables.sql b/src/test/regress/sql/session_variables.sql
index c5420183d94..3d0eb8c7ad9 100644
--- a/src/test/regress/sql/session_variables.sql
+++ b/src/test/regress/sql/session_variables.sql
@@ -1401,3 +1401,24 @@ DROP SCHEMA testvar;
 SET session_variables_ambiguity_warning TO DEFAULT;
 SET session_variables_use_fence_warning_guard TO DEFAULT;
 SET search_path TO DEFAULT;
+
+CREATE VARIABLE var1 bigint;
+
+CREATE TABLE var_tab_test_table(a int);
+
+INSERT INTO var_tab_test_table SELECT * FROM generate_series(1,10);
+
+VACUUM ANALYZE var_tab_test_table;
+
+EXPLAIN (COSTS OFF) LET var1 = (SELECT count(*) FROM var_tab_test_table);
+
+-- should be NULL
+SELECT var1;
+
+EXPLAIN (COSTS OFF, TIMING OFF, ANALYZE, SUMMARY OFF, BUFFERS OFF) LET var1 = (SELECT count(*) FROM var_tab_test_table);
+
+-- should be 10
+SELECT var1;
+
+DROP VARIABLE var1;
+DROP TABLE var_tab_test_table;
-- 
2.48.1