Make local variable to evaluate dynamically (only if it appears on Ta…

…rgetList) (#3011)

Currently, declared local variables are assumed to be static by analyser and hence it is replaced with const node by
optimiser. Currently, the parameter value will be supplied through pltsql_param_fetch(..) which is setup during
pltsql_estate_setup as part of ParamListInfo setup. This hook is purely used during planning phase and is not used
during query execution.

If parameter value is required during execution then values will be read from estate->ndatums through
pltsql_param_eval_var callback for EEOP_PARAM_CALLBACK step. This estate->ndatums setup during
pltsql_estate_setup and values are copied through copy_pltsql_datums(…). And currently, any assignments to declare
variables are implemented using "select target" which will be executed at the end of executor. Hence, any use of
variables will read original (value copied during pltsql_estate_setup(...)) and it would not reflect any dynamic change in
the value of any variable.

In order to make this behaviour dynamic, this commit introduces internal function sys.pltsql_assign_var(dno, expr).
Argument dno is item number of declared variable and expr could be anything whose value will be assigned to declared
variable pointed by dno during execution. We will use this function to rewrite any variables assignment operation during
ANTLR parsing as stated in following example,

select @A = expr

Will be written to

select @A = sys.pltsql_assign_var(dno, expr)

Internally, sys.pltsql_assign_var invokes exec_assign_value(…) to assign appropriate value directly to variable by
updating estate->ndatums array and hence any subsequent read will read latest value.

Engine PR: babelfish-for-postgresql/postgresql_modified_for_babelfish#469
Task: BABEL-5325
Signed-off-by: Dipesh Dhameliya <[email protected]>

contrib/babelfishpg_tds/src/backend/tds/tdsresponse.c

@@ -543,6 +543,19 @@ resolve_numeric_typmod_from_exp(Plan *plan, Node *expr)
 		return -1;
 	switch (nodeTag(expr))
 	{
+		case T_Param:
+			{
+				Param *param = (Param *) expr;
+				if (!is_numeric_datatype(param->paramtype))
+				{
+					/* typmod is undefined */
+					return -1;
+				}
+				else
+				{
+					return param->paramtypmod;
+				}
+			}
 		case T_Const:
 			{
 				Const	   *con = (Const *) expr;

contrib/babelfishpg_tsql/sql/sys_functions.sql

@@ -5416,3 +5416,7 @@ CREATE OR REPLACE AGGREGATE sys.string_agg(sys.NVARCHAR, sys.VARCHAR) (
     PARALLEL = SAFE
 );
 
+/* Helper function to update local variables dynamically during execution */
+CREATE OR REPLACE FUNCTION sys.pltsql_assign_var(dno INT, val ANYELEMENT)
+RETURNS ANYELEMENT
+AS 'babelfishpg_tsql', 'pltsql_assign_var' LANGUAGE C PARALLEL UNSAFE;

contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--4.4.0--4.5.0.sql

@@ -11,6 +11,11 @@ SELECT set_config('search_path', 'sys, '||current_setting('search_path'), false)
  * final behaviour.
  */
 
+/* Helper function to update local variables dynamically during execution */
+CREATE OR REPLACE FUNCTION sys.pltsql_assign_var(dno INT, val ANYELEMENT)
+RETURNS ANYELEMENT
+AS 'babelfishpg_tsql', 'pltsql_assign_var' LANGUAGE C PARALLEL UNSAFE;
+
 -- This is a temporary procedure which is only meant to be called during upgrade
 CREATE OR REPLACE PROCEDURE sys.babelfish_revoke_guest_from_mapped_logins()
 LANGUAGE C

contrib/babelfishpg_tsql/src/collation.c

@@ -1136,6 +1136,16 @@ pltsql_planner_node_transformer(PlannerInfo *root,
 								Node *expr,
 								int kind)
 {
+	/*
+	 * check if this is called to reset saved expression kind. Quickly return if so.
+	 */
+	if (kind == -1)
+	{
+		Assert(expr == NULL);
+		saved_expr_kind = -1;
+		return NULL;
+	}
+
 	/*
 	 * Fall out quickly if expression is empty.
 	 */
@@ -1144,6 +1154,7 @@ pltsql_planner_node_transformer(PlannerInfo *root,
 
 	if (EXPRKIND_TARGET == kind)
 	{
+		saved_expr_kind = EXPRKIND_TARGET;
 		/*
 		 * If expr is NOT a Boolean expression then recurse through its
 		 * expresion tree

contrib/babelfishpg_tsql/src/pl_exec.c

@@ -62,14 +62,14 @@
 #include "guc.h"
 #include "multidb.h"
 #include "session.h"
-#include "guc.h"
 #include "catalog.h"
 
 uint64		rowcount_var = 0;
 List	   *columns_updated_list = NIL;
 static char *original_query_string = NULL;
 
 int			fetch_status_var = 0;
+int			saved_expr_kind = -1;
 
 typedef struct
 {
@@ -4368,6 +4368,7 @@ pltsql_estate_setup(PLtsql_execstate *estate,
 	pltsql_init_exec_error_data(&(es_cs_entry->error_data));
 	es_cs_entry->next = exec_state_call_stack;
 	exec_state_call_stack = es_cs_entry;
+	saved_expr_kind = -1;
 }
 
 /* ----------
@@ -7858,12 +7859,22 @@ pltsql_param_fetch(ParamListInfo params,
 		}
 	}
 
+	if (saved_expr_kind == EXPRKIND_TARGET)
+	{
+		/* Let extension to set value of param dynamically during execution when variables appears in TargetList */
+		prm->pflags = 0;
+	}
+	else
+	{
+		/* For other cases, for example, Quals, we can always mark params as "const" for executor's purposes */
+		prm->pflags = PARAM_FLAG_CONST;
+	}
+
 	/* Return "no such parameter" if not ok */
 	if (!ok)
 	{
 		prm->value = (Datum) 0;
 		prm->isnull = true;
-		prm->pflags = 0;
 		prm->ptype = InvalidOid;
 		return prm;
 	}
@@ -7872,8 +7883,6 @@ pltsql_param_fetch(ParamListInfo params,
 	exec_eval_datum(estate, datum,
 					&prm->ptype, &prmtypmod,
 					&prm->value, &prm->isnull);
-	/* We can always mark params as "const" for executor's purposes */
-	prm->pflags = PARAM_FLAG_CONST;
 
 	/*
 	 * If it's a read/write expanded datum, convert reference to read-only,
@@ -10428,3 +10437,38 @@ pltsql_exec_function_cleanup(PLtsql_execstate *estate, PLtsql_function *func, Er
 	}
 	PG_END_TRY();
 }
+
+PG_FUNCTION_INFO_V1(pltsql_assign_var);
+
+/*
+ * pltsql_assign_var - Helper function to update local variables dynamically during execution.
+ * Any statement which updates local variables as part of TargetList will be re-written using
+ * this function. for example,
+ * @var = expr will be re-written to @var=sys.pltsql_assign_var(dno, cast((expr) as type)).
+ */
+Datum
+pltsql_assign_var(PG_FUNCTION_ARGS)
+{
+	int dno = PG_GETARG_INT32(0);
+	Datum data = PG_GETARG_DATUM(1);
+	Oid valtype = get_fn_expr_argtype(fcinfo->flinfo, 1);
+	bool isNull = PG_ARGISNULL(1);
+	int32 valtypmod = -1;
+	PLtsql_datum *target;
+	MemoryContext oldcontext;
+
+	PLtsql_execstate *estate = get_current_tsql_estate();
+	Assert(estate != NULL);
+	oldcontext = MemoryContextSwitchTo(estate->datum_context);
+	target = estate->datums[dno];
+
+	/* we will reuse exec_assign_value function here provided in pl_exec.c */
+	exec_assign_value(estate, target, data, isNull, valtype, valtypmod);
+
+	MemoryContextSwitchTo(oldcontext);
+
+	if (isNull)
+		PG_RETURN_NULL();
+
+	PG_RETURN_DATUM(data);
+}

contrib/babelfishpg_tsql/src/pltsql.h

@@ -2229,6 +2229,13 @@ extern bool pltsql_trace_exec_codes;
 extern bool pltsql_trace_exec_counts;
 extern bool pltsql_trace_exec_time;
 
+/*
+ * saved_expr_kind - special context to store which kind of expression if being processed.
+ * This is useful specially when handling declared variables because variables are dynamic only when it appears
+ * in the TargetList. Should be folded as const otherwise.
+ */
+extern int saved_expr_kind;
+
 /*
  * Functions in cursor.c
  */
@@ -2294,4 +2301,9 @@ extern void	exec_alter_role_cmd(char *query_str, RoleSpec *role);
 extern bool validate_special_function(char *proc_nsname, char *proc_name,  List* fargs, int nargs, Oid *input_typeids, bool num_args_match);
 extern void init_special_function_list(void);
 
+/*
+ * Function in pltsql_ruleutils.c
+ */
+extern char *tsql_format_type_extended(Oid type_oid, int32 typemod, bits16 flags);
+
 #endif							/* PLTSQL_H */

contrib/babelfishpg_tsql/src/pltsql_ruleutils.c

@@ -345,7 +345,6 @@ static char *tsql_get_constraintdef_worker(Oid constraintId, bool fullCommand,
 static text *tsql_get_expr_worker(text *expr, Oid relid, const char *relname,
 								  int prettyFlags);
 static char *tsql_printTypmod(const char *typname, int32 typmod, Oid typmodout);
-static char *tsql_format_type_extended(Oid type_oid, int32 typemod, bits16 flags);
 int			tsql_print_function_arguments(StringInfo buf, HeapTuple proctup,
 										  bool print_table_args, bool print_defaults, int **typmod_arr_arg, bool *has_tvp);
 char	   *tsql_quote_qualified_identifier(const char *qualifier, const char *ident);
@@ -2793,7 +2792,7 @@ find_recursive_union(deparse_namespace *dpns, WorkTableScan *wtscan)
  *
  * Returns a palloc'd string.
  */
-static char *
+char *
 tsql_format_type_extended(Oid type_oid, int32 typemod, bits16 flags)
 {
 	HeapTuple	tuple;