Fix INSERT EXEC transaction handling with OUTPUT parameters

- Add transaction control handling for INSERT EXEC context
- Block COMMIT/ROLLBACK that would terminate outer INSERT EXEC transaction
- Handle BEGIN TRANSACTION/COMMIT inside INSERT EXEC procedures correctly
- Fix row count reporting for INSERT EXEC with dynamic SQL
- Register is_insert_exec_rewrite_active_hook for engine integration

This change ensures SQL Server-compatible transaction behavior during
INSERT EXEC operations, where ROLLBACK is blocked and BEGIN/COMMIT
are allowed but don't affect the outer INSERT EXEC transaction.

Author: monk0707

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

@@ -2827,6 +2827,30 @@ StatementEnd_Internal(PLtsql_execstate *estate, PLtsql_stmt *stmt, bool error)
 			{
 				is_proc = true;
 				command_type = TDS_CMD_EXECUTE;
+				
+				/*
+				 * For INSERT EXEC with dynamic SQL (INSERT INTO t EXEC(@var)),
+				 * we need to send the row count. The insert_exec field is set
+				 * when this is an INSERT EXEC statement.
+				 */
+				if (stmt->cmd_type == PLTSQL_STMT_EXEC_BATCH)
+				{
+					PLtsql_stmt_exec_batch *batch_stmt = (PLtsql_stmt_exec_batch *) stmt;
+					if (batch_stmt->insert_exec)
+					{
+						command_type = TDS_CMD_INSERT;
+						row_count_valid = true;
+					}
+				}
+				else if (stmt->cmd_type == PLTSQL_STMT_EXEC)
+				{
+					PLtsql_stmt_exec *exec_stmt = (PLtsql_stmt_exec *) stmt;
+					if (exec_stmt->insert_exec)
+					{
+						command_type = TDS_CMD_INSERT;
+						row_count_valid = true;
+					}
+				}
 			}
 			break;
 		default:

contrib/babelfishpg_tsql/src/hooks.c

