Added ANTLR parse tree persistent caching for procedures

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

Author: manisha-deshpande

contrib/babelfishpg_tsql/sql/ownership.sql

@@ -41,6 +41,7 @@ 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 JSONB DEFAULT NULL, -- JSONB serialized ANTLR parse tree for caching
 	PRIMARY KEY(funcname, nspname, funcsignature)
 );
 GRANT SELECT ON sys.babelfish_function_ext TO PUBLIC;

contrib/babelfishpg_tsql/src/catalog.h

@@ -277,7 +277,8 @@ typedef FormData_bbf_servers_def *Form_bbf_servers_def;
 #define Anum_bbf_function_ext_create_date 8
 #define Anum_bbf_function_ext_modify_date 9
 #define Anum_bbf_function_ext_definition 10
-#define BBF_FUNCTION_EXT_NUM_COLS 10
+#define Anum_bbf_function_ext_antlr_parse_tree 11
+#define BBF_FUNCTION_EXT_NUM_COLS 11
 #define FLAG_IS_ANSI_NULLS_ON (1<<0)
 #define FLAG_USES_QUOTED_IDENTIFIER (1<<1)
 #define FLAG_CREATED_WITH_RECOMPILE (1<<2)

contrib/babelfishpg_tsql/src/hooks.c

@@ -101,6 +101,7 @@
 #include "bbf_parallel_query.h"
 #include "extendedproperty.h"
 #include "utils/xml.h"
+#include "pltsql_serialize_jsonb.h"
 
 #ifdef USE_LIBXML
 #include <libxml/tree.h>
