v20241228-0010-EXPLAIN-LET-support.patch

From 0e74b9c8a69d604a795ce3884b464e7d16e0226b 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 10/22] 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                | 31 ++++++++++++--
 src/backend/commands/prepare.c                |  5 ++-
 src/backend/parser/gram.y                     |  3 +-
 src/include/commands/explain.h                |  3 +-
 .../regress/expected/session_variables.out    | 40 +++++++++++++++++++
 src/test/regress/sql/session_variables.sql    | 21 ++++++++++
 7 files changed, 97 insertions(+), 9 deletions(-)

diff --git a/doc/src/sgml/ref/explain.sgml b/doc/src/sgml/ref/explain.sgml
index 6361a14e65..4292615b44 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 a201ed3082..e7ad5243eb 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -18,6 +18,7 @@
 #include "commands/createas.h"
 #include "commands/defrem.h"
 #include "commands/prepare.h"
+#include "executor/svariableReceiver.h"
 #include "foreign/fdwapi.h"
 #include "jit/jit.h"
 #include "libpq/pqformat.h"
@@ -519,8 +520,9 @@ standard_ExplainOneQuery(Query *query, int cursorOptions,
 	}
 
 	/* run it (if needed) and produce output */
-	ExplainOnePlan(plan, into, es, queryString, params, queryEnv,
-				   &planduration, (es->buffers ? &bufusage : NULL),
+	ExplainOnePlan(plan, into, query->resultVariable, es, queryString,
+				   params, queryEnv, &planduration,
+				   (es->buffers ? &bufusage : NULL),
 				   es->memory ? &mem_counters : NULL);
 }
 
@@ -618,6 +620,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)
@@ -641,8 +662,8 @@ ExplainOneUtility(Node *utilityStmt, IntoClause *into, ExplainState *es,
  * to call it.
  */
 void
-ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into, ExplainState *es,
-			   const char *queryString, ParamListInfo params,
+ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into, Oid targetvar,
+			   ExplainState *es, const char *queryString, ParamListInfo params,
 			   QueryEnvironment *queryEnv, const instr_time *planduration,
 			   const BufferUsage *bufusage,
 			   const MemoryContextCounters *mem_counters)
@@ -691,6 +712,8 @@ ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into, ExplainState *es,
 	 */
 	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 e9e393aa02..af42496263 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -662,8 +662,9 @@ 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, pstate->p_queryEnv,
-						   &planduration, (es->buffers ? &bufusage : NULL),
+			ExplainOnePlan(pstmt, 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 fb8c3d608b..fbdc240170 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -12129,7 +12129,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 aa5872bc15..f5b08857ae 100644
--- a/src/include/commands/explain.h
+++ b/src/include/commands/explain.h
@@ -103,7 +103,8 @@ extern void ExplainOneUtility(Node *utilityStmt, IntoClause *into,
 							  ExplainState *es, ParseState *pstate,
 							  ParamListInfo params);
 
-extern void ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into,
+extern void ExplainOnePlan(PlannedStmt *plannedstmt,
+						   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 8e4b44763d..478631f98e 100644
--- a/src/test/regress/expected/session_variables.out
+++ b/src/test/regress/expected/session_variables.out
@@ -1950,3 +1950,43 @@ NOTICE:  20 t
 SET session_variables_use_fence_context_guard TO DEFAULT;
 DROP VARIABLE var1;
 DROP TABLE vartest_tab1;
+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 loops=1)
+   InitPlan 1
+     ->  Aggregate (actual rows=1 loops=1)
+           ->  Seq Scan on var_tab_test_table (actual rows=10 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 447b042e70..0dfc53de44 100644
--- a/src/test/regress/sql/session_variables.sql
+++ b/src/test/regress/sql/session_variables.sql
@@ -1328,3 +1328,24 @@ SET session_variables_use_fence_context_guard TO DEFAULT;
 
 DROP VARIABLE var1;
 DROP TABLE vartest_tab1;
+
+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.47.1