Addressed PR comments and fixed regression test failures

Task: BABEL-6037
Signed-off-by: Manisha Deshpande <mmdeshp@amazon.com>

Author: manisha-deshpande

contrib/babelfishpg_tsql/sql/ownership.sql

@@ -41,8 +41,8 @@ CREATE TABLE sys.babelfish_function_ext (
 	create_date SYS.DATETIME NOT NULL,
 	modify_date SYS.DATETIME NOT NULL,
 	definition sys.NTEXT DEFAULT NULL,
-	antlr_parse_tree_text TEXT DEFAULT NULL, -- Native PG nodeToString() serialized parse tree [TODO: convert to BINARY TEXT]
-	antlr_parse_tree_datums TEXT DEFAULT NULL, -- Native PG nodeToString() serialized datums array [TODO: convert to BINARY TEXT]
+	antlr_parse_tree_text TEXT DEFAULT NULL, -- Native PG nodeToString() serialized parse tree
+	antlr_parse_tree_datums TEXT DEFAULT NULL, -- Native PG nodeToString() serialized datums array
   antlr_parse_tree_modify_date SYS.DATETIME DEFAULT NULL,
 	antlr_parse_tree_bbf_version TEXT DEFAULT NULL,
 	PRIMARY KEY(funcname, nspname, funcsignature)

contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--5.5.0--5.6.0.sql

@@ -232,28 +232,11 @@ $$
 $$;
 
 -- BABELFISH_FUNCTION_EXT
-DROP TABLE IF EXISTS sys.babelfish_function_ext;
-
-CREATE TABLE sys.babelfish_function_ext (
- nspname NAME NOT NULL,
- funcname NAME NOT NULL,
- orig_name sys.NVARCHAR(128), -- original input name of users
- funcsignature TEXT NOT NULL COLLATE "C",
- default_positions TEXT COLLATE "C",
- flag_validity BIGINT,
- flag_values BIGINT,
- create_date SYS.DATETIME NOT NULL,
- modify_date SYS.DATETIME NOT NULL,
- definition sys.NTEXT DEFAULT NULL,
- antlr_parse_tree_text TEXT DEFAULT NULL, -- Native PG nodeToString() serialized parse tree
- antlr_parse_tree_datums TEXT DEFAULT NULL, -- Native PG nodeToString() serialized datums array
- antlr_parse_tree_modify_date SYS.DATETIME DEFAULT NULL,
- antlr_parse_tree_bbf_version TEXT DEFAULT NULL,
- PRIMARY KEY(funcname, nspname, funcsignature)
-);
-GRANT SELECT ON sys.babelfish_function_ext TO PUBLIC;
+ALTER TABLE sys.babelfish_function_ext ADD COLUMN IF NOT EXISTS antlr_parse_tree_text TEXT DEFAULT NULL;
+ALTER TABLE sys.babelfish_function_ext ADD COLUMN IF NOT EXISTS antlr_parse_tree_datums TEXT DEFAULT NULL;
+ALTER TABLE sys.babelfish_function_ext ADD COLUMN IF NOT EXISTS antlr_parse_tree_modify_date SYS.DATETIME DEFAULT NULL;
+ALTER TABLE sys.babelfish_function_ext ADD COLUMN IF NOT EXISTS antlr_parse_tree_bbf_version TEXT DEFAULT NULL;
 
-SELECT pg_catalog.pg_extension_config_dump('sys.babelfish_function_ext', '');
 -- Please add your SQLs here
 
 -- Deprecate and drop old aggregates first (they depend on the function)

contrib/babelfishpg_tsql/src/guc.c

@@ -72,7 +72,7 @@ char	   *pltsql_host_service_pack_level = NULL;
 
 bool		pltsql_enable_create_alter_view_from_pg = false;
 bool		pltsql_enable_alter_owner_from_pg = false;
-bool		pltsql_enable_procedure_parse_cache = false;
+bool		pltsql_enable_routine_parse_cache = false;
 
 static const struct config_enum_entry explain_format_options[] = {
 	{"text", EXPLAIN_FORMAT_TEXT, false},
@@ -1144,10 +1144,10 @@ define_custom_variables(void)
 	/*
 	 * Enable/disable procedure ANTLR parse result caching for cross-session (sys.babelfish_function_ext->antlr_parse_tree) performance optimization.
 	 */
-	DefineCustomBoolVariable("babelfishpg_tsql.enable_procedure_parse_cache",
+	DefineCustomBoolVariable("babelfishpg_tsql.enable_routine_parse_cache",
 							 gettext_noop("Enables caching of ANTLR parse results for stored procedures across sessions"),
 							 NULL,
-							 &pltsql_enable_procedure_parse_cache,
+							 &pltsql_enable_routine_parse_cache,
 							 false,
 							 PGC_USERSET,
 							 GUC_NOT_IN_SAMPLE | GUC_DISALLOW_IN_FILE | GUC_DISALLOW_IN_AUTO_FILE,

contrib/babelfishpg_tsql/src/guc.h

@@ -18,7 +18,7 @@ typedef enum IsolationOptions
 extern bool pltsql_fmtonly;
 extern bool pltsql_enable_create_alter_view_from_pg;
 extern bool	pltsql_enable_alter_owner_from_pg;
-extern bool pltsql_enable_procedure_parse_cache;
+extern bool pltsql_enable_routine_parse_cache;
 extern bool pltsql_enable_linked_servers;
 extern bool pltsql_enable_ownership_chaining;
 extern bool pltsql_allow_windows_login;

contrib/babelfishpg_tsql/src/hooks.c

@@ -4335,7 +4335,7 @@ pltsql_store_func_default_positions(ObjectAddress address, List *parameters, con
 		PLtsql_function *function = NULL;
 
 		/* Check if procedure parse caching is enabled via GUC */
-		if (!pltsql_enable_procedure_parse_cache)
+		if (!pltsql_enable_routine_parse_cache)
 		{
 			/* GUC disabled - skip cache storage, set columns to NULL */
 			elog(DEBUG1, "Parse result caching skipped for function %u: GUC disabled",
@@ -4566,11 +4566,12 @@ pltsql_update_func_cache_entry(HeapTuple proctup, PLtsql_function *function)
 	bool		new_record_nulls[BBF_FUNCTION_EXT_NUM_COLS];
 	bool		new_record_replaces[BBF_FUNCTION_EXT_NUM_COLS];
 	Datum		modify_date;
+	bool		snapshot_pushed = false;
 
 	if (babelfish_dump_restore)
 		return;
 
-	if (!pltsql_enable_procedure_parse_cache)
+	if (!pltsql_enable_routine_parse_cache)
 		return;
 
 	if (!OidIsValid(get_bbf_function_ext_idx_oid()))
@@ -4580,6 +4581,17 @@ pltsql_update_func_cache_entry(HeapTuple proctup, PLtsql_function *function)
 	if (!HeapTupleIsValid(oldtup))
 		return;
 
+	/*
+	 * CatalogTupleUpdate requires an active snapshot. During after-trigger
+	 * execution the snapshot may already have been popped, so push one if
+	 * needed. This mirrors the pattern at pltsql_store_func_default_positions.
+	 */
+	if (!ActiveSnapshotSet())
+	{
+		PushActiveSnapshot(GetTransactionSnapshot());
+		snapshot_pushed = true;
+	}
+
 	rel = table_open(get_bbf_function_ext_oid(), RowExclusiveLock);
 	dsc = RelationGetDescr(rel);
 
@@ -4604,6 +4616,9 @@ pltsql_update_func_cache_entry(HeapTuple proctup, PLtsql_function *function)
 	heap_freetuple(oldtup);
 	heap_freetuple(newtup);
 	table_close(rel, RowExclusiveLock);
+
+	if (snapshot_pushed)
+		PopActiveSnapshot();
 }
 
 /*
@@ -4615,7 +4630,7 @@ pltsql_update_func_cache_entry(HeapTuple proctup, PLtsql_function *function)
  * expensive ANTLR parsing entirely for first time execution of routines in new sessions.
  *
  * Validation checks (all must pass for a cache hit):
- *   1. GUC babelfishpg_tsql.enable_procedure_parse_cache is enabled
+ *   1. GUC babelfishpg_tsql.enable_routine_parse_cache is enabled
  *   2. bbf_version matches current BABELFISH_VERSION_STR (detects MVU)
  *   3. antlr_parse_tree_modify_date >= modify_date (detects guc-disabled ALTER after caching)
  *   4. antlr_parse_tree_text is not NULL and deserializes successfully
@@ -4641,9 +4656,9 @@ pltsql_restore_func_parse_result(HeapTuple proctup)
 		return NULL;
 
 	/* Check if procedure parse caching is enabled via GUC */
-	if (!pltsql_enable_procedure_parse_cache)
+	if (!pltsql_enable_routine_parse_cache)
 	{
-		elog(DEBUG1, "Parse result cache retrieval skipped: GUC babelfishpg_tsql.enable_procedure_parse_cache is disabled");
+		elog(DEBUG1, "Parse result cache retrieval skipped: GUC babelfishpg_tsql.enable_routine_parse_cache is disabled");
 		return NULL;
 	}
 

contrib/babelfishpg_tsql/src/pl_comp.c

@@ -930,6 +930,8 @@ do_compile(FunctionCallInfo fcinfo,
 
 			if (cached_result)
 			{
+				int		ci;
+
 				elog(LOG, "do_compile: Using cached parse result (ndatums=%d), skipping ANTLR parsing", cached_result->ndatums);
 				pltsql_parse_result = cached_result->parse_tree;
 
@@ -939,12 +941,110 @@ do_compile(FunctionCallInfo fcinfo,
 
 				/* Mark that this function was loaded from persistent cache */
 				function->from_cache = true;
-				
+
+				/*
+				 * Re-derive found_varno and fetch_status_varno from cached
+				 * datums by scanning refnames.
+				 *
+				 * The cache was serialized from the validator-compiled function
+				 * which may have a different datum layout than what do_compile
+				 * built above (e.g., for ITVFs the validator serializes BEFORE
+				 * pg_proc is updated with TABLE-mode columns, so the cache has
+				 * fewer datums than the runtime parameter loop creates).
+				 * Scanning by refname makes us layout-independent.
+				 *
+				 * in_arg_varnos is NOT re-derived here — the parameter loop
+				 * above already set it correctly from pg_proc metadata, and
+				 * scanning '@'-prefixed refnames would also match local
+				 * variables, overflowing the pronargs-sized buffer.
+				 */
+				function->found_varno = -1;
+				function->fetch_status_varno = -1;
+
+				for (ci = 0; ci < cached_result->ndatums; ci++)
+				{
+					PLtsql_datum *d = cached_result->datums[ci];
+
+					if (d->dtype == PLTSQL_DTYPE_VAR)
+					{
+						PLtsql_var *v = (PLtsql_var *) d;
+
+						if (v->refname && strcmp(v->refname, "found") == 0)
+							function->found_varno = d->dno;
+						else if (v->refname && strcmp(v->refname, "@@fetch_status") == 0)
+							function->fetch_status_varno = d->dno;
+					}
+				}
+
+				/*
+				 * Do NOT re-derive in_arg_varnos from cached datums.
+				 * The parameter loop above already populated in_arg_varnos
+				 * correctly from pg_proc metadata. Re-scanning cached datums
+				 * for '@'-prefixed refnames would also match local variables
+				 * (e.g. @sql, @precision), overflowing the in_arg_varnos
+				 * buffer which is sized for pronargs only.
+				 */
+
+				Assert(function->found_varno >= 0);
+				Assert(function->fetch_status_varno >= 0);
+
+				/*
+				 * Re-derive out_param_varno from cached datums.
+				 *
+				 * The validator already built the OUT row (PLtsql_row) and it was
+				 * serialized into the cache. The cached parse tree's RETURN nodes
+				 * reference it by dno. We just need to find it and set
+				 * function->out_param_varno so the runtime can use it.
+				 *
+				 * For multiple OUT args (or procedure with any OUT): find the
+				 * PLtsql_row datum in the cache, then rebuild its rowtupdesc
+				 * (which is not serialized — marked read_write_ignore).
+				 * For single OUT arg (non-procedure function): point directly
+				 * to the matching cached datum using the dno from the parameter
+				 * loop (which assigned the same sequential dnos as the validator).
+				 */
+				if (num_out_args > 1 ||
+					(num_out_args == 1 && function->fn_prokind == PROKIND_PROCEDURE))
+				{
+					for (ci = 0; ci < cached_result->ndatums; ci++)
+					{
+						if (cached_result->datums[ci]->dtype == PLTSQL_DTYPE_ROW)
+						{
+							PLtsql_row *row = (PLtsql_row *) cached_result->datums[ci];
+							int fi;
+
+							function->out_param_varno = row->dno;
+
+							/*
+							 * Rebuild rowtupdesc — it is NULL after deserialization
+							 * (TupleDesc is not serializable). Walk the ROW's varnos
+							 * to find each member VAR's type info, same as
+							 * build_row_from_vars does.
+							 */
+							row->rowtupdesc = CreateTemplateTupleDesc(row->nfields);
+							for (fi = 0; fi < row->nfields; fi++)
+							{
+								PLtsql_var *fvar = (PLtsql_var *) cached_result->datums[row->varnos[fi]];
+
+								TupleDescInitEntry(row->rowtupdesc, fi + 1,
+												   row->fieldnames[fi],
+												   fvar->datatype->typoid,
+												   fvar->datatype->atttypmod,
+												   0);
+								TupleDescInitEntryCollation(row->rowtupdesc, fi + 1,
+															fvar->datatype->collation);
+							}
+							break;
+						}
+					}
+				}
+				else if (num_out_args == 1)
+					function->out_param_varno = out_arg_variables[0]->dno;
+
 				/* Free the wrapper structure (but not the data we transferred) */
 				pfree(cached_result);
 
 				parse_rc = 0;
-				// [TODO}: uncomment to execute deserialized parse result
 				goto skip_antlr_parsing;
 			}
 		}
@@ -976,8 +1076,9 @@ do_compile(FunctionCallInfo fcinfo,
 	/*
 	 * Multi-Statement Table-Valued Function: 1) Add a declare table statement
 	 * to the beginning 2) Add a return table statement to the end
+	 * Skip when restoring from cache — the cached parse tree already contains these.
 	 */
-	if (function->is_mstvf)
+	if (function->is_mstvf && !function->from_cache)
 	{
 		/*
 		 * ANTLR parser would return a stmt list like INIT->BLOCK, where BLOCK
@@ -994,9 +1095,11 @@ do_compile(FunctionCallInfo fcinfo,
 	 * control to fall off the end without an explicit RETURN statement. The
 	 * easiest way to implement this is to add a RETURN statement to the end
 	 * of the statement list during parsing.
+	 * Skip when restoring from cache — the cached parse tree already has the dummy return.
 	 */
-	if (num_out_args > 0 || function->fn_rettype == VOIDOID ||
-		function->fn_retset)
+	if (!function->from_cache &&
+		(num_out_args > 0 || function->fn_rettype == VOIDOID ||
+		function->fn_retset))
 		add_dummy_return(function);
 
 	/*
@@ -1031,7 +1134,7 @@ do_compile(FunctionCallInfo fcinfo,
 	 * handles cache writes for DDL. Calling both in the same transaction causes
 	 * "tuple already updated by self" errors.
 	 */
-	if (!forValidator && pltsql_enable_procedure_parse_cache && !function->from_cache)
+	if (!forValidator && pltsql_enable_routine_parse_cache && !function->from_cache)
 		pltsql_update_func_cache_entry(procTup, function);
 
 	/* Debug dump for completed functions */
@@ -3220,7 +3323,7 @@ pltsql_HashTableLookup(PLtsql_func_hashkey *func_key)
 	if (hentry)
 	{
 		/* If GUC is disabled and function was loaded from cache, return NULL to force re-parse */
-		if (!pltsql_enable_procedure_parse_cache && hentry->function->from_cache)
+		if (!pltsql_enable_routine_parse_cache && hentry->function->from_cache)
 		{
 			elog(DEBUG1, "pltsql_HashTableLookup: GUC disabled, invalidating cached function %u", (unsigned int) func_key->funcOid);
 			return NULL;

test/JDBC/expected/BABEL-6037-vu-cleanup.out

@@ -53,5 +53,17 @@ GO
 DROP TABLE IF EXISTS dbo.dep_test_table;
 GO
 
+-- Drop Test 11: Procedure with single OUT parameter
+DROP PROCEDURE IF EXISTS dbo.babel_6037_out_single;
+GO
+
+-- Drop Test 12: Procedure with multiple OUT parameters
+DROP PROCEDURE IF EXISTS dbo.babel_6037_out_multi;
+GO
+
+-- Drop Test 13: MSTVF
+DROP FUNCTION IF EXISTS dbo.babel_6037_mstvf;
+GO
+
 PRINT 'Cleanup completed successfully';
 GO

test/JDBC/expected/BABEL-6037-vu-prepare.out

@@ -1,7 +1,7 @@
 
 -- BABEL-6037 POC Test - Prepare
 -- Creates test procedures for parse tree serialization/deserialization testing
-SELECT set_config('babelfishpg_tsql.enable_procedure_parse_cache', 'on', false);
+SELECT set_config('babelfishpg_tsql.enable_routine_parse_cache', 'on', false);
 GO
 ~~START~~
 text
@@ -214,7 +214,48 @@ BEGIN
 END;
 GO
 
-SELECT set_config('babelfishpg_tsql.enable_procedure_parse_cache', 'off', false);
+-- Test 11: Procedure with single OUT parameter
+-- Tests out_param_varno re-derivation for single OUT arg on a procedure
+-- (validator builds PLtsql_row, serialized into cache).
+CREATE PROCEDURE dbo.babel_6037_out_single
+    @in_val INT,
+    @out_val INT OUT
+AS
+BEGIN
+    SET @out_val = @in_val * 10;
+END;
+GO
+
+-- Test 12: Procedure with multiple OUT parameters
+-- For procedures with >1 OUT args, the validator builds a PLtsql_row datum
+-- that gets serialized into the cache. The cache-hit path must find it
+-- and set function->out_param_varno.
+CREATE PROCEDURE dbo.babel_6037_out_multi
+    @in_val INT,
+    @out_val INT OUT,
+    @out_msg VARCHAR(100) OUT
+AS
+BEGIN
+    SET @out_val = @in_val * 2;
+    SET @out_msg = 'Result: ' + CAST(@out_val AS VARCHAR);
+END;
+GO
+
+-- Test 13: Multi-statement table-valued function (MSTVF)
+-- MSTVFs use out_param_varno at runtime (pl_exec-2.c:1445) to build the
+-- result table. The cache must preserve the ROW datum so the runtime can
+-- find it.
+CREATE FUNCTION dbo.babel_6037_mstvf(@in_val INT)
+RETURNS @result TABLE (id INT, val VARCHAR(50))
+AS
+BEGIN
+    INSERT INTO @result VALUES (@in_val, 'row1');
+    INSERT INTO @result VALUES (@in_val + 1, 'row2');
+    RETURN;
+END;
+GO
+
+SELECT set_config('babelfishpg_tsql.enable_routine_parse_cache', 'off', false);
 GO
 ~~START~~
 text