@@ -4102,6 +4103,86 @@ pltsql_store_func_default_positions(ObjectAddress address, List *parameters, con
 		new_record_nulls[Anum_bbf_function_ext_definition - 1] = true;
 	new_record_replaces[Anum_bbf_function_ext_default_positions - 1] = true;
 
+	/*
+	 * Store serialized ANTLR parse result for cross-session caching.
+	 * Look up the compiled function and serialize its parse tree.
+	 */
+	{
+		PLtsql_func_hashkey hashkey;
+		PLtsql_function *function = NULL;
+		char	   *serialized_parse_tree = NULL;
+		PLtsql_serial_status status;
+
+		/* [TODO]: complete hashkey logic
+		 * Build a hashkey to look up the function in the hash table.
+		 * Note: This uses a simplified hashkey that may not match all
+		 * function overloads perfectly (e.g., default parameter values).
+		 * This is acceptable since we're just trying to find the compiled
+		 * function to serialize its parse tree.
+		 */
+		MemSet(&hashkey, 0, sizeof(PLtsql_func_hashkey));
+		hashkey.funcOid = address.objectId;
+		hashkey.isTrigger = false;
+		hashkey.isEventTrigger = false;
+		hashkey.trigOid = 0;
+		hashkey.inputCollation = 0;
+
+		if (form_proctup->pronargs > 0)
+		{
+			/* Copy argument types from proc tuple */
+			memcpy(hashkey.argtypes, form_proctup->proargtypes.values,
+				   form_proctup->pronargs * sizeof(Oid));
+		}
+
+		/* Look up the compiled function in the hash table */
+		function = pltsql_HashTableLookup(&hashkey);
+
+		if (function != NULL && function->action != NULL)
+		{
+			//[TODO] Save engine version in additional column
+			/* Serialize the parse tree to JSON string */
+			serialized_parse_tree = pltsql_serialize_stmt_block_jsonb(function->action, &status);
+
+			if (status == PLTSQL_SERIAL_SUPPORTED && serialized_parse_tree != NULL)
+			{
+				/* Convert JSON string to JSONB datum and store directly */
+				new_record[Anum_bbf_function_ext_antlr_parse_tree - 1] = DirectFunctionCall1(jsonb_in, 
+																							   CStringGetDatum(serialized_parse_tree));
+				new_record_nulls[Anum_bbf_function_ext_antlr_parse_tree - 1] = false;
+				new_record_replaces[Anum_bbf_function_ext_antlr_parse_tree - 1] = true;
+				
+				elog(DEBUG1, "Stored cached parse result for function %u (%zu bytes JSON)",
+					 address.objectId, strlen(serialized_parse_tree));
+			}
+			else
+			{
+				/* Serialization failed or unsupported - skip caching */
+				if (status == PLTSQL_SERIAL_UNSUPPORTED)
+				{
+					elog(DEBUG1, "Parse result caching skipped for function %u: contains unsupported node types",
+						 address.objectId);
+				}
+				else if (status == PLTSQL_SERIAL_ERROR)
+				{
+					elog(DEBUG1, "Failed to serialize parse result for function %u",
+						 address.objectId);
+				}
+				/* Set to NULL */
+				new_record_nulls[Anum_bbf_function_ext_antlr_parse_tree - 1] = true;
+				new_record_replaces[Anum_bbf_function_ext_antlr_parse_tree - 1] = true;
+			}
+
+			if (serialized_parse_tree)
+				pfree(serialized_parse_tree);
+		}
+		else
+		{
+			/* Function not found in hash table or no action - set to NULL */
+			new_record_nulls[Anum_bbf_function_ext_antlr_parse_tree - 1] = true;
+			new_record_replaces[Anum_bbf_function_ext_antlr_parse_tree - 1] = true;
+		}
+	}
+
 	oldtup = get_bbf_function_tuple_from_proctuple(proctup);
 
 	if (HeapTupleIsValid(oldtup))
@@ -4139,6 +4220,81 @@ pltsql_store_func_default_positions(ObjectAddress address, List *parameters, con
 	table_close(bbf_function_ext_rel, RowExclusiveLock);
 }
 
+/*
+ * pltsql_restore_func_parse_result
+ *      Attempt to restore a cached parse result from babelfish_function_ext
+ *
+ * Retrieves a previously serialized ANTLR parse tree from the catalog and
+ * deserializes it back into a PLtsql_stmt_block structure. This allows
+ * procedures to skip ANTLR parsing on subsequent executions in new sessions.
+ *
+ * Parameters:
+ *   proctup - HeapTuple from pg_proc for the function
+ *
+ * Returns:
+ *   PLtsql_stmt_block * - Deserialized parse tree if cache hit
+ *   NULL - If no cache entry exists, deserialization fails, or during restore
+ */
+PLtsql_stmt_block *
+pltsql_restore_func_parse_result(HeapTuple proctup)
+{
+	HeapTuple	bbffunctuple;
+	Datum		cached_parse_tree_datum;
+	PLtsql_stmt_block *deserialized;
+	PLtsql_serial_status status;
+	bool		isnull;
+	Jsonb	   *jsonb_data;
+
+	/* Disallow during restore */
+	if (babelfish_dump_restore)
+		return NULL;
+
+	/* 1. Get babelfish_function_ext tuple using existing helper */
+	bbffunctuple = get_bbf_function_tuple_from_proctuple(proctup);
+	
+	if (!HeapTupleIsValid(bbffunctuple))
+	{
+		elog(DEBUG1, "pltsql_restore_func_parse_result: No babelfish_function_ext entry found");
+		return NULL;
+	}
+
+	/* 2. Fetch cached_parse_tree column (JSONB type) */
+	cached_parse_tree_datum = SysCacheGetAttr(
+		PROCNAMENSPSIGNATURE,
+		bbffunctuple,
+		Anum_bbf_function_ext_antlr_parse_tree,
+		&isnull
+	);
+
+
+	if (isnull)
+	{
+		elog(DEBUG1, "pltsql_restore_func_parse_result: No cached parse result in catalog");
+		return NULL;
+	}
+
+	/* 3. Get Jsonb pointer directly - detoasts automatically */
+	jsonb_data = DatumGetJsonbP(cached_parse_tree_datum);
+	
+	elog(DEBUG1, "pltsql_restore_func_parse_result: Found cached parse result");
+
+	/* 4. Deserialize directly from Jsonb - efficient, no string conversion */
+	deserialized = pltsql_deserialize_stmt_block_from_jsonb(jsonb_data, &status);
+				
+	if (status != PLTSQL_SERIAL_SUPPORTED)
+	{
+		elog(DEBUG1, "pltsql_restore_func_parse_result: Deserialization failed with status %d", status);
+		return NULL;
+	}
+
+	elog(DEBUG1, "pltsql_restore_func_parse_result: Deserialization succeeded (cmd_type=%d, lineno=%d)", 
+		 deserialized->cmd_type, deserialized->lineno);
+	
+	/* Return deserialized result */
+	heap_freetuple(bbffunctuple);
+	return deserialized;
+}
+
 /*
  * Update 'function_args' in 'sys.babelfish_schema_permissions' 
  */

contrib/babelfishpg_tsql/src/hooks.h

@@ -6,6 +6,7 @@
 #include "tcop/cmdtag.h"
 #include "utils/pg_locale.h"
 #include "utils/xml.h"
+#include "pltsql.h"
 
 extern IsExtendedCatalogHookType PrevIsExtendedCatalogHook;
 extern IsToastRelationHookType PrevIsToastRelationHook;
@@ -25,6 +26,7 @@ extern void pltsql_store_func_default_positions(ObjectAddress address,
                                                 const char *queryString,
                                                 int origname_location,
                                                 bool with_recompile);
+extern PLtsql_stmt_block *pltsql_restore_func_parse_result(HeapTuple proctup);
 extern void alter_bbf_schema_permissions_catalog(ObjectWithArgs *owa, 
                                                     List *parameters,
                                                     int objtypeInt);

contrib/babelfishpg_tsql/src/pl_comp.c

@@ -41,6 +41,8 @@
 
 #include "pltsql.h"
 #include "pltsql-2.h"
+#include "pltsql_serialize_jsonb.h"
+#include "hooks.h"
 #include "analyzer.h"
 #include "catalog.h"
 #include "codegen.h"
@@ -137,7 +139,7 @@ static void pltsql_resolve_polymorphic_argtypes(int numargs,
 												Oid *argtypes, char *argmodes,
 												Node *call_expr, bool forValidator,
 												const char *proname);
-static PLtsql_function *pltsql_HashTableLookup(PLtsql_func_hashkey *func_key);
+PLtsql_function *pltsql_HashTableLookup(PLtsql_func_hashkey *func_key);
 static void pltsql_HashTableInsert(PLtsql_function *function,
 								   PLtsql_func_hashkey *func_key);
 static void pltsql_HashTableDelete(PLtsql_function *function);
@@ -251,7 +253,7 @@ pltsql_compile(FunctionCallInfo fcinfo, bool forValidator)
 									 forValidator);
 
 		/*
-		 * Do the hard part.
+		 * Compile the function (will attempt to use cached parse result if available).
 		 */
 		function = do_compile(fcinfo, procTup, function,
 							  &hashkey, forValidator);
