From 40dfcdc1d9fc623ca6d38a818a2ad087e25b2265 Mon Sep 17 00:00:00 2001
From: Laurenz Albe <laurenz.albe@cybertec.at>
Date: Wed, 13 Nov 2024 15:07:03 +0100
Subject: [PATCH 16/22] allow read an value of session variable directly from
 expression executor

The expression executor can be called directly (not inside query executor)
when expression is evaluated as simple expression in plpgsql or when expression
is an argument of CALL or EXECUTE statements. For these cases we need to allow direct
access to content of session variables from executor. We can do it, because
we know so the expression is not evaluated under query and cannot be
evaluated inside parallel worker.

This patch enables session variables for simple expression evaluation.

This patch enables usage session variables in arguments of CALL stmt.
---
 src/backend/commands/prepare.c                | 10 ---
 src/backend/executor/execExpr.c               | 85 +++++++++++++++----
 src/backend/executor/execExprInterp.c         | 14 +++
 src/backend/jit/llvm/llvmjit_expr.c           |  6 ++
 src/backend/parser/analyze.c                  | 10 ---
 src/include/executor/execExpr.h               | 11 ++-
 .../src/expected/plpgsql_session_variable.out |  3 +-
 src/pl/plpgsql/src/pl_exec.c                  |  3 +-
 .../regress/expected/session_variables.out    | 41 ++++-----
 src/test/regress/sql/session_variables.sql    | 12 +--
 10 files changed, 125 insertions(+), 70 deletions(-)

diff --git a/src/backend/commands/prepare.c b/src/backend/commands/prepare.c
index 4ea2bfea4e..70f0309999 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -338,16 +338,6 @@ EvaluateParams(ParseState *pstate, PreparedStatement *pstmt, List *params,
 		i++;
 	}
 
-	/*
-	 * The arguments of EXECUTE are evaluated by a direct expression
-	 * executor call.  This mode doesn't support session variables yet.
-	 * It will be enabled later.
-	 */
-	if (pstate->p_hasSessionVariables)
-		ereport(ERROR,
-				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-				 errmsg("session variable cannot be used as an argument")));
-
 	/* Prepare the expressions for execution */
 	exprstates = ExecPrepareExprList(params, estate);
 
diff --git a/src/backend/executor/execExpr.c b/src/backend/executor/execExpr.c
index 1334842399..bcb425ea32 100644
--- a/src/backend/executor/execExpr.c
+++ b/src/backend/executor/execExpr.c
@@ -34,6 +34,7 @@
 #include "catalog/objectaccess.h"
 #include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
+#include "catalog/pg_variable.h"
 #include "commands/session_variable.h"
 #include "executor/execExpr.h"
 #include "executor/nodeSubplan.h"
@@ -1008,22 +1009,74 @@ ExecInitExprRec(Expr *node, ExprState *state,
 								es_num_session_variables = state->parent->state->es_num_session_variables;
 							}
 
-							Assert(es_session_variables);
-
-							/* parameter sanity checks */
-							if (param->paramid >= es_num_session_variables)
-								elog(ERROR, "paramid of PARAM_VARIABLE param is out of range");
-
-							var = &es_session_variables[param->paramid];
-
-							/*
-							 * In this case, pass the value like a
-							 * constant.
-							 */
-							scratch.opcode = EEOP_CONST;
-							scratch.d.constval.value = var->value;
-							scratch.d.constval.isnull = var->isnull;
-							ExprEvalPushStep(state, &scratch);
+							if (es_session_variables)
+							{
+								/*
+								 * This path is used when expression is
+								 * evaluated inside query evaluation. For
+								 * ensuring of immutability of session variable
+								 * inside query we use special buffer with copy
+								 * of used session variables. This method helps
+								 * with parallel execution too.
+								 */
+
+								/* parameter sanity checks */
+								if (param->paramid >= es_num_session_variables)
+									elog(ERROR, "paramid of PARAM_VARIABLE param is out of range");
+
+								var = &es_session_variables[param->paramid];
+
+								/*
+								 * In this case, pass the value like a
+								 * constant.
+								 */
+								scratch.opcode = EEOP_CONST;
+								scratch.d.constval.value = var->value;
+								scratch.d.constval.isnull = var->isnull;
+								ExprEvalPushStep(state, &scratch);
+							}
+							else
+							{
+								Oid			varid = param->paramvarid;
+
+								Assert(!IsParallelWorker());
+
+								/*
+								 * In some cases (plpgsql's simple expression
+								 * evaluation or evaluation of CALL arguments),
+								 * the expression executor is called directly.
+								 * We can allow direct access to session
+								 * variables (copy only), because we know, so
+								 * outside is not any query (and expression
+								 * cannot be executed parallel).
+								 *
+								 * In this case we should to do aclcheck,
+								 * because usual aclcheck from
+								 * standard_ExecutorStart is not executed in
+								 * this case. Fortunately it is just once per
+								 * transaction.
+								 *
+								 * Don't do permission check, when variable is
+								 * used like base node for assignment indirection.
+								 */
+								if (!param->parambasenode)
+								{
+									AclResult	aclresult;
+
+									aclresult = object_aclcheck(
+													  VariableRelationId, varid,
+													  GetUserId(), ACL_SELECT);
+
+									if (aclresult != ACLCHECK_OK)
+										aclcheck_error(aclresult,
+													   OBJECT_VARIABLE,
+													   get_session_variable_name(varid));
+								}
+
+								scratch.opcode = EEOP_PARAM_VARIABLE;
+								scratch.d.vparam.varid = varid;
+								ExprEvalPushStep(state, &scratch);
+							}
 						}
 						break;
 
