v20240918-0008-EXPLAIN-LET-support.patch

From 870edbf19646535fd78fde9a2c4de2a7198bd030 Mon Sep 17 00:00:00 2001
From: "okbob@github.com" <pavel.stehule@gmail.com>
Date: Tue, 21 May 2024 18:28:07 +0200
Subject: [PATCH 08/20] 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 db9d3a8549a..caccc70658c 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 aaec439892c..1dc4dc406cf 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"
@@ -509,8 +510,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);
 }
 
@@ -594,6 +596,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,
+						queryString, params, queryEnv);
+	}
 	else
 	{
 		if (es->format == EXPLAIN_FORMAT_TEXT)
@@ -617,8 +638,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)
@@ -667,6 +688,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 719e63fe823..2046fa8a543 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -663,8 +663,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, queryEnv,
-						   &planduration, (es->buffers ? &bufusage : NULL),
+			ExplainOnePlan(pstmt, into, InvalidOid, es, query_string, paramLI,
+						   queryEnv, &planduration,
+						   (es->buffers ? &bufusage : NULL),
 						   es->memory ? &mem_counters : NULL);
 		else
 			ExplainOneUtility(pstmt->utilityStmt, into, es, query_string,
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 9c43b280258..72ae3b11994 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -12097,7 +12097,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 3ab0aae78f7..e1f04c0d00d 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, const char *queryString,
 							  ParamListInfo params, QueryEnvironment *queryEnv);
 
-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 660423d5de8..8c5d8a5eb29 100644
--- a/src/test/regress/expected/session_variables.out
+++ b/src/test/regress/expected/session_variables.out
@@ -1280,3 +1280,43 @@ SET session_variables_ambiguity_warning TO off;
 DROP TABLE public.xxtab;
 DROP SCHEMA xxtab CASCADE;
 NOTICE:  drop cascades to session variable xxtab.avar
+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) 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 391c1349ffa..3649fc1316b 100644
--- a/src/test/regress/sql/session_variables.sql
+++ b/src/test/regress/sql/session_variables.sql
@@ -873,3 +873,24 @@ SET session_variables_ambiguity_warning TO off;
 
 DROP TABLE public.xxtab;
 DROP SCHEMA xxtab CASCADE;
+
+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) LET var1 = (SELECT count(*) FROM var_tab_test_table);
+
+-- should be 10
+SELECT var1;
+
+DROP VARIABLE var1;
+DROP TABLE var_tab_test_table;
-- 
2.46.1