@@ -510,6 +512,14 @@ do_compile(FunctionCallInfo fcinfo,
 												 function->fn_input_collation,
 												 NULL);
 
+				// /* Debug logging for populateTaxInvDetails parameters */
+				// if (function->fn_signature && 
+				// 	strstr(function->fn_signature, "populatetaxinvdetails") != NULL)
+				// {
+				// 	elog(NOTICE, "Creating param %d: argtypeid=%u, typmod=%d, argdtype->typoid=%u, argdtype->typname=%s",
+				// 		 i, argtypeid, typmod, argdtype->typoid, argdtype->typname);
+				// }
+
 				/* Disallow pseudotype argument */
 				/* (note we already replaced polymorphic types) */
 				/* (build_variable would do this, but wrong message) */
@@ -907,7 +917,30 @@ do_compile(FunctionCallInfo fcinfo,
 	 * Now parse the function's text
 	 */
 	{
-		ANTLR_result result = antlr_parser_cpp(proc_source);
+		ANTLR_result result;
+		PLtsql_stmt_block *cached_parse_result = NULL;
+		
+		/*
+		 * Try to restore cached parse result from previous compilation.
+		 * This allows us to skip expensive ANTLR parsing.
+		 * Skip during validation (forValidator=true) as we're just checking syntax.
+		 */
+		if (!forValidator)
+		{
+			cached_parse_result = pltsql_restore_func_parse_result(procTup);
+			
+			if (cached_parse_result) //[TODO] Retrieve antlr parse tree engine version and compare with current engine version
+			{
+				elog(DEBUG1, "do_compile: Using cached parse result, skipping ANTLR parsing");
+				pltsql_parse_result = cached_parse_result;
+				parse_rc = 0;
+				// [TODO}: uncomment to execute deserialized parse result
+				// goto skip_antlr_parsing;
+			}
+		}
+		
+		/* No cache hit - do ANTLR parsing */
+		result = antlr_parser_cpp(proc_source);
 
 		if (result.success)
 		{
@@ -919,7 +952,8 @@ do_compile(FunctionCallInfo fcinfo,
 			parse_rc = 1;		/* invalid input */
 		}
 	}
-
+// skip_antlr_parsing:
+	/* Continue with normal flow */
 	if (parse_rc != 0)
 		elog(ERROR, "pltsql parser returned %d", parse_rc);
 	function->action = pltsql_parse_result;
@@ -3148,7 +3182,7 @@ pltsql_HashTableInit(void)
 								   HASH_ELEM | HASH_BLOBS);
 }
 
-static PLtsql_function *
+PLtsql_function *
 pltsql_HashTableLookup(PLtsql_func_hashkey *func_key)
 {
 	pltsql_HashEnt *hentry;

contrib/babelfishpg_tsql/src/pltsql.h

@@ -2083,6 +2083,7 @@ extern PLtsql_condition *pltsql_parse_err_condition(char *condname);
 extern void pltsql_adddatum(PLtsql_datum *newdatum);
 extern int	pltsql_add_initdatums(int **varnos);
 extern void pltsql_HashTableInit(void);
+extern PLtsql_function *pltsql_HashTableLookup(PLtsql_func_hashkey *func_key);
 extern void reset_cache(void);
 
 /*