Route INSERT EXEC flush through execute_batch and fix rows-affected for sp_executesql. Flush the buffered INSERT EXEC temp table into the target via execute_batch. Since the flush re-enters through the inline handler (which pushes its own estate), redirect table-variable resolution, the per-statement implicit-transaction decision, and ownership chaining back to the caller's estate via insert_exec_flush_estate, and report rows-affected from the DestReceiver's captured-row count.

Author: monk0707

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

@@ -2781,25 +2781,35 @@ StatementEnd_Internal(PLtsql_execstate *estate, PLtsql_stmt *stmt, bool error)
 		case PLTSQL_STMT_EXEC_BATCH:
 		case PLTSQL_STMT_EXEC_SP:
 			{
+				/* INSERT EXEC can target any of EXEC, EXEC_BATCH or EXEC_SP */
+				void *insert_exec = NULL;
+
 				is_proc = true;
 				command_type = TDS_CMD_EXECUTE;
+
+				switch (stmt->cmd_type)
+				{
+					case PLTSQL_STMT_EXEC:
+						insert_exec = ((PLtsql_stmt_exec *) stmt)->insert_exec;
+						break;
+					case PLTSQL_STMT_EXEC_BATCH:
+						insert_exec = ((PLtsql_stmt_exec_batch *) stmt)->insert_exec;
+						break;
+					case PLTSQL_STMT_EXEC_SP:
+						insert_exec = ((PLtsql_stmt_exec_sp *) stmt)->insert_exec;
+						break;
+					default:
+						break;
+				}
+
 				/*
 				 * For INSERT EXEC, report the row count set in
 				 * flush_insert_exec_temp_table(). Suppress it when an error is
 				 * pending: the rows are rolled back, no count is sent to the
 				 * client, and a counted DONE left pending here would otherwise
 				 * carry a stale count into the following error DONE token.
 				 */
-				if (!markErrorFlag &&
-					stmt->cmd_type == PLTSQL_STMT_EXEC &&
-					((PLtsql_stmt_exec *) stmt)->insert_exec != NULL)
-				{
-					command_type = TDS_CMD_INSERT;
-					row_count_valid = true;
-				}
-				else if (!markErrorFlag &&
-					stmt->cmd_type == PLTSQL_STMT_EXEC_BATCH &&
-					((PLtsql_stmt_exec_batch *) stmt)->insert_exec != NULL)
+				if (!markErrorFlag && insert_exec != NULL)
 				{
 					command_type = TDS_CMD_INSERT;
 					row_count_valid = true;

contrib/babelfishpg_tsql/src/pl_exec-2.c

@@ -867,9 +867,6 @@ exec_stmt_exec(PLtsql_execstate *estate, PLtsql_stmt_exec *stmt)
 	/* whether procedure was created WITH RECOMPILE */
 	bool created_with_recompile = false;		
 
-	/* INSERT EXEC handling - temp table lifecycle */
-	bool insert_exec_setup_done = false;
-
 	/*
 	 * We need to disable the explain gucs incase of sp_reset_connection
 	 * execution otherwise we will get explain output for it which is
@@ -901,8 +898,8 @@ exec_stmt_exec(PLtsql_execstate *estate, PLtsql_stmt_exec *stmt)
 		 * Setup INSERT EXEC (new path): create temp table to capture procedure
 		 * output. After procedure completes, temp table is flushed to target.
 		 */
-		if (pltsql_enable_new_insert_exec)
-			insert_exec_setup_done = insert_exec_setup(estate, stmt->insert_exec, true);
+		if (stmt->insert_exec != NULL)
+			insert_exec_setup(estate, stmt->insert_exec, true);
 
 		if (IS_TDS_CONN())
 		{
@@ -1338,12 +1335,12 @@ exec_stmt_exec(PLtsql_execstate *estate, PLtsql_stmt_exec *stmt)
 		exec_eval_cleanup(estate);
 		SPI_freetuptable(SPI_tuptable);
 
-		if (insert_exec_setup_done)
+		if (stmt->insert_exec != NULL)
 			insert_exec_flush_and_cleanup(estate, stmt->insert_exec);
 	}
 	PG_FINALLY();
 	{
-		if (insert_exec_setup_done)
+		if (stmt->insert_exec != NULL)
 			pltsql_insert_exec_reset_all();
 
 		/*
@@ -1543,9 +1540,6 @@ exec_stmt_exec_batch(PLtsql_execstate *estate, PLtsql_stmt_exec_batch *stmt)
 	char	   *old_db_name = get_cur_db_name();
 	char	   *cur_db_name = NULL;
 
-	/* INSERT EXEC handling - temp table lifecycle */
-	bool insert_exec_setup_done = false;
-
 	LOCAL_FCINFO(fcinfo, 1);
 
 	/*
@@ -1568,8 +1562,8 @@ exec_stmt_exec_batch(PLtsql_execstate *estate, PLtsql_stmt_exec_batch *stmt)
 		 * output. No implicit transaction for dynamic SQL (different semantics
 		 * than stored procs).
 		 */
-		if (pltsql_enable_new_insert_exec)
-			insert_exec_setup_done = insert_exec_setup(estate, stmt->insert_exec, false);
+		if (stmt->insert_exec != NULL)
+			insert_exec_setup(estate, stmt->insert_exec, false);
 
 		/* Get the C-String representation */
 		querystr = convert_value_to_string(estate, query, restype);
@@ -1592,12 +1586,12 @@ exec_stmt_exec_batch(PLtsql_execstate *estate, PLtsql_stmt_exec_batch *stmt)
 		if (fcinfo->isnull)
 			elog(ERROR, "pltsql_inline_handler failed");
 
-		if (insert_exec_setup_done)
+		if (stmt->insert_exec != NULL)
 			insert_exec_flush_and_cleanup(estate, stmt->insert_exec);
 	}
 	PG_FINALLY();
 	{
-		if (insert_exec_setup_done)
+		if (stmt->insert_exec != NULL)
 			pltsql_insert_exec_reset_all();
 
 		/* Restore past settings */
@@ -2224,9 +2218,6 @@ exec_stmt_exec_sp(PLtsql_execstate *estate, PLtsql_stmt_exec_sp *stmt)
 				int			save_nestlevel;
 				int			scope_level;
 				InlineCodeBlockArgs *args = NULL;
-				
-				/* INSERT EXEC handling - temp table lifecycle */
-				bool insert_exec_setup_done = false;
 
 				batch = exec_eval_expr(estate, stmt->query, &isnull1, &restype1, &restypmod1);
 				if (isnull1)
@@ -2279,8 +2270,8 @@ exec_stmt_exec_sp(PLtsql_execstate *estate, PLtsql_stmt_exec_sp *stmt)
 					 * The procedure output will be redirected to this temp table.
 					 * After procedure completes, we flush temp table to target and cleanup.
 					 */
-					if (pltsql_enable_new_insert_exec)
-						insert_exec_setup_done = insert_exec_setup(estate, stmt->insert_exec, true);
+					if (stmt->insert_exec != NULL)
+						insert_exec_setup(estate, stmt->insert_exec, true);
 
 					if (strcmp(batchstr, "") != 0)	/* check edge cases for
 													 * sp_executesql */
@@ -2293,12 +2284,12 @@ exec_stmt_exec_sp(PLtsql_execstate *estate, PLtsql_stmt_exec_sp *stmt)
 						exec_assign_value(estate, estate->datums[stmt->return_code_dno], Int32GetDatum(ret), false, INT4OID, 0);
 					}
 
-					if (insert_exec_setup_done)
+					if (stmt->insert_exec != NULL)
 						insert_exec_flush_and_cleanup(estate, stmt->insert_exec);
 				}
 				PG_FINALLY();
 				{
-					if (insert_exec_setup_done)
+					if (stmt->insert_exec != NULL)
 						pltsql_insert_exec_reset_all();
 
 					pltsql_revert_guc(save_nestlevel);

contrib/babelfishpg_tsql/src/pl_exec.c

@@ -4990,9 +4990,17 @@ exec_stmt_execsql(PLtsql_execstate *estate,
 		{
 			/* Open nesting level in engine */
 			BeginCompositeTriggers(CurrentMemoryContext);
-			/* TSQL commands must run inside an explicit transaction */
+			/*
+			 * TSQL commands must run inside an explicit transaction.
+			 *
+			 * Skip this for the INSERT EXEC flush statement
+			 * The flush runs while the INSERT EXEC context is still active, 
+			 * so the matching per-statement commit further below is suppressed.
+			 * The flush is a single INSERT that runs correctly under autocommit.
+			 */
 			if (!pltsql_disable_batch_auto_commit && support_tsql_trans &&
-				stmt->txn_data == NULL && !IsTransactionBlockActive())
+				stmt->txn_data == NULL && !IsTransactionBlockActive() &&
+				insert_exec_flush_estate == NULL)
 			{
 				MemoryContext oldCxt = CurrentMemoryContext;
 

contrib/babelfishpg_tsql/src/pl_handler.c

@@ -3067,6 +3067,16 @@ bbf_table_var_lookup(const char *relname, Oid relnamespace)
 	PLtsql_tbl *tbl;
 	PLtsql_execstate *estate = get_current_tsql_estate();
 
+	/*
+	 * During an INSERT EXEC flush the query runs through execute_batch/the
+	 * inline handler, which pushes its own (empty) estate. insert_exec_flush_estate
+	 * points us back at the estate that actually declared the target table
+	 * variable, so an "@tv" flush target resolves to its backing table.
+	 * Outside the flush it is NULL and we use the current (topmost) estate.
+	 */
+	estate = insert_exec_flush_estate ? insert_exec_flush_estate
+									  : get_current_tsql_estate();
+
 	if (prev_relname_lookup_hook)
 		relid = (*prev_relname_lookup_hook) (relname, relnamespace);
 	else

contrib/babelfishpg_tsql/src/pl_insert_exec.c

@@ -73,9 +73,16 @@ static void insertexec_destroy(DestReceiver *self);
  * Global INSERT EXEC context pointer - defined in pltsql.h, declared here.
  */
 InsertExecContext *insert_exec_ctx = NULL;
+PLtsql_execstate *insert_exec_flush_estate = NULL;
 
 extern void exec_set_rowcount(uint64 rowno);
 extern void exec_set_found(PLtsql_execstate *estate, bool state);
+/*
+ * The flush INSERT is routed through execute_batch (the top-level batch entry
+ * point). It runs through the same econtext setup as a normal T-SQL batch.
+ */
+extern int	execute_batch(PLtsql_execstate *estate, char *batch, InlineCodeBlockArgs *args, List *params);
+extern InlineCodeBlockArgs *create_args(int numargs);
 
 /*
  * Build a comma-separated list of quoted column identifiers from the parser's
@@ -523,6 +530,8 @@ flush_insert_exec_temp_table(PLtsql_execstate *estate, const char *target_schema
 	const char	   *temp_name;
 	Relation		temp_rel;
 	char		   *qualified_target;
+	InlineCodeBlockArgs	*flush_args;
+	PLtsql_execstate *flush_estate_saved;
 
 	if (insert_exec_ctx == NULL)
 		return;
@@ -581,23 +590,38 @@ flush_insert_exec_temp_table(PLtsql_execstate *estate, const char *target_schema
 
 	pfree(qualified_target);
 
-	/* Route through SPI_execute to run the flush INSERT in the current
-	 * transaction/SPI context. execute_batch cannot be used here: it enters
-	 * pltsql_inline_handler as an independent batch that manages its own
-	 * transaction boundaries, which aborts when the flush runs mid-INSERT-EXEC
-	 * (notably for table-variable targets whose lifetime is tied to the
-	 * surrounding transaction). */
-	rc = SPI_execute(flush_query.data, false, 0);
+	/*
+	 * Run the flush through execute_batch, publishing the caller's estate via
+	 * insert_exec_flush_estate for the duration so the inline handler's own
+	 * empty estate does not shadow it.
+	 */
+	flush_args = create_args(0);
+
+	flush_estate_saved = insert_exec_flush_estate;
+	insert_exec_flush_estate = estate;
+	PG_TRY();
+	{
+		rc = execute_batch(estate, flush_query.data, flush_args, NULL);
+	}
+	PG_CATCH();
+	{
+		insert_exec_flush_estate = flush_estate_saved;
+		PG_RE_THROW();
+	}
+	PG_END_TRY();
+	insert_exec_flush_estate = flush_estate_saved;
 
 	pfree(flush_query.data);
 
-	if (rc != SPI_OK_INSERT && rc != SPI_OK_INSERT_RETURNING)
+	if (rc != PLTSQL_RC_OK)
 		elog(ERROR, "INSERT EXEC failed due to error while flushing temp table to target table");
 
-	/* Update rowcount and FOUND for T-SQL compatibility */
-	estate->eval_processed = SPI_processed;
-	exec_set_rowcount(SPI_processed);
-	exec_set_found(estate, SPI_processed != 0);
+	/*
+	 * Report rows-affected from the DestReceiver's captured-row count
+	 */
+	estate->eval_processed = insert_exec_ctx->rows_processed;
+	exec_set_rowcount(insert_exec_ctx->rows_processed);
+	exec_set_found(estate, insert_exec_ctx->rows_processed != 0);
 }
 
 /*
@@ -758,6 +782,9 @@ insertexec_receive(TupleTableSlot *slot, DestReceiver *self)
 	/* Insert the projected tuple */
 	table_tuple_insert(temp_rel, insert_slot, myState->cid, 0, NULL);
 
+	/* INSERT EXEC rows-affected count */
+	insert_exec_ctx->rows_processed++;
+
 	/* Close immediately - do not hold open across subtransaction boundaries */
 	table_close(temp_rel, RowExclusiveLock);
 

contrib/babelfishpg_tsql/src/pltsql.h

@@ -2534,10 +2534,19 @@ typedef struct InsertExecContext
 	PLExecStateCallStack *call_stack_entry;	/* Call stack entry when INSERT EXEC started */
 	Oid			target_rel_oid;			/* OID of target table - lock held to detect schema changes */
 	bool		is_target_relation_modified;	/* Set by bbf_object_access_hook when target table is altered */
+	uint64		rows_processed;			/* Rows captured by the DestReceiver = INSERT EXEC rows-affected */
 } InsertExecContext;
 
 extern InsertExecContext *insert_exec_ctx;
 
+/*
+ * Set only during an INSERT EXEC flush. The flush runs through the inline
+ * handler, which pushes its own empty estate; this points back at the estate
+ * that declared the flush target so table-variable lookup, the implicit-
+ * transaction decision, and ownership chaining all resolve against the caller.
+ */
+extern PLtsql_execstate *insert_exec_flush_estate;
+
 extern Oid create_insert_exec_temp_table(const char *target_table, const char *column_list, const char *schema_name_in, const char *db_name_in);
 extern void pltsql_set_insert_exec_context_info(const char *target_table);
 extern void pltsql_insert_exec_reset_all(void);

contrib/babelfishpg_tsql/src/pltsql_utils.c

@@ -3065,6 +3065,15 @@ get_current_func_oid(void)
 	if (!pltsql_support_tsql_transactions())
 		return InvalidOid;
 
+	/*
+	 * During an INSERT EXEC flush the inline handler pushes its own (anonymous
+	 * batch) estate, so fall back to insert_exec_flush_estate (the procedure
+	 * that issued the INSERT EXEC) to keep ownership chaining intact.
+	 */
+	if (insert_exec_flush_estate != NULL)
+		return (insert_exec_flush_estate->func) ?
+			insert_exec_flush_estate->func->fn_oid : InvalidOid;
+
 	/*
 	* Fetch the top procedure excution state from execution state call stack
 	* and get the owner of that procedure. Top entry in stack will have

contrib/babelfishpg_tsql/src/tsqlIface.cpp

@@ -2163,11 +2163,11 @@ class tsqlBuilder : public tsqlCommonMutator
 	 */
 	void handleInsertExec(TSqlParser::Dml_statementContext *ctx)
 	{
-		/* INSERT EXEC is not allowed in functions unless target is a table variable */
+		/* INSERT EXEC is not allowed in a function */
 		if (is_compiling_create_function())
 		{
 			auto ddl_object = ctx->insert_statement()->ddl_object();
-			if (ddl_object && !ddl_object->local_id())
+			if (ddl_object)
 				throw PGErrorWrapperException(ERROR, ERRCODE_INVALID_FUNCTION_DEFINITION,
 					"'INSERT EXEC' cannot be used within a function", getLineAndPos(ddl_object));
 		}