INSERT-EXECUTE Redesign

contrib/babelfishpg_tds/error_mapping.txt

@@ -187,4 +187,6 @@ XX000 ERRCODE_INTERNAL_ERROR	"The table-valued parameter \"%s\" must be declared
 22008 ERRCODE_DATETIME_VALUE_OUT_OF_RANGE	"Adding a value to a \'%s\' column caused an overflow." SQL_ERROR_517 16
 42P01 ERRCODE_UNDEFINED_TABLE	"FOR JSON AUTO requires at least one table for generating JSON objects. Use FOR JSON PATH or add a FROM clause with a table name."		SQL_ERROR_13600	16
 42P01 ERRCODE_FEATURE_NOT_SUPPORTED	"sub-select and values for json auto are not currently supported."		SQL_ERROR_13600	16
+54000 ERRCODE_PROGRAM_LIMIT_EXCEEDED	"nested INSERT ... EXECUTE statements are not allowed"	SQL_ERROR_8164	16
+0A000 ERRCODE_FEATURE_NOT_SUPPORTED	"INSERT EXEC failed because the stored procedure altered the schema of the target table."	SQL_ERROR_556	16
 

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

@@ -190,9 +190,11 @@ SendPendingDone(bool more)
 
 			/*
 			 * If a statement throws an error, the row count should be always
-			 * 0.
+			 * 0. Reset it here in case a previous statement (e.g., INSERT
+			 * inside a procedure during INSERT EXEC) had set a non-zero count
+			 * before the error was thrown.
 			 */
-			Assert(TdsPendingDoneRowCnt == 0);
+			TdsPendingDoneRowCnt = 0;
 		}
 
 		TDS_DEBUG(TDS_DEBUG3, "SendPendingDone: putbytes");
@@ -2725,20 +2727,29 @@ StatementEnd_Internal(PLtsql_execstate *estate, PLtsql_stmt *stmt, bool error)
 								 * is inside the procedure of an INSERT-EXEC,
 								 * or if the INSERT itself is an INSERT-EXEC
 								 * and it just returned error.
+								 * 
+								 * Check both estate->insert_exec (old approach)
+								 * and pltsql_insert_exec_active() (new DestReceiver approach).
 								 */
 								row_count_valid = !estate->insert_exec &&
+									!(pltsql_plugin_handler_ptr->pltsql_insert_exec_active &&
+									  pltsql_plugin_handler_ptr->pltsql_insert_exec_active()) &&
 									!(markErrorFlag &&
 									  ((PLtsql_stmt_execsql *) stmt)->insert_exec);
 							}
 							else if (plansource->commandTag == CMDTAG_UPDATE)
 							{
 								command_type = TDS_CMD_UPDATE;
-								row_count_valid = !estate->insert_exec;
+								row_count_valid = !estate->insert_exec &&
+									!(pltsql_plugin_handler_ptr->pltsql_insert_exec_active &&
+									  pltsql_plugin_handler_ptr->pltsql_insert_exec_active());
 							}
 							else if (plansource->commandTag == CMDTAG_DELETE)
 							{
 								command_type = TDS_CMD_DELETE;
-								row_count_valid = !estate->insert_exec;
+								row_count_valid = !estate->insert_exec &&
+									!(pltsql_plugin_handler_ptr->pltsql_insert_exec_active &&
+									  pltsql_plugin_handler_ptr->pltsql_insert_exec_active());
 							}
 
 							/*
@@ -2748,7 +2759,9 @@ StatementEnd_Internal(PLtsql_execstate *estate, PLtsql_stmt *stmt, bool error)
 							else if (plansource->commandTag == CMDTAG_SELECT)
 							{
 								command_type = TDS_CMD_SELECT;
-								row_count_valid = !estate->insert_exec;
+								row_count_valid = !estate->insert_exec &&
+									!(pltsql_plugin_handler_ptr->pltsql_insert_exec_active &&
+									  pltsql_plugin_handler_ptr->pltsql_insert_exec_active());
 							}
 						}
 					}
@@ -2771,6 +2784,23 @@ StatementEnd_Internal(PLtsql_execstate *estate, PLtsql_stmt *stmt, bool error)
 			{
 				is_proc = true;
 				command_type = TDS_CMD_EXECUTE;
+
+				/*
+				 * For INSERT EXEC, we need to report the row count.
+				 * The row count is set in flush_insert_exec_temp_table().
+				 */
+				if (stmt->cmd_type == PLTSQL_STMT_EXEC &&
+					((PLtsql_stmt_exec *) stmt)->insert_exec)
+				{
+					command_type = TDS_CMD_INSERT;
+					row_count_valid = true;
+				}
+				else if (stmt->cmd_type == PLTSQL_STMT_EXEC_BATCH &&
+					((PLtsql_stmt_exec_batch *) stmt)->insert_exec)
+				{
+					command_type = TDS_CMD_INSERT;
+					row_count_valid = true;
+				}
 			}
 			break;
 		default:

contrib/babelfishpg_tsql/src/hooks.c

@@ -802,13 +802,81 @@ pltsql_bbfCustomProcessUtility(ParseState *pstate, PlannedStmt *pstmt, const cha
 		}
 		case T_TransactionStmt:
 		{
-			if (NestedTranCount > 0 || (sql_dialect == SQL_DIALECT_TSQL && !IsTransactionBlockActive()))
+			TransactionStmt *stmt = (TransactionStmt *) parsetree;
+			/*
+			 * Call PLTsqlProcessTransaction when:
+			 * 1. NestedTranCount > 0 (explicit transaction)
+			 * 2. TSQL dialect and no active transaction block
+			 * 3. INSERT EXEC is active and it's a ROLLBACK (to block with error 3915)
+			 * 4. INSERT EXEC is active and it's a COMMIT (to block with error 3916)
+			 * 5. INSERT EXEC is active and it's a BEGIN TRAN (to properly track NestedTranCount)
+			 */
+			if (NestedTranCount > 0 || 
+				(sql_dialect == SQL_DIALECT_TSQL && !IsTransactionBlockActive()) ||
+				(pltsql_insert_exec_active() && 
+				 (stmt->kind == TRANS_STMT_ROLLBACK || stmt->kind == TRANS_STMT_ROLLBACK_TO ||
+				  stmt->kind == TRANS_STMT_COMMIT || stmt->kind == TRANS_STMT_BEGIN)))
 			{
 				PLTsqlProcessTransaction(parsetree, params, qc);
 				return true;
 			}
 			break;
 		}
+		case T_DropStmt:
+		{
+			/*
+			 * Check if DROP TABLE is trying to drop the INSERT EXEC target table.
+			 * SQL Server error 556: "cannot DROP TABLE because it is being used 
+			 * by active queries in this session"
+			 */
+			DropStmt *stmt = (DropStmt *) parsetree;
+
+			if (stmt->removeType == OBJECT_TABLE && pltsql_insert_exec_active())
+			{
+				const char *target_table = pltsql_get_insert_exec_target_table();
+
+				if (target_table != NULL)
+				{
+					ListCell *cell;
+
+					foreach(cell, stmt->objects)
+					{
+						List *names = (List *) lfirst(cell);
+						char *table_name = NULL;
+
+						/* Get the table name from the list (last element) */
+						if (list_length(names) > 0)
+						{
+							table_name = strVal(llast(names));
+						}
+
+						if (table_name != NULL)
+						{
+							/*
+							 * Compare table names (case-insensitive for temp tables).
+							 * The target_table might include schema prefix, so we need
+							 * to compare just the table name part.
+							 */
+							const char *target_name = strrchr(target_table, '.');
+							if (target_name != NULL)
+								target_name++;  /* Skip the dot */
+							else
+								target_name = target_table;
+
+							if (pg_strcasecmp(table_name, target_name) == 0)
+							{
+								ereport(ERROR,
+										(errcode(ERRCODE_OBJECT_IN_USE),
+										 errmsg("cannot %s \"%s\" because it is being used by active queries in this session",
+												"DROP TABLE", table_name)));
+							}
+						}
+					}
+				}
+			}
+			/* Let the default handler process the DROP */
+			break;
+		}
 		default:
 			return false;
 			break;