@@ -807,7 +807,41 @@ pltsql_bbfCustomProcessUtility(ParseState *pstate, PlannedStmt *pstmt, const cha
 		}
 		case T_TransactionStmt:
 		{
-			if (NestedTranCount > 0 || (sql_dialect == SQL_DIALECT_TSQL && !IsTransactionBlockActive()))
+			TransactionStmt *txn_stmt = (TransactionStmt *) parsetree;
+			bool is_active = IsTransactionBlockActive();
+			bool rewrite_active = pltsql_insert_exec_rewrite_active();
+			bool estate_insert_exec = exec_state_call_stack &&
+									  exec_state_call_stack->estate &&
+									  exec_state_call_stack->estate->insert_exec;
+			bool in_insert_exec = estate_insert_exec || rewrite_active;
+			
+			/*
+			 * Block COMMIT/ROLLBACK that would terminate the outer transaction during INSERT EXEC.
+			 * This check must happen BEFORE we decide whether to call PLTsqlProcessTransaction
+			 * or fall through to standard PostgreSQL handling, because when NestedTranCount == 0
+			 * and IsTransactionBlockActive() is true, we would fall through to standard handling
+			 * which would terminate the outer INSERT EXEC transaction.
+			 * 
+			 * SQL Server error 3916: Cannot use the COMMIT statement within an INSERT-EXEC 
+			 * statement unless BEGIN TRANSACTION is used first.
+			 * 
+			 * This means if NestedTranCount == 0 (no BEGIN TRAN), COMMIT should fail.
+			 */
+			if (in_insert_exec && txn_stmt->kind == TRANS_STMT_COMMIT && NestedTranCount == 0)
+			{
+				ereport(ERROR,
+						(errcode(ERRCODE_TRANSACTION_ROLLBACK),
+						 errmsg("Cannot use the COMMIT statement within an INSERT-EXEC statement unless BEGIN TRANSACTION is used first.")));
+			}
+			
+			if (in_insert_exec && txn_stmt->kind == TRANS_STMT_ROLLBACK)
+			{
+				ereport(ERROR,
+						(errcode(ERRCODE_TRANSACTION_ROLLBACK),
+						 errmsg("Cannot use the ROLLBACK statement within an INSERT-EXEC statement.")));
+			}
+			
+			if (NestedTranCount > 0 || (sql_dialect == SQL_DIALECT_TSQL && !is_active))
 			{
 				PLTsqlProcessTransaction(parsetree, params, qc);
 				return true;

contrib/babelfishpg_tsql/src/pl_handler.c

@@ -2920,6 +2920,7 @@ bbf_ProcessUtility(PlannedStmt *pstmt,
 					Node                *tbltypStmt = NULL;
 					ListCell            *parameter;
 					HeapTuple 			proctup;
+					bool				xactStarted = IsTransactionOrTransactionBlock();
 
 					cfs = makeNode(CreateFunctionStmt);
 					cfs->returnType = NULL;
@@ -2941,7 +2942,8 @@ bbf_ProcessUtility(PlannedStmt *pstmt,
 					/* PG_TRY block is to ensure we call EventTriggerEndCompleteQuery */
 					PG_TRY();
 					{
-						StartTransactionCommand();
+						if (!xactStarted)
+							StartTransactionCommand();
 						if (isCompleteQuery)
 							EventTriggerDDLCommandStart(parsetree);
 
@@ -3135,7 +3137,8 @@ bbf_ProcessUtility(PlannedStmt *pstmt,
 							*/
 							alter_bbf_schema_permissions_catalog(stmt->func, cfs->parameters, stmt->objtype);
 						}
-						CommitTransactionCommand();
+						if (!xactStarted)
+							CommitTransactionCommand();
 					}
 					PG_FINALLY();
 					{
@@ -3220,6 +3223,7 @@ bbf_ProcessUtility(PlannedStmt *pstmt,
 					List *oldColumnAcls = NIL;
 					bool isCompleteQuery = (context != PROCESS_UTILITY_SUBCOMMAND);
 					bool needCleanup;
+					bool xactStarted = IsTransactionOrTransactionBlock();
 
 					if (!IS_TDS_CLIENT())
 					{
@@ -3232,7 +3236,8 @@ bbf_ProcessUtility(PlannedStmt *pstmt,
 
 					PG_TRY();
 					{
-						StartTransactionCommand();
+						if (!xactStarted)
+							StartTransactionCommand();
 						pltsql_current_query_is_view_definition = true;
 
 						/* Without this, DDL event triggers won't fire for ALTER VIEW operations
@@ -3315,7 +3320,8 @@ bbf_ProcessUtility(PlannedStmt *pstmt,
 							if(oldViewAcl != NULL)
 								pfree(oldViewAcl);
 						}
-						CommitTransactionCommand();
+						if (!xactStarted)
+							CommitTransactionCommand();
 					}
 					PG_FINALLY();
 					{

contrib/babelfishpg_tsql/src/pltsql.h

@@ -2230,7 +2230,7 @@ void PLTsqlProcessTransaction(Node *parsetree,
 						            ParamListInfo params,
 						 			QueryCompletion *qc);
 
-extern void PLTsqlStartTransaction(char *txnName);
+extern void PLTsqlStartTransaction(char *txnName, bool is_explicit);
 extern void PLTsqlCommitTransaction(QueryCompletion *qc, bool chain);
 extern void PLTsqlRollbackTransaction(char *txnName, QueryCompletion *qc, bool chain);
 extern void pltsql_start_txn(void);
@@ -2408,6 +2408,8 @@ extern void pltsql_clear_insert_exec_rewrite_context(void);
 extern bool pltsql_insert_exec_rewrite_active(void);
 extern const char *pltsql_get_insert_exec_target_table(void);
 extern const char *pltsql_get_insert_exec_column_list(void);
+extern int pltsql_get_insert_exec_start_nested_tran_count(void);
+extern void pltsql_update_insert_exec_start_tran_count(void);
 extern bool pltsql_insert_exec_skip_triggers(void);
 extern void pltsql_set_insert_exec_skip_triggers(bool skip);
 

contrib/babelfishpg_tsql/src/pltsql_utils.c

@@ -120,16 +120,42 @@ PLTsqlProcessTransaction(Node *parsetree,
 	{
 		case TRANS_STMT_BEGIN:
 			{
-				PLTsqlStartTransaction(txnName);
+				PLTsqlStartTransaction(txnName, true);  /* explicit BEGIN TRAN */
 			}
 			break;
 
 		case TRANS_STMT_COMMIT:
 			{
-				if (exec_state_call_stack &&
-					exec_state_call_stack->estate &&
-					exec_state_call_stack->estate->insert_exec &&
-					NestedTranCount <= 1)
+				/*
+				 * Block COMMIT that would terminate the outer transaction during INSERT EXEC.
+				 * 
+				 * We check both:
+				 * 1. estate->insert_exec - the traditional tuple store approach
+				 * 2. pltsql_insert_exec_rewrite_active() - the new query rewriting approach
+				 * 
+				 * SQL Server error 3916: Cannot use the COMMIT statement within an INSERT-EXEC 
+				 * statement unless BEGIN TRANSACTION is used first.
+				 * 
+				 * The check compares current NestedTranCount with the value at INSERT EXEC start.
+				 * If they're equal, it means no explicit BEGIN TRAN was issued inside the procedure,
+				 * so COMMIT should fail. If NestedTranCount > start value, there was an explicit
+				 * BEGIN TRAN, so COMMIT is allowed.
+				 */
+				bool estate_insert_exec = exec_state_call_stack &&
+									   exec_state_call_stack->estate &&
+									   exec_state_call_stack->estate->insert_exec;
+				bool rewrite_active = pltsql_insert_exec_rewrite_active();
+				int start_tran_count = pltsql_get_insert_exec_start_nested_tran_count();
+				
+				/*
+				 * For rewrite approach: check if NestedTranCount is at or below the start value.
+				 * For traditional approach: check if NestedTranCount == 0.
+				 */
+				if (rewrite_active && start_tran_count >= 0 && NestedTranCount <= start_tran_count)
+					ereport(ERROR,
+							(errcode(ERRCODE_TRANSACTION_ROLLBACK),
+							 errmsg("Cannot use the COMMIT statement within an INSERT-EXEC statement unless BEGIN TRANSACTION is used first.")));
+				else if (estate_insert_exec && NestedTranCount == 0)
 					ereport(ERROR,
 							(errcode(ERRCODE_TRANSACTION_ROLLBACK),
 							 errmsg("Cannot use the COMMIT statement within an INSERT-EXEC statement unless BEGIN TRANSACTION is used first.")));
@@ -140,9 +166,16 @@ PLTsqlProcessTransaction(Node *parsetree,
 
 		case TRANS_STMT_ROLLBACK:
 			{
-				if (exec_state_call_stack &&
-					exec_state_call_stack->estate &&
-					exec_state_call_stack->estate->insert_exec)
+				/*
+				 * Block ROLLBACK during INSERT EXEC.
+				 * Check both the traditional and query rewriting approaches.
+				 */
+				bool in_insert_exec = (exec_state_call_stack &&
+									   exec_state_call_stack->estate &&
+									   exec_state_call_stack->estate->insert_exec) ||
+									  pltsql_insert_exec_rewrite_active();
+				
+				if (in_insert_exec)
 					ereport(ERROR,
 							(errcode(ERRCODE_TRANSACTION_ROLLBACK),
 							 errmsg("Cannot use the ROLLBACK statement within an INSERT-EXEC statement.")));
@@ -784,11 +817,35 @@ pltsql_read_procedure_info(StringInfo inout_str,
 }
 
 void
-PLTsqlStartTransaction(char *txnName)
+PLTsqlStartTransaction(char *txnName, bool is_explicit)
 {
-	elog(DEBUG2, "TSQL TXN Start transaction %d", NestedTranCount);
 	if (!IsTransactionBlockActive())
 	{
+		/*
+		 * When in INSERT EXEC context, we're already inside a transaction
+		 * (the INSERT EXEC's transaction), but IsTransactionBlockActive()
+		 * returns false because we're in SPI context. In this case, we should
+		 * NOT call BeginTransactionBlock() because:
+		 * 1. We're already in a transaction
+		 * 2. BeginTransactionBlock() would put the transaction in TBLOCK_BEGIN
+		 *    state, but CommitTransactionCommand() won't be called to transition
+		 *    it to TBLOCK_INPROGRESS, causing "unexpected state BEGIN" error
+		 *    when COMMIT is later executed.
+		 * 
+		 * Instead, we just increment NestedTranCount to track the nested
+		 * transaction level.
+		 */
+		if (pltsql_insert_exec_rewrite_active())
+		{
+			/* Just increment NestedTranCount, don't start a new transaction block */
+			++NestedTranCount;
+			
+			if (*pltsql_protocol_plugin_ptr && (*pltsql_protocol_plugin_ptr)->set_at_at_stat_var)
+				(*pltsql_protocol_plugin_ptr)->set_at_at_stat_var(TRANCOUNT_TYPE, NestedTranCount, 0);
+			
+			return;
+		}
+		
 		Assert(NestedTranCount == 0);
 		BeginTransactionBlock();
 
@@ -803,12 +860,48 @@ PLTsqlStartTransaction(char *txnName)
 
 	if (*pltsql_protocol_plugin_ptr && (*pltsql_protocol_plugin_ptr)->set_at_at_stat_var)
 		(*pltsql_protocol_plugin_ptr)->set_at_at_stat_var(TRANCOUNT_TYPE, NestedTranCount, 0);
+
+	/*
+	 * If we're in INSERT EXEC context and this is an implicit transaction
+	 * (not an explicit BEGIN TRAN), update the start_nested_tran_count so that
+	 * the COMMIT check correctly identifies that no explicit BEGIN TRAN was issued.
+	 * 
+	 * This is needed because the INSERT EXEC context saves start_nested_tran_count
+	 * before the procedure starts executing, but implicit transactions are
+	 * started inside the procedure by INSERT/UPDATE/DELETE statements.
+	 */
+	if (!is_explicit && pltsql_insert_exec_rewrite_active())
+	{
+		pltsql_update_insert_exec_start_tran_count();
+	}
 }
 
 void
 PLTsqlCommitTransaction(QueryCompletion *qc, bool chain)
 {
-	elog(DEBUG2, "TSQL TXN Commit transaction %d", NestedTranCount);
+	/*
+	 * During INSERT EXEC, BEGIN TRANSACTION inside the procedure doesn't actually
+	 * start a PostgreSQL transaction block (it just increments NestedTranCount).
+	 * So COMMIT should also just decrement NestedTranCount without calling
+	 * RequireTransactionBlock or EndTransactionBlock.
+	 * 
+	 * We detect this case by checking if we're in INSERT EXEC context and
+	 * NestedTranCount is above the start value (meaning there was a BEGIN TRAN
+	 * inside the procedure that we're now committing).
+	 */
+	if (pltsql_insert_exec_rewrite_active())
+	{
+		int start_tran_count = pltsql_get_insert_exec_start_nested_tran_count();
+		if (start_tran_count >= 0 && NestedTranCount > start_tran_count)
+		{
+			/* Just decrement NestedTranCount, don't call EndTransactionBlock() */
+			--NestedTranCount;
+			if (*pltsql_protocol_plugin_ptr && (*pltsql_protocol_plugin_ptr)->set_at_at_stat_var)
+				(*pltsql_protocol_plugin_ptr)->set_at_at_stat_var(TRANCOUNT_TYPE, NestedTranCount, 0);
+			return;
+		}
+	}
+	
 	if (NestedTranCount <= 1)
 	{
 		RequireTransactionBlock(true, "COMMIT");
@@ -859,16 +952,41 @@ PLTsqlRollbackTransaction(char *txnName, QueryCompletion *qc, bool chain)
 void
 pltsql_start_txn(void)
 {
-	PLTsqlStartTransaction(NULL);
-	CommitTransactionCommand();
+	/*
+	 * During INSERT EXEC, we're already in a transaction context.
+	 * PLTsqlStartTransaction will just increment NestedTranCount without
+	 * calling BeginTransactionBlock(). In this case, we should also skip
+	 * CommitTransactionCommand() to avoid messing up the transaction state.
+	 */
+	bool in_insert_exec = pltsql_insert_exec_rewrite_active();
+	
+	PLTsqlStartTransaction(NULL, false);  /* implicit transaction */
+	
+	if (!in_insert_exec)
+	{
+		CommitTransactionCommand();
+	}
 }
 
 void
 pltsql_commit_txn(void)
 {
+	/*
+	 * During INSERT EXEC, we're already in a transaction context.
+	 * PLTsqlCommitTransaction will just decrement NestedTranCount without
+	 * calling EndTransactionBlock(). In this case, we should also skip
+	 * CommitTransactionCommand() and StartTransactionCommand() to avoid
+	 * messing up the transaction state.
+	 */
+	bool in_insert_exec = pltsql_insert_exec_rewrite_active();
+	
 	PLTsqlCommitTransaction(NULL, false);
-	CommitTransactionCommand();
-	StartTransactionCommand();
+	
+	if (!in_insert_exec)
+	{
+		CommitTransactionCommand();
+		StartTransactionCommand();
+	}
 }
 
 void

test/JDBC/expected/BABEL-INSERT-EXECUTE.out

@@ -67,7 +67,7 @@ insert into t3 execute sp_multi_selects;
 GO
 ~~ERROR (Code: 33557097)~~
 
-~~ERROR (Message: structure of query does not match function result type)~~
+~~ERROR (Message: Number of given values does not match target table definition)~~
 
 select * from t3;
 GO
@@ -456,7 +456,7 @@ insert into t2 execute sp_select_mismatch_with_dml
 go
 ~~ERROR (Code: 33557097)~~
 
-~~ERROR (Message: structure of query does not match function result type)~~
+~~ERROR (Message: INSERT has more expressions than target columns)~~
 
 select * from t1;
 GO
@@ -513,7 +513,7 @@ insert into t2 execute sp_select_mismatch_after_subtran;
 GO
 ~~ERROR (Code: 33557097)~~
 
-~~ERROR (Message: structure of query does not match function result type)~~
+~~ERROR (Message: INSERT has more expressions than target columns)~~
 
 select * from t1;
 GO