diff --git a/mssql_python/cursor.py b/mssql_python/cursor.py index 12d630d2..7ae9938e 100644 --- a/mssql_python/cursor.py +++ b/mssql_python/cursor.py @@ -884,6 +884,53 @@ def next(self): """ return self.__next__() + def _buffer_intermediate_results(self): + """ + Buffer intermediate results automatically. + + This method skips "rows affected" messages and empty result sets, + positioning the cursor on the first meaningful result set that contains + actual data. This eliminates the need for SET NOCOUNT ON detection. + """ + try: + # Keep advancing through result sets until we find one with actual data + # or reach the end + while True: + # Check if current result set has actual columns/data + if self.description and len(self.description) > 0: + # We have a meaningful result set with columns, stop here + break + + # Try to advance to next result set + try: + ret = ddbc_bindings.DDBCSQLMoreResults(self.hstmt) + + # If no more result sets, we're done + if ret == ddbc_sql_const.SQL_NO_DATA.value: + break + + # Check for errors + check_error(ddbc_sql_const.SQL_HANDLE_STMT.value, self.hstmt, ret) + + # Update description for the new result set + column_metadata = [] + try: + ddbc_bindings.DDBCSQLDescribeCol(self.hstmt, column_metadata) + self._initialize_description(column_metadata) + except Exception: + # If describe fails, it's likely there are no results (e.g., for INSERT) + self.description = None + + except Exception: + # If we can't advance further, stop + break + + except Exception as e: + log('warning', "Exception occurred during `_buffer_intermediate_results` %s", e) + # If anything goes wrong during buffering, continue with current state + # This ensures we don't break existing functionality + pass + def execute( self, operation: str, @@ -965,6 +1012,7 @@ def execute( # Executing a new statement. Reset is_stmt_prepared to false self.is_stmt_prepared = [False] + log('debug', "Executing query: %s", operation) for i, param in enumerate(parameters): log('debug', @@ -1005,9 +1053,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 @@ -1019,12 +1067,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 + 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 @@ -2183,11 +2235,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 \ No newline at end of file + raise diff --git a/tests/test_004_cursor.py b/tests/test_004_cursor.py index 38832131..9fc5a473 100644 --- a/tests/test_004_cursor.py +++ b/tests/test_004_cursor.py @@ -10744,6 +10744,149 @@ def test_datetime_string_parameter_binding(cursor, db_connection): drop_table_if_exists(cursor, table_name) db_connection.commit() +def test_multi_statement_query(cursor, db_connection): + """Test multi-statement query with temp tables""" + table_name = "#temp1" + try: + drop_table_if_exists(cursor, table_name) + # Single SQL with multiple statements - tests pyODBC-style buffering + multi_statement_sql = f""" + SELECT 1 as col1, 'test' as col2 INTO {table_name}; + SELECT * FROM {table_name}; + """ + + cursor.execute(multi_statement_sql) + results = cursor.fetchall() + + assert len(results) > 0, "Multi-statement query should return results" + assert results[0][1] == 'test', "Should return string 'test'" + + except Exception as e: + pytest.fail(f"Multi-statement query test failed: {e}") + finally: + drop_table_if_exists(cursor, table_name) + db_connection.commit() + +def test_multiple_result_sets_with_nextset(cursor, db_connection): + """Test multiple result sets with multiple select statements on temp tables with nextset()""" + table_name1 = "#TempData1" + table_name2 = "#TempData2" + try: + drop_table_if_exists(cursor, table_name1) + drop_table_if_exists(cursor, table_name2) + + # Create temp tables and execute multiple SELECT statements + multi_select_sql = f""" + CREATE TABLE {table_name1} (id INT, name NVARCHAR(50)); + INSERT INTO {table_name1} VALUES (1, 'First'), (2, 'Second'); + + CREATE TABLE {table_name2} (id INT, value INT); + INSERT INTO {table_name2} VALUES (1, 100), (2, 200); + + SELECT id, name FROM {table_name1} ORDER BY id; + SELECT id, value FROM {table_name2} ORDER BY id; + SELECT t1.name, t2.value FROM {table_name1} t1 JOIN {table_name2} t2 ON t1.id = t2.id ORDER BY t1.id; + """ + + cursor.execute(multi_select_sql) + + # First result set + results1 = cursor.fetchall() + assert len(results1) == 2, "First result set should have 2 rows" + assert results1[0][1] == 'First', "First row should contain 'First'" + + # Move to second result set + assert cursor.nextset() is True, "Should have second result set" + results2 = cursor.fetchall() + assert len(results2) == 2, "Second result set should have 2 rows" + assert results2[0][1] == 100, "First row should contain value 100" + + # Move to third result set + assert cursor.nextset() is True, "Should have third result set" + results3 = cursor.fetchall() + assert len(results3) == 2, "Third result set should have 2 rows" + assert results3[0][0] == 'First', "First row should contain 'First'" + assert results3[0][1] == 100, "First row should contain value 100" + + # Check if there are more result sets (there shouldn't be any more SELECT results) + next_result = cursor.nextset() + if next_result is not None: + # If there are more, they should be empty (from CREATE/INSERT statements) + remaining_results = cursor.fetchall() + assert len(remaining_results) == 0, "Any remaining result sets should be empty" + + except Exception as e: + pytest.fail(f"Multiple result sets with nextset test failed: {e}") + finally: + drop_table_if_exists(cursor, table_name1) + drop_table_if_exists(cursor, table_name2) + db_connection.commit() + +def test_semicolons_in_string_literals(cursor, db_connection): + """Test semicolons in string literals to ensure no false positives in buffering logic""" + table_name = "#StringTest" + try: + drop_table_if_exists(cursor, table_name) + # SQL with semicolons inside string literals - should not be treated as statement separators + sql_with_semicolons = f""" + CREATE TABLE {table_name} (id INT, data NVARCHAR(200)); + INSERT INTO {table_name} VALUES + (1, 'Value with; semicolon inside'), + (2, 'Another; value; with; multiple; semicolons'), + (3, 'Normal value'); + SELECT id, data, 'Status: OK; Processing: Complete' as status_message + FROM {table_name} + WHERE data LIKE '%semicolon%' OR data = 'Normal value' + ORDER BY id; + """ + + cursor.execute(sql_with_semicolons) + results = cursor.fetchall() + + assert len(results) == 3, "Should return 3 rows" + assert 'semicolon inside' in results[0][1], "Should preserve semicolon in string literal" + assert 'multiple; semicolons' in results[1][1], "Should preserve multiple semicolons in string literal" + assert 'Status: OK; Processing: Complete' in results[0][2], "Should preserve semicolons in status message" + + except Exception as e: + pytest.fail(f"Semicolons in string literals test failed: {e}") + finally: + drop_table_if_exists(cursor, table_name) + db_connection.commit() + +def test_multi_statement_batch_final_non_select(cursor, db_connection): + """Test multi-statement batch where the final statement is not a SELECT""" + table_name = "#BatchTest" + try: + drop_table_if_exists(cursor, table_name) + # Multi-statement batch ending with non-SELECT statement + multi_statement_non_select = f""" + CREATE TABLE {table_name} (id INT, name NVARCHAR(50), created_at DATETIME); + INSERT INTO {table_name} VALUES (1, 'Test1', GETDATE()), (2, 'Test2', GETDATE()); + SELECT COUNT(*) as record_count FROM {table_name}; + UPDATE {table_name} SET name = name + '_updated' WHERE id IN (1, 2); + """ + + cursor.execute(multi_statement_non_select) + + # Should be able to fetch results from the SELECT statement + results = cursor.fetchall() + assert len(results) == 1, "Should return 1 row from COUNT query" + assert results[0][0] == 2, "Should count 2 records" + + # Verify the UPDATE was executed by checking the updated records + cursor.execute(f"SELECT name FROM {table_name} ORDER BY id") + updated_results = cursor.fetchall() + assert len(updated_results) == 2, "Should have 2 updated records" + assert updated_results[0][0] == 'Test1_updated', "First record should be updated" + assert updated_results[1][0] == 'Test2_updated', "Second record should be updated" + + except Exception as e: + pytest.fail(f"Multi-statement batch with final non-SELECT test failed: {e}") + finally: + drop_table_if_exists(cursor, table_name) + db_connection.commit() + def test_close(db_connection): """Test closing the cursor""" try: @@ -10753,4 +10896,4 @@ def test_close(db_connection): except Exception as e: pytest.fail(f"Cursor close test failed: {e}") finally: - cursor = db_connection.cursor() + cursor = db_connection.cursor() \ No newline at end of file