v20241229-0018-plpgsql-implementation-for-LET-statement.patch

From 6233ef3351890d3e7903af8a01a869eb0d3149fa Mon Sep 17 00:00:00 2001
From: "okbob@github.com" <okbob@github.com>
Date: Sat, 20 Jan 2024 08:56:17 +0100
Subject: [PATCH 18/22] plpgsql implementation for LET statement

PLpgSQL allows to call expression executor directly (for simple expression).
This possibility strongly reduces overhead related to execution. Reimplementation
of LET statement (inside PLpgSQL) allows to use this possibility, and strongly
increase performance:

CREATE VARIABLE svar int;

DO $$
BEGIN
  FOR i IN 1..10000
  LOOP
    LET svar = i;
  END LOOP;
END;
$$;

From 120ms to 8ms (with assertions) (this is best case, but without it, the LET statement
can be bottle neck).

An alternative can be reimplementation of LET statement inside expression
executor, and then SQL LET command can be executed in simple expression
execution, but this increase an complexity of executor, but still the
benefits is only inside plpgsql, so it is better to this optimization
inside plpgsql.
---
 src/backend/commands/session_variable.c       | 16 ++++++
 src/backend/executor/spi.c                    |  3 +
 src/backend/parser/analyze.c                  | 12 ++++
 src/backend/parser/gram.y                     |  8 +++
 src/backend/parser/parser.c                   |  1 +
 src/include/commands/session_variable.h       |  1 +
 src/include/nodes/parsenodes.h                |  2 +
 src/include/parser/parser.h                   |  4 ++
 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/pl_unreserved_kwlist.h     |  1 +
 src/pl/plpgsql/src/plpgsql.h                  | 12 ++++
 .../regress/expected/session_variables.out    |  6 +-
 src/tools/pgindent/typedefs.list              |  1 +
 16 files changed, 170 insertions(+), 5 deletions(-)

diff --git a/src/backend/commands/session_variable.c b/src/backend/commands/session_variable.c
index c7c7aa57f6..5e02cb990c 100644
--- a/src/backend/commands/session_variable.c
+++ b/src/backend/commands/session_variable.c
@@ -956,6 +956,22 @@ SetSessionVariable(Oid varid, Datum value, bool isNull)
 	set_session_variable(svar, value, isNull);
 }
 
+/*
+ * Wrapper around SetSessionVariable with permission checks.
+ */
+void
+SetSessionVariableWithSecurityCheck(Oid varid, Datum value, bool isNull)
+{
+	AclResult	aclresult;
+
+	/* is the caller allowed to update the session variable? */
+	aclresult = object_aclcheck(VariableRelationId, varid, GetUserId(), ACL_UPDATE);
+	if (aclresult != ACLCHECK_OK)
+		aclcheck_error(aclresult, OBJECT_VARIABLE, get_session_variable_name(varid));
+
+	SetSessionVariable(varid, value, isNull);
+}
+
 /*
  * Returns a copy of the value stored in a variable.
  */
diff --git a/src/backend/executor/spi.c b/src/backend/executor/spi.c
index 659bae0fbc..3417ed49a5 100644
--- a/src/backend/executor/spi.c
+++ b/src/backend/executor/spi.c
@@ -2999,6 +2999,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/analyze.c b/src/backend/parser/analyze.c
index e41465b327..af402b6547 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -2013,6 +2013,18 @@ transformLetStmt(ParseState *pstate, LetStmt *stmt)
 
 	stmt->query = (Node *) query;
 
+	/*
+	 * Inside PL/pgSQL we don't want to execute LET statement as utility
+	 * command, because it disallow to execute expression as simple
+	 * expression. So for PL/pgSQL we have extra path, and we return SELECT.
+	 * Then it can be executed by exec_eval_expr. Result is dirrectly assigned
+	 * to target session variable inside PL/pgSQL LET statement handler. This
+	 * is extra code, extra path, but possibility to get faster execution is
+	 * too attractive.
+	 */
+	if (stmt->plpgsql_mode)
+		return query;
+
 	/* represent the command as a utility Query */
 	result = makeNode(Query);
 	result->commandType = CMD_UTILITY;
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 4256d51ef2..5134657e61 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -816,6 +816,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 */
@@ -947,6 +948,13 @@ parse_toplevel:
 				pg_yyget_extra(yyscanner)->parsetree =
 					list_make1(makeRawStmt((Node *) n, @2));
 			}
+			| 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 118488c3f3..f9430a3291 100644
--- a/src/backend/parser/parser.c
+++ b/src/backend/parser/parser.c
@@ -62,6 +62,7 @@ raw_parser(const char *str, RawParseMode mode)
 			[RAW_PARSE_PLPGSQL_ASSIGN1] = MODE_PLPGSQL_ASSIGN1,
 			[RAW_PARSE_PLPGSQL_ASSIGN2] = MODE_PLPGSQL_ASSIGN2,
 			[RAW_PARSE_PLPGSQL_ASSIGN3] = MODE_PLPGSQL_ASSIGN3,
+			[RAW_PARSE_PLPGSQL_LET] = MODE_PLPGSQL_LET,
 		};
 
 		yyextra.have_lookahead = true;
diff --git a/src/include/commands/session_variable.h b/src/include/commands/session_variable.h
index 00451115dd..1fb97b7fe1 100644
--- a/src/include/commands/session_variable.h
+++ b/src/include/commands/session_variable.h
@@ -28,6 +28,7 @@ extern void AtEOSubXact_SessionVariables(bool isCommit, SubTransactionId mySubid
 										 SubTransactionId parentSubid);
 
 extern void SetSessionVariable(Oid varid, Datum value, bool isNull);
