v20230401-1-0004-support-of-LET-command-in-PLpgSQL.patch

From 50cd4960af0380db04ffa22cf9cfb906fbc48ba8 Mon Sep 17 00:00:00 2001
From: "okbob@github.com" <okbob@github.com>
Date: Mon, 4 Apr 2022 20:32:45 +0200
Subject: [PATCH 04/10] support of LET command in PLpgSQL

The support for LET command as PLpgSQL command does assigning of expression
evaluation to session variable. With extra support in PLpgSQL we can implement
SQL LET command like usual PostgreSQL's command (and the assigning of result
to session variable is not side effect of query evaluation).
---
 src/backend/executor/spi.c              |  3 ++
 src/backend/parser/gram.y               |  8 ++++
 src/backend/parser/parser.c             |  3 +-
 src/include/parser/parser.h             |  6 ++-
 src/pl/plpgsql/src/pl_exec.c            | 55 +++++++++++++++++++++++++
 src/pl/plpgsql/src/pl_funcs.c           | 24 +++++++++++
 src/pl/plpgsql/src/pl_gram.y            | 28 ++++++++++++-
 src/pl/plpgsql/src/pl_reserved_kwlist.h |  1 +
 src/pl/plpgsql/src/plpgsql.h            | 14 ++++++-
 src/tools/pgindent/typedefs.list        |  1 +
 10 files changed, 139 insertions(+), 4 deletions(-)

diff --git a/src/backend/executor/spi.c b/src/backend/executor/spi.c
index e3a170c38b..1459a4d1f8 100644
--- a/src/backend/executor/spi.c
+++ b/src/backend/executor/spi.c
@@ -2967,6 +2967,9 @@ _SPI_error_callback(void *arg)
 			case RAW_PARSE_PLPGSQL_ASSIGN3:
 				errcontext("PL/pgSQL assignment \"%s\"", query);
 				break;
+			case RAW_PARSE_PLPGSQL_LET:
+				errcontext("LET statement \"%s\"", query);
+				break;
 			default:
 				errcontext("SQL statement \"%s\"", query);
 				break;
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index a847c11a5a..8c6c411117 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -812,6 +812,7 @@ static Node *makeRecursiveViewSelect(char *relname, List *aliases, Node *query);
 %token		MODE_PLPGSQL_ASSIGN1
 %token		MODE_PLPGSQL_ASSIGN2
 %token		MODE_PLPGSQL_ASSIGN3
+%token		MODE_PLPGSQL_LET
 
 
 /* Precedence: lowest to highest */
@@ -927,6 +928,13 @@ parse_toplevel:
 				pg_yyget_extra(yyscanner)->parsetree =
 					list_make1(makeRawStmt((Node *) n, 0));
 			}
+			| MODE_PLPGSQL_LET LetStmt
+			{
+				LetStmt *n = (LetStmt *) $2;
+				n->plpgsql_mode = true;
+				pg_yyget_extra(yyscanner)->parsetree =
+					list_make1(makeRawStmt((Node *) n, 0));
+			}
 		;
 
 /*
diff --git a/src/backend/parser/parser.c b/src/backend/parser/parser.c
index 65eb087657..ef911076ef 100644
--- a/src/backend/parser/parser.c
+++ b/src/backend/parser/parser.c
@@ -61,7 +61,8 @@ raw_parser(const char *str, RawParseMode mode)
 			MODE_PLPGSQL_EXPR,	/* RAW_PARSE_PLPGSQL_EXPR */
 			MODE_PLPGSQL_ASSIGN1,	/* RAW_PARSE_PLPGSQL_ASSIGN1 */
 			MODE_PLPGSQL_ASSIGN2,	/* RAW_PARSE_PLPGSQL_ASSIGN2 */
-			MODE_PLPGSQL_ASSIGN3	/* RAW_PARSE_PLPGSQL_ASSIGN3 */
+			MODE_PLPGSQL_ASSIGN3,	/* RAW_PARSE_PLPGSQL_ASSIGN3 */
+			MODE_PLPGSQL_LET	/* RAW_PARSE_PLPGSQL_LET */
 		};
 
 		yyextra.have_lookahead = true;