diff --git a/src/backend/executor/execExprInterp.c b/src/backend/executor/execExprInterp.c
index b2c00a0a1b..bfda74dce0 100644
--- a/src/backend/executor/execExprInterp.c
+++ b/src/backend/executor/execExprInterp.c
@@ -59,6 +59,7 @@
 #include "access/heaptoast.h"
 #include "catalog/pg_type.h"
 #include "commands/sequence.h"
+#include "commands/session_variable.h"
 #include "executor/execExpr.h"
 #include "executor/nodeSubplan.h"
 #include "funcapi.h"
@@ -512,6 +513,7 @@ ExecInterpExpr(ExprState *state, ExprContext *econtext, bool *isnull)
 		&&CASE_EEOP_PARAM_EXEC,
 		&&CASE_EEOP_PARAM_EXTERN,
 		&&CASE_EEOP_PARAM_CALLBACK,
+		&&CASE_EEOP_PARAM_VARIABLE,
 		&&CASE_EEOP_PARAM_SET,
 		&&CASE_EEOP_CASE_TESTVAL,
 		&&CASE_EEOP_MAKE_READONLY,
@@ -1161,6 +1163,18 @@ ExecInterpExpr(ExprState *state, ExprContext *econtext, bool *isnull)
 			EEO_NEXT();
 		}
 
+		EEO_CASE(EEOP_PARAM_VARIABLE)
+		{
+			/*
+			 * Direct access to session variable (without buffering). Because
+			 * returned value can be used (without an assignement) after the
+			 * referenced session variables is updated, we have to use an copy
+			 * of stored value every time.
+			 */
+			*op->resvalue = GetSessionVariable(op->d.vparam.varid, op->resnull);
+			EEO_NEXT();
+		}
+
 		EEO_CASE(EEOP_PARAM_SET)
 		{
 			/* out of line, unlikely to matter performance-wise */
diff --git a/src/backend/jit/llvm/llvmjit_expr.c b/src/backend/jit/llvm/llvmjit_expr.c
index 6c4915e373..0b73c5f354 100644
--- a/src/backend/jit/llvm/llvmjit_expr.c
+++ b/src/backend/jit/llvm/llvmjit_expr.c
@@ -1136,6 +1136,12 @@ llvm_compile_expr(ExprState *state)
 					break;
 				}
 
+			case EEOP_PARAM_VARIABLE:
+				build_EvalXFunc(b, mod, "ExecEvalParamVariable",
+								v_state, op, v_econtext);
+				LLVMBuildBr(b, opblocks[opno + 1]);
+				break;
+
 			case EEOP_PARAM_SET:
 				build_EvalXFunc(b, mod, "ExecEvalParamSet",
 								v_state, op, v_econtext);
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index a03860d4b8..b8df16e4f0 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -3469,16 +3469,6 @@ transformCallStmt(ParseState *pstate, CallStmt *stmt)
 							 true,
 							 stmt->funccall->location);
 
-	/*
-	 * The arguments of CALL statement are evaluated by a direct expression
-	 * executor call.  This path is unsupported yet, so block it.  It will be
-	 * enabled later.
-	 */
-	if (pstate->p_hasSessionVariables)
-		ereport(ERROR,
-				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-				 errmsg("session variable cannot be used as an argument")));
-
 	assign_expr_collations(pstate, node);
 
 	fexpr = castNode(FuncExpr, node);
diff --git a/src/include/executor/execExpr.h b/src/include/executor/execExpr.h
index 8019e5490e..fadcde1ed6 100644
--- a/src/include/executor/execExpr.h
+++ b/src/include/executor/execExpr.h
@@ -156,10 +156,11 @@ typedef enum ExprEvalOp
 	EEOP_BOOLTEST_IS_FALSE,
 	EEOP_BOOLTEST_IS_NOT_FALSE,
 