+extern void SetSessionVariableWithSecurityCheck(Oid varid, Datum value, bool isNull);
 extern Datum GetSessionVariable(Oid varid, bool *isNull, Oid *typid);
 extern Datum GetSessionVariableWithTypeCheck(Oid varid, bool *isNull, Oid expected_typid);
 
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 77714850dd..a453760811 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -2124,6 +2124,8 @@ typedef struct LetStmt
 	NodeTag		type;
 	List	   *target;			/* target variable */
 	Node	   *query;			/* source expression */
+	bool		plpgsql_mode;	/* true, when command will be executed
+								 * (parsed) by plpgsql runtime */
 	ParseLoc	location;
 } LetStmt;
 
diff --git a/src/include/parser/parser.h b/src/include/parser/parser.h
index be184ec506..efbc76f0ad 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
 {
@@ -42,6 +45,7 @@ typedef enum
 	RAW_PARSE_PLPGSQL_ASSIGN1,
 	RAW_PARSE_PLPGSQL_ASSIGN2,
 	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 e31206e7f4..857ff00026 100644
--- a/src/pl/plpgsql/src/pl_exec.c
+++ b/src/pl/plpgsql/src/pl_exec.c
@@ -22,6 +22,7 @@
 #include "access/tupconvert.h"
 #include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
+#include "commands/session_variable.h"
 #include "executor/execExpr.h"
 #include "executor/spi.h"
 #include "executor/tstoreReceiver.h"
@@ -323,6 +324,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,
@@ -2116,6 +2119,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;
@@ -4999,6 +5006,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 eeb7c4d7c0..1b39351db7 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";
@@ -370,6 +372,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);
 
 
@@ -459,6 +462,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;
@@ -713,6 +719,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)
 {
@@ -815,6 +827,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);
 
 
@@ -914,6 +927,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;
@@ -1590,6 +1606,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 11714430c1..9eb3e6ebb5 100644
--- a/src/pl/plpgsql/src/pl_gram.y
+++ b/src/pl/plpgsql/src/pl_gram.y
@@ -200,7 +200,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
@@ -307,6 +307,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
@@ -876,6 +877,8 @@ proc_stmt		: pl_block ';'
 						{ $$ = $1; }
 				| stmt_rollback
 						{ $$ = $1; }
+				| stmt_let
+						{ $$ = $1; }
 				;
 
 stmt_perform	: K_PERFORM
@@ -992,6 +995,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,
+													   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 d338e9f637..be94aa751a 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/pl_unreserved_kwlist.h b/src/pl/plpgsql/src/pl_unreserved_kwlist.h
index 670b4cf0c1..602628d7d2 100644
--- a/src/pl/plpgsql/src/pl_unreserved_kwlist.h
+++ b/src/pl/plpgsql/src/pl_unreserved_kwlist.h
@@ -69,6 +69,7 @@ PG_KEYWORD("info", K_INFO)
 PG_KEYWORD("insert", K_INSERT)
 PG_KEYWORD("is", K_IS)
 PG_KEYWORD("last", K_LAST)
+PG_KEYWORD("let", K_LET)
 PG_KEYWORD("log", K_LOG)
 PG_KEYWORD("merge", K_MERGE)
 PG_KEYWORD("message", K_MESSAGE)
diff --git a/src/pl/plpgsql/src/plpgsql.h b/src/pl/plpgsql/src/plpgsql.h
index 50c3b28472..bb400629ad 100644
--- a/src/pl/plpgsql/src/plpgsql.h
+++ b/src/pl/plpgsql/src/plpgsql.h
@@ -128,6 +128,7 @@ typedef enum PLpgSQL_stmt_type
 	PLPGSQL_STMT_CALL,
 	PLPGSQL_STMT_COMMIT,
 	PLPGSQL_STMT_ROLLBACK,
+	PLPGSQL_STMT_LET,
 } PLpgSQL_stmt_type;
 
 /*
@@ -520,6 +521,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/test/regress/expected/session_variables.out b/src/test/regress/expected/session_variables.out
index 908118fbb4..e0c07e0d92 100644
--- a/src/test/regress/expected/session_variables.out
+++ b/src/test/regress/expected/session_variables.out
@@ -1007,8 +1007,7 @@ ERROR:  permission denied for session variable var1
 CONTEXT:  PL/pgSQL function inline_code_block line 1 at RAISE
 DO $$ BEGIN LET var1.a = var1.a + 10; END $$;
 ERROR:  permission denied for session variable var1
-CONTEXT:  SQL statement "LET var1.a = var1.a + 10"
-PL/pgSQL function inline_code_block line 1 at SQL statement
+CONTEXT:  PL/pgSQL function inline_code_block line 1 at LET
 SET ROLE TO DEFAULT;
 GRANT SELECT ON VARIABLE var1 TO regress_var_test_role;
 SET ROLE TO regress_var_test_role;
@@ -1049,8 +1048,7 @@ ERROR:  permission denied for session variable var1
 CONTEXT:  PL/pgSQL function inline_code_block line 1 at RAISE
 DO $$ BEGIN LET var1.a = var1.a + 10; END $$;
 ERROR:  permission denied for session variable var1
-CONTEXT:  SQL statement "LET var1.a = var1.a + 10"
-PL/pgSQL function inline_code_block line 1 at SQL statement
+CONTEXT:  PL/pgSQL function inline_code_block line 1 at LET
 SET ROLE TO DEFAULT;
 DROP VARIABLE var1;
 DROP ROLE regress_var_test_role;
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 5c875ae684..7ae5b88feb 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -1893,6 +1893,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.47.1