Make local variable to evaluate dynamically (only if it appears on TargetList)

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

@@ -5422,3 +5422,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

@@ -54,6 +54,11 @@ BEGIN
 END;
 $$;
 
+/* 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;
+
 
 -- After upgrade, always run analyze for all babelfish catalogs.
 CALL sys.analyze_babelfish_catalogs();

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

@@ -932,6 +932,11 @@ END;
 $$;
 GRANT EXECUTE on PROCEDURE sys.sp_rename(IN sys.nvarchar(776), IN sys.SYSNAME, IN sys.varchar(13)) TO PUBLIC;
 
+/* 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;
+
 -- Drops the temporary procedure used by the upgrade script.
 -- Please have this be one of the last statements executed in this upgrade script.
 DROP PROCEDURE sys.babelfish_drop_deprecated_object(varchar, varchar, varchar);

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,
@@ -10423,3 +10432,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

@@ -2264,6 +2264,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
  */
@@ -2329,4 +2336,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;