v20240125-0003-EXPLAIN-LET-support.patch

From fe97583f3f7fd05a5f21b517b8e0d13116eaeb83 Mon Sep 17 00:00:00 2001
From: "okbob@github.com" <okbob@github.com>
Date: Fri, 19 Jan 2024 08:28:25 +0100
Subject: [PATCH 03/18] 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                |  2 +-
 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, 93 insertions(+), 7 deletions(-)

diff --git a/doc/src/sgml/ref/explain.sgml b/doc/src/sgml/ref/explain.sgml
index 5ba6486da1..71750c91e4 100644
--- a/doc/src/sgml/ref/explain.sgml
+++ b/doc/src/sgml/ref/explain.sgml
@@ -96,7 +96,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 843472e6dd..d68e0b000a 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -19,6 +19,7 @@
 #include "commands/defrem.h"
 #include "commands/prepare.h"
 #include "executor/nodeHash.h"
+#include "executor/svariableReceiver.h"
 #include "foreign/fdwapi.h"
 #include "jit/jit.h"
 #include "nodes/extensible.h"
@@ -416,7 +417,7 @@ ExplainOneQuery(Query *query, int cursorOptions,
 		}
 
 		/* run it (if needed) and produce output */
-		ExplainOnePlan(plan, into, es, queryString, params, queryEnv,
+		ExplainOnePlan(plan, into, query->resultVariable, es, queryString, params, queryEnv,
 					   &planduration, (es->buffers ? &bufusage : NULL));
 	}
 }
@@ -501,6 +502,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)
@@ -524,8 +544,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)
 {
@@ -568,6 +588,8 @@ ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into, ExplainState *es,
 	 */
 	if (into)
 		dest = CreateIntoRelDestReceiver(into);
+	else if (OidIsValid(targetvar))
+		dest = CreateVariableDestReceiver(targetvar);
 	else
 		dest = None_Receiver;
 
diff --git a/src/backend/commands/prepare.c b/src/backend/commands/prepare.c
index 703976f633..f315720bb7 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -639,7 +639,7 @@ 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, InvalidOid, es, query_string, paramLI, queryEnv,
 						   &planduration, (es->buffers ? &bufusage : NULL));
 		else
 			ExplainOneUtility(pstmt->utilityStmt, into, es, query_string,
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 7a26c2c9a5..fcb14acf28 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -12019,7 +12019,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 1b44d483d6..7d576b0016 100644
--- a/src/include/commands/explain.h
+++ b/src/include/commands/explain.h
@@ -88,7 +88,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 e0d5ac12ab..1cff6c7291 100644
--- a/src/test/regress/expected/session_variables.out
+++ b/src/test/regress/expected/session_variables.out
@@ -942,3 +942,43 @@ DROP VARIABLE :"DBNAME".:"DBNAME".b;
 DROP VARIABLE :"DBNAME".:"DBNAME".:"DBNAME";
 DROP SCHEMA :"DBNAME";
 RESET search_path;
+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 (returns $0)
+     ->  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 (returns $0)
+     ->  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 1f72b405a2..4f2ea892e5 100644
--- a/src/test/regress/sql/session_variables.sql
+++ b/src/test/regress/sql/session_variables.sql
@@ -663,3 +663,24 @@ DROP VARIABLE :"DBNAME".:"DBNAME".:"DBNAME";
 DROP SCHEMA :"DBNAME";
 
 RESET search_path;
+
+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.43.0