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,87 @@ 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:
+		{
+			DropStmt *drop = (DropStmt *) parsetree;
+			if (drop->removeType == OBJECT_TABLE && pltsql_insert_exec_active())
+			{
+				Oid target_oid = pltsql_get_insert_exec_target_rel_oid();
+
+				if (OidIsValid(target_oid))
+				{
+					/*
+					 * Regular table: use OID comparison so schema-qualified names
+					 * like schema1.orders and schema2.orders are distinguished.
+					 */
+					ListCell *lc;
+					foreach(lc, drop->objects)
+					{
+						List     *names    = lfirst_node(List, lc);
+						RangeVar *rv       = makeRangeVarFromNameList(names);
+						Oid       drop_oid = RangeVarGetRelid(rv, NoLock, true /* missing_ok */);
+
+						if (OidIsValid(drop_oid) && drop_oid == target_oid)
+							ereport(ERROR,
+									(errcode(ERRCODE_OBJECT_IN_USE),
+									 errmsg("cannot %s \"%s\" because it is "
+											"being used by active queries in "
+											"this session",
+											"DROP TABLE", get_rel_name(drop_oid))));
+					}
+				}
+				else
+				{
+					/*
+					 * Temp table / table variable: session-local, no persistent OID.
+					 * Fall back to unqualified name comparison (safe over-block).
+					 */
+					const char *target      = pltsql_get_insert_exec_target_table();
+					const char *dot         = target ? strrchr(target, '.') : NULL;
+					const char *unqualified = dot ? (dot + 1) : target;
+					ListCell   *lc;
+
+					if (unqualified != NULL)
+					{
+						foreach(lc, drop->objects)
+						{
+							List *names   = lfirst_node(List, lc);
+							char *relname = strVal(llast(names));
+
+							if (pg_strcasecmp(relname, unqualified) == 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", relname)));
+						}
+					}
+				}
+			}
+			break;
+		}
+
 		default:
 			return false;
 			break;