-	/* evaluate PARAM_EXEC/EXTERN parameters */
+	/* evaluate PARAM_EXEC/EXTERN/VARIABLE parameters */
 	EEOP_PARAM_EXEC,
 	EEOP_PARAM_EXTERN,
 	EEOP_PARAM_CALLBACK,
+	EEOP_PARAM_VARIABLE,
 	/* set PARAM_EXEC value */
 	EEOP_PARAM_SET,
 
@@ -410,6 +411,12 @@ typedef struct ExprEvalStep
 			Oid			paramtype;	/* OID of parameter's datatype */
 		}			cparam;
 
+		/* for EEOP_PARAM_VARIABLE */
+		struct
+		{
+			Oid			varid;		/* OID of assigned variable */
+		}			vparam;
+
 		/* for EEOP_CASE_TESTVAL/DOMAIN_TESTVAL */
 		struct
 		{
@@ -832,6 +839,8 @@ extern void ExecEvalParamSet(ExprState *state, ExprEvalStep *op,
 							 ExprContext *econtext);
 extern void ExecEvalParamExtern(ExprState *state, ExprEvalStep *op,
 								ExprContext *econtext);
+extern void ExecEvalParamVariable(ExprState *state, ExprEvalStep *op,
+								  ExprContext *econtext);
 extern void ExecEvalCoerceViaIOSafe(ExprState *state, ExprEvalStep *op);
 extern void ExecEvalSQLValueFunction(ExprState *state, ExprEvalStep *op);
 extern void ExecEvalCurrentOfExpr(ExprState *state, ExprEvalStep *op);
diff --git a/src/pl/plpgsql/src/expected/plpgsql_session_variable.out b/src/pl/plpgsql/src/expected/plpgsql_session_variable.out
index a6523f7afe..ecac64fcfb 100644
--- a/src/pl/plpgsql/src/expected/plpgsql_session_variable.out
+++ b/src/pl/plpgsql/src/expected/plpgsql_session_variable.out
@@ -211,8 +211,7 @@ SET ROLE TO regress_var_exec_role;
 -- should fail
 SELECT var_read_func();
 ERROR:  permission denied for session variable plpgsql_sv_var1
-CONTEXT:  PL/pgSQL expression "plpgsql_sv_var1"
-PL/pgSQL function var_read_func() line 3 at RAISE
+CONTEXT:  PL/pgSQL function var_read_func() line 3 at RAISE
 SET ROLE TO DEFAULT;
 SET ROLE TO regress_var_owner_role;
 GRANT SELECT ON VARIABLE plpgsql_sv_var1 TO regress_var_reader_role;
diff --git a/src/pl/plpgsql/src/pl_exec.c b/src/pl/plpgsql/src/pl_exec.c
index 3f901cd8fc..e5b0da04e3 100644
--- a/src/pl/plpgsql/src/pl_exec.c
+++ b/src/pl/plpgsql/src/pl_exec.c
@@ -8099,8 +8099,7 @@ exec_is_simple_query(PLpgSQL_expr *expr)
 		query->sortClause ||
 		query->limitOffset ||
 		query->limitCount ||
-		query->setOperations ||
-		query->hasSessionVariables)
+		query->setOperations)
 		return false;
 
 	/*
diff --git a/src/test/regress/expected/session_variables.out b/src/test/regress/expected/session_variables.out
index 09bb4b2fb2..b5ab05840a 100644
--- a/src/test/regress/expected/session_variables.out
+++ b/src/test/regress/expected/session_variables.out
@@ -483,8 +483,7 @@ ERROR:  permission denied for session variable var1
 CONTEXT:  SQL function "sqlfx" statement 1
 SELECT plpgsqlfx(20);
 ERROR:  permission denied for session variable var1
-CONTEXT:  PL/pgSQL expression "$1 + var1"
-PL/pgSQL function plpgsqlfx(integer) line 1 at RETURN
+CONTEXT:  PL/pgSQL function plpgsqlfx(integer) line 1 at RETURN
 -- should be ok
 SELECT sqlfx_sd(20);
  sqlfx_sd 
@@ -780,27 +779,28 @@ $$;
 NOTICE:  result array: {1,2,3,4,5,6,7,8,9,10}
 DROP VARIABLE var1;
 DROP VARIABLE var2;
--- CALL statement is not supported yet
--- requires direct access to session variable from expression executor
+-- CALL statement is supported
 CREATE VARIABLE v int;
+LET v = 1;
 CREATE PROCEDURE p(arg int) AS $$ BEGIN RAISE NOTICE '%', arg; END $$ LANGUAGE plpgsql;
--- should not crash (but is not supported yet)
+-- should to work
 CALL p(v);
-ERROR:  session variable cannot be used as an argument
+NOTICE:  1
 DO $$ BEGIN CALL p(v); END $$;
-ERROR:  session variable cannot be used as an argument
-CONTEXT:  SQL statement "CALL p(v)"
-PL/pgSQL function inline_code_block line 1 at CALL
+NOTICE:  1
 DROP PROCEDURE p(int);
 DROP VARIABLE v;
--- EXECUTE statement is not supported yet
--- requires direct access to session variable from expression executor
+-- EXECUTE statement
 CREATE VARIABLE v int;
 LET v = 20;
 PREPARE ptest(int) AS SELECT $1;
--- should fail
+-- should be ok
 EXECUTE ptest(v);
-ERROR:  session variable cannot be used as an argument
+ ?column? 
+----------
+       20
+(1 row)
+
 DEALLOCATE ptest;
 DROP VARIABLE v;
 -- test search path
@@ -1123,12 +1123,10 @@ LET var1.a = (SELECT * FROM (SELECT count(*) FROM vartesttab WHERE xcol = var1.a
 ERROR:  permission denied for session variable var1
 DO $$ BEGIN RAISE NOTICE '%', var1; END $$;
 ERROR:  permission denied for session variable var1
-CONTEXT:  PL/pgSQL expression "var1"
-PL/pgSQL function inline_code_block line 1 at RAISE
+CONTEXT:  PL/pgSQL function inline_code_block line 1 at RAISE
 DO $$ BEGIN RAISE NOTICE '%', var1.a; END $$;
 ERROR:  permission denied for session variable var1
-CONTEXT:  PL/pgSQL expression "var1.a"
-PL/pgSQL function inline_code_block line 1 at RAISE
+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"
@@ -1170,12 +1168,10 @@ LET var1.a = (SELECT * FROM (SELECT count(*) FROM vartesttab WHERE xcol = var1.a
 ERROR:  permission denied for session variable var1
 DO $$ BEGIN RAISE NOTICE '%', var1; END $$;
 ERROR:  permission denied for session variable var1
-CONTEXT:  PL/pgSQL expression "var1"
-PL/pgSQL function inline_code_block line 1 at RAISE
+CONTEXT:  PL/pgSQL function inline_code_block line 1 at RAISE
 DO $$ BEGIN RAISE NOTICE '%', var1.a; END $$;
 ERROR:  permission denied for session variable var1
-CONTEXT:  PL/pgSQL expression "var1.a"
-PL/pgSQL function inline_code_block line 1 at RAISE
+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"
@@ -2080,8 +2076,7 @@ $$;
 ERROR:  session variable "public.var1" is not used inside variable fence
 DETAIL:  There is a risk of unwanted usage of session variable.
 HINT:  Use variable fence "VARIABLE(varname) for access to variable".
-CONTEXT:  PL/pgSQL expression "var1"
-PL/pgSQL function inline_code_block line 4 at RAISE
+CONTEXT:  PL/pgSQL function inline_code_block line 4 at RAISE
 -- should be ok
 SELECT VARIABLE(var1), VARIABLE(var1);
  var1 | var1 
diff --git a/src/test/regress/sql/session_variables.sql b/src/test/regress/sql/session_variables.sql
index b83f2a6bbb..8a122aeb1e 100644
--- a/src/test/regress/sql/session_variables.sql
+++ b/src/test/regress/sql/session_variables.sql
@@ -548,13 +548,14 @@ $$;
 DROP VARIABLE var1;
 DROP VARIABLE var2;
 
--- CALL statement is not supported yet
--- requires direct access to session variable from expression executor
+-- CALL statement is supported
 CREATE VARIABLE v int;
 
+LET v = 1;
+
 CREATE PROCEDURE p(arg int) AS $$ BEGIN RAISE NOTICE '%', arg; END $$ LANGUAGE plpgsql;
 
--- should not crash (but is not supported yet)
+-- should to work
 CALL p(v);
 
 DO $$ BEGIN CALL p(v); END $$;
@@ -562,13 +563,12 @@ DO $$ BEGIN CALL p(v); END $$;
 DROP PROCEDURE p(int);
 DROP VARIABLE v;
 
--- EXECUTE statement is not supported yet
--- requires direct access to session variable from expression executor
+-- EXECUTE statement
 CREATE VARIABLE v int;
 LET v = 20;
 PREPARE ptest(int) AS SELECT $1;
 
--- should fail
+-- should be ok
 EXECUTE ptest(v);
 
 DEALLOCATE ptest;
-- 
2.48.0