diff --git a/src/include/parser/parser.h b/src/include/parser/parser.h
index 8d90064d87..0548ad8370 100644
--- a/src/include/parser/parser.h
+++ b/src/include/parser/parser.h
@@ -33,6 +33,9 @@
  * RAW_PARSE_PLPGSQL_ASSIGNn: parse a PL/pgSQL assignment statement,
  * and return a one-element List containing a RawStmt node.  "n"
  * gives the number of dotted names comprising the target ColumnRef.
+ *
+ * RAW_PARSE_PLPGSQL_LET: parse a LET statement, and return a
+ * one-element List containing a RawStmt node.
  */
 typedef enum
 {
@@ -41,7 +44,8 @@ typedef enum
 	RAW_PARSE_PLPGSQL_EXPR,
 	RAW_PARSE_PLPGSQL_ASSIGN1,
 	RAW_PARSE_PLPGSQL_ASSIGN2,
-	RAW_PARSE_PLPGSQL_ASSIGN3
+	RAW_PARSE_PLPGSQL_ASSIGN3,
+	RAW_PARSE_PLPGSQL_LET
 } RawParseMode;
 
 /* Values for the backslash_quote GUC */
diff --git a/src/pl/plpgsql/src/pl_exec.c b/src/pl/plpgsql/src/pl_exec.c
index b0a2cac227..d0c3558a26 100644
--- a/src/pl/plpgsql/src/pl_exec.c
+++ b/src/pl/plpgsql/src/pl_exec.c
@@ -24,6 +24,7 @@
 #include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
 #include "commands/defrem.h"
+#include "commands/session_variable.h"
 #include "executor/execExpr.h"
 #include "executor/spi.h"
 #include "executor/tstoreReceiver.h"
@@ -318,6 +319,8 @@ static int	exec_stmt_commit(PLpgSQL_execstate *estate,
 							 PLpgSQL_stmt_commit *stmt);
 static int	exec_stmt_rollback(PLpgSQL_execstate *estate,
 							   PLpgSQL_stmt_rollback *stmt);
+static int	exec_stmt_let(PLpgSQL_execstate *estate,
+						  PLpgSQL_stmt_let *let);
 
 static void plpgsql_estate_setup(PLpgSQL_execstate *estate,
 								 PLpgSQL_function *func,
@@ -2109,6 +2112,10 @@ exec_stmts(PLpgSQL_execstate *estate, List *stmts)
 				rc = exec_stmt_rollback(estate, (PLpgSQL_stmt_rollback *) stmt);
 				break;
 
+			case PLPGSQL_STMT_LET:
+				rc = exec_stmt_let(estate, (PLpgSQL_stmt_let *) stmt);
+				break;
+
 			default:
 				/* point err_stmt to parent, since this one seems corrupt */
 				estate->err_stmt = save_estmt;
@@ -4977,6 +4984,54 @@ exec_stmt_rollback(PLpgSQL_execstate *estate, PLpgSQL_stmt_rollback *stmt)
 	return PLPGSQL_RC_OK;
 }
 
+/* ----------
+ * exec_stmt_let			Evaluate an expression and
+ *					put the result into a session variable.
+ * ----------
+ */
+static int
+exec_stmt_let(PLpgSQL_execstate *estate, PLpgSQL_stmt_let *stmt)
+{
+	bool		isNull;
+	Oid			valtype;
+	int32		valtypmod;
+	Datum		value;
+	Oid			varid;
+
+	List	   *plansources;
+	CachedPlanSource *plansource;
+
+	value = exec_eval_expr(estate,
+						   stmt->expr,
+						   &isNull,
+						   &valtype,
+						   &valtypmod);
+
+	/*
+	 * Oid of target session variable is stored in Query structure. It is
+	 * safer to read this value directly from the plan than to hold this value
+	 * in the plpgsql context, because it's not necessary to handle
+	 * invalidation of the cached value. Next operations are read only without
+	 * any allocations, so we can expect that retrieving varid from Query
+	 * should be fast.
+	 */
+	plansources = SPI_plan_get_plan_sources(stmt->expr->plan);
+	if (list_length(plansources) != 1)
+		elog(ERROR, "unexpected length of plansources of query for LET statement");
+
+	plansource = (CachedPlanSource *) linitial(plansources);
+	if (list_length(plansource->query_list) != 1)
+		elog(ERROR, "unexpected length of plansource of query for LET statement");
+
+	varid = linitial_node(Query, plansource->query_list)->resultVariable;
+	if (!OidIsValid(varid))
+		elog(ERROR, "oid of target session variable is not valid");
+
+	SetSessionVariableWithSecurityCheck(varid, value, isNull);
+
+	return PLPGSQL_RC_OK;
+}
+
 /* ----------
  * exec_assign_expr			Put an expression's result into a variable.
  * ----------
diff --git a/src/pl/plpgsql/src/pl_funcs.c b/src/pl/plpgsql/src/pl_funcs.c
index 5a6eadccd5..b7513ec0c9 100644
--- a/src/pl/plpgsql/src/pl_funcs.c
+++ b/src/pl/plpgsql/src/pl_funcs.c
@@ -288,6 +288,8 @@ plpgsql_stmt_typename(PLpgSQL_stmt *stmt)
 			return "COMMIT";
 		case PLPGSQL_STMT_ROLLBACK:
 			return "ROLLBACK";
+		case PLPGSQL_STMT_LET:
+			return "LET";
 	}
 
 	return "unknown";
@@ -368,6 +370,7 @@ static void free_perform(PLpgSQL_stmt_perform *stmt);
 static void free_call(PLpgSQL_stmt_call *stmt);
 static void free_commit(PLpgSQL_stmt_commit *stmt);
 static void free_rollback(PLpgSQL_stmt_rollback *stmt);
+static void free_let(PLpgSQL_stmt_let *stmt);
 static void free_expr(PLpgSQL_expr *expr);
 
 
@@ -457,6 +460,9 @@ free_stmt(PLpgSQL_stmt *stmt)
 		case PLPGSQL_STMT_ROLLBACK:
 			free_rollback((PLpgSQL_stmt_rollback *) stmt);
 			break;
+		case PLPGSQL_STMT_LET:
+			free_let((PLpgSQL_stmt_let *) stmt);
+			break;
 		default:
 			elog(ERROR, "unrecognized cmd_type: %d", stmt->cmd_type);
 			break;
@@ -711,6 +717,12 @@ free_getdiag(PLpgSQL_stmt_getdiag *stmt)
 {
 }
 
+static void
+free_let(PLpgSQL_stmt_let *stmt)
+{
+	free_expr(stmt->expr);
+}
+
 static void
 free_expr(PLpgSQL_expr *expr)
 {
@@ -813,6 +825,7 @@ static void dump_perform(PLpgSQL_stmt_perform *stmt);
 static void dump_call(PLpgSQL_stmt_call *stmt);
 static void dump_commit(PLpgSQL_stmt_commit *stmt);
 static void dump_rollback(PLpgSQL_stmt_rollback *stmt);
+static void dump_let(PLpgSQL_stmt_let *stmt);
 static void dump_expr(PLpgSQL_expr *expr);
 
 
@@ -912,6 +925,9 @@ dump_stmt(PLpgSQL_stmt *stmt)
 		case PLPGSQL_STMT_ROLLBACK:
 			dump_rollback((PLpgSQL_stmt_rollback *) stmt);
 			break;
+		case PLPGSQL_STMT_LET:
+			dump_let((PLpgSQL_stmt_let *) stmt);
+			break;
 		default:
 			elog(ERROR, "unrecognized cmd_type: %d", stmt->cmd_type);
 			break;
@@ -1588,6 +1604,14 @@ dump_getdiag(PLpgSQL_stmt_getdiag *stmt)
 	printf("\n");
 }
 
+static void
+dump_let(PLpgSQL_stmt_let *stmt)
+{
+	dump_ind();
+	dump_expr(stmt->expr);
+	printf("\n");
+}
+
 static void
 dump_expr(PLpgSQL_expr *expr)
 {
diff --git a/src/pl/plpgsql/src/pl_gram.y b/src/pl/plpgsql/src/pl_gram.y
index edeb72c380..77eb0580d7 100644
--- a/src/pl/plpgsql/src/pl_gram.y
+++ b/src/pl/plpgsql/src/pl_gram.y
@@ -195,7 +195,7 @@ static	void			check_raise_parameters(PLpgSQL_stmt_raise *stmt);
 %type <stmt>	stmt_return stmt_raise stmt_assert stmt_execsql
 %type <stmt>	stmt_dynexecute stmt_for stmt_perform stmt_call stmt_getdiag
 %type <stmt>	stmt_open stmt_fetch stmt_move stmt_close stmt_null
-%type <stmt>	stmt_commit stmt_rollback
+%type <stmt>	stmt_commit stmt_rollback stmt_let
 %type <stmt>	stmt_case stmt_foreach_a
 
 %type <list>	proc_exceptions
@@ -302,6 +302,7 @@ static	void			check_raise_parameters(PLpgSQL_stmt_raise *stmt);
 %token <keyword>	K_INTO
 %token <keyword>	K_IS
 %token <keyword>	K_LAST
+%token <keyword>	K_LET
 %token <keyword>	K_LOG
 %token <keyword>	K_LOOP
 %token <keyword>	K_MERGE
@@ -869,6 +870,8 @@ proc_stmt		: pl_block ';'
 						{ $$ = $1; }
 				| stmt_rollback
 						{ $$ = $1; }
+				| stmt_let
+						{ $$ = $1; }
 				;
 
 stmt_perform	: K_PERFORM
@@ -985,6 +988,29 @@ stmt_assign		: T_DATUM
 					}
 				;
 
+stmt_let		: K_LET
+					{
+						PLpgSQL_stmt_let *new;
+						RawParseMode pmode;
+
+						pmode = RAW_PARSE_PLPGSQL_LET;
+
+						new = palloc0(sizeof(PLpgSQL_stmt_let));
+						new->cmd_type = PLPGSQL_STMT_LET;
+						new->lineno   = plpgsql_location_to_lineno(@1);
+						new->stmtid = ++plpgsql_curr_compile->nstatements;
+
+						/* Push back the head name to include it in the stmt */
+						plpgsql_push_back_token(K_LET);
+						new->expr = read_sql_construct(';', 0, 0, ";",
+													   pmode,
+													   false, true, true,
+													   NULL, NULL);
+
+						$$ = (PLpgSQL_stmt *)new;
+					}
+				;
+
 stmt_getdiag	: K_GET getdiag_area_opt K_DIAGNOSTICS getdiag_list ';'
 					{
 						PLpgSQL_stmt_getdiag *new;
diff --git a/src/pl/plpgsql/src/pl_reserved_kwlist.h b/src/pl/plpgsql/src/pl_reserved_kwlist.h
index c116abbb7a..90eb1ca8aa 100644
--- a/src/pl/plpgsql/src/pl_reserved_kwlist.h
+++ b/src/pl/plpgsql/src/pl_reserved_kwlist.h
@@ -40,6 +40,7 @@ PG_KEYWORD("from", K_FROM)
 PG_KEYWORD("if", K_IF)
 PG_KEYWORD("in", K_IN)
 PG_KEYWORD("into", K_INTO)
+PG_KEYWORD("let", K_LET)
 PG_KEYWORD("loop", K_LOOP)
 PG_KEYWORD("not", K_NOT)
 PG_KEYWORD("null", K_NULL)
diff --git a/src/pl/plpgsql/src/plpgsql.h b/src/pl/plpgsql/src/plpgsql.h
index 355c9f678d..9d6d47b3a3 100644
--- a/src/pl/plpgsql/src/plpgsql.h
+++ b/src/pl/plpgsql/src/plpgsql.h
@@ -127,7 +127,8 @@ typedef enum PLpgSQL_stmt_type
 	PLPGSQL_STMT_PERFORM,
 	PLPGSQL_STMT_CALL,
 	PLPGSQL_STMT_COMMIT,
-	PLPGSQL_STMT_ROLLBACK
+	PLPGSQL_STMT_ROLLBACK,
+	PLPGSQL_STMT_LET
 } PLpgSQL_stmt_type;
 
 /*
@@ -519,6 +520,17 @@ typedef struct PLpgSQL_stmt_assign
 	PLpgSQL_expr *expr;
 } PLpgSQL_stmt_assign;
 
+/*
+ * Let statement
+ */
+typedef struct PLpgSQL_stmt_let
+{
+	PLpgSQL_stmt_type cmd_type;
+	int			lineno;
+	unsigned int stmtid;
+	PLpgSQL_expr *expr;
+} PLpgSQL_stmt_let;
+
 /*
  * PERFORM statement
  */
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 6ede939ede..c020a2f298 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -1795,6 +1795,7 @@ PLpgSQL_stmt_forq
 PLpgSQL_stmt_fors
 PLpgSQL_stmt_getdiag
 PLpgSQL_stmt_if
+PLpgSQL_stmt_let
 PLpgSQL_stmt_loop
 PLpgSQL_stmt_open
 PLpgSQL_stmt_perform
-- 
2.40.0