Optimized fix with proper testing

Author: arvis108

PR_SUMMARY.md

@@ -965,10 +965,6 @@ def execute(
 # Executing a new statement. Reset is_stmt_prepared to false
             self.is_stmt_prepared = [False]
 
-        # Enhanced multi-statement handling - pyodbc approach
-        # Apply SET NOCOUNT ON to all multi-statement queries to prevent result set issues
-        if self._is_multistatement_query(operation):
-            operation = self._add_nocount_to_multistatement_sql(operation)
 
         log('debug', "Executing query: %s", operation)
         for i, param in enumerate(parameters):
@@ -1010,9 +1006,9 @@ def execute(
 
         self.last_executed_stmt = operation
 
-        # Update rowcount after execution
+        # Update rowcount after execution (before buffering)
         # TODO: rowcount return code from SQL needs to be handled
-        self.rowcount = ddbc_bindings.DDBCSQLRowCount(self.hstmt)
+        initial_rowcount = ddbc_bindings.DDBCSQLRowCount(self.hstmt)
 
         # Initialize description after execution
         # After successful execution, initialize description if there are results
@@ -1024,12 +1020,16 @@ def execute(
             # If describe fails, it's likely there are no results (e.g., for INSERT)
             self.description = None
 
-        # Reset rownumber for new result set (only for SELECT statements)
+        # Buffer intermediate results automatically (pyODBC-style approach)
+        self._buffer_intermediate_results()
+        
+        # Set final rowcount based on result type (preserve original rowcount for non-SELECT)
         if self.description:  # If we have column descriptions, it's likely a SELECT
             self.rowcount = -1
             self._reset_rownumber()
         else:
-            self.rowcount = ddbc_bindings.DDBCSQLRowCount(self.hstmt)
+            # For non-SELECT statements (INSERT/UPDATE/DELETE), preserve the original rowcount
+            self.rowcount = initial_rowcount
             self._clear_rownumber()
 
         # After successful execution, initialize description if there are results
@@ -2188,11 +2188,11 @@ def tables(self, table=None, catalog=None, schema=None, tableType=None):
                 ("table_type", str, None, 128, 128, 0, False),
                 ("remarks", str, None, 254, 254, 0, True)
             ]
-            
+        
             # Use the helper method to prepare the result set
             return self._prepare_metadata_result_set(fallback_description=fallback_description)
 
         except Exception as e:
             # Log the error and re-raise
             log('error', f"Error executing tables query: {e}")
-            raise
+            raise

mssql_python/cursor.py

@@ -10754,3 +10754,30 @@ def test_close(db_connection):
         pytest.fail(f"Cursor close test failed: {e}")
     finally:
         cursor = db_connection.cursor()
+
+def test_multi_statement_query(cursor, db_connection):
+    """Test multi-statement query with temp tables"""
+    try:
+        # Single SQL with multiple statements - tests pyODBC-style buffering
+        multi_statement_sql = """
+        CREATE TABLE #TestData (id INT, name NVARCHAR(50), value INT);
+        INSERT INTO #TestData VALUES (1, 'Test1', 100), (2, 'Test2', 200);
+        SELECT COALESCE(name, 'DEFAULT') as result_name, SUM(value) as total_value INTO #TempResult FROM #TestData GROUP BY name;
+        SELECT result_name, total_value, 'SUCCESS' as status FROM #TempResult ORDER BY result_name;
+        """
+        
+        cursor.execute(multi_statement_sql)
+        results = cursor.fetchall()
+        
+        assert len(results) > 0, "Multi-statement query should return results"
+        assert results[0][2] == 'SUCCESS', "Should return success status"
+        
+    except Exception as e:
+        pytest.fail(f"Multi-statement query test failed: {e}")
+    finally:
+        try:
+            cursor.execute("DROP TABLE #TestData")
+            cursor.execute("DROP TABLE #TempResult")
+            db_connection.commit()
+        except:
+            pass