From 43a2c93fa479e11b2b7b239c65d10354ba5ade55 Mon Sep 17 00:00:00 2001 From: Aakash Saravanan Date: Tue, 10 Nov 2026 03:47:21 +0000 Subject: [PATCH 1/8] fix: populate rowcount after async DML completion Carry affected-row counts from kernel handles and Thrift operation status into async result sets, then publish them on the cursor. Preserve zero, unknown counts, and fetched counter rows. Signed-off-by: Aakash Saravanan --- src/databricks/sql/backend/kernel/client.py | 4 +- src/databricks/sql/backend/thrift_backend.py | 6 +++ src/databricks/sql/client.py | 3 ++ tests/unit/test_client.py | 23 ++++++++++ tests/unit/test_kernel_client.py | 42 +++++++++++++++++ tests/unit/test_thrift_backend.py | 47 +++++++++++++++++++- 6 files changed, 123 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 0ca7ba6a8..012710540 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -966,7 +966,9 @@ def get_execution_result( # telemetry — an accepted, narrow gap limited to the case where # result-set construction fails after a successful await. try: - return self._make_result_set(stream, cursor, command_id) + result_set = self._make_result_set(stream, cursor, command_id) + result_set.num_modified_rows = getattr(handle, "num_modified_rows", None) + return result_set except Exception as exc: raise _wrap_kernel_exception("get_execution_result", exc) from exc diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 19e378292..9439861e1 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -931,6 +931,11 @@ def get_execution_result( if not thrift_handle: raise ValueError("Not a valid Thrift command ID") + operation_status = self._poll_for_status(thrift_handle) + self._check_command_not_in_error_or_closed_state( + thrift_handle, operation_status + ) + req = ttypes.TFetchResultsReq( operationHandle=ttypes.TOperationHandle( thrift_handle.operationId, @@ -981,6 +986,7 @@ def get_execution_result( is_staging_operation=is_staging_operation, arrow_schema_bytes=schema_bytes, result_format=t_result_set_metadata_resp.resultFormat, + num_modified_rows=operation_status.numModifiedRows, ) return ThriftResultSet( diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 6e03ed227..0ae0afea1 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1701,6 +1701,9 @@ def get_async_execution_result(self): self.active_command_id, self ) + num_modified_rows = self.active_result_set.num_modified_rows + self.rowcount = num_modified_rows if num_modified_rows is not None else -1 + if self.active_result_set and self.active_result_set.is_staging_operation: self._handle_staging_operation( staging_allowed_local_path=self.connection.staging_allowed_local_path diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 047252488..b9c42dd7c 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -303,6 +303,29 @@ def test_rowcount_resets_between_statements(self): cursor.execute("SELECT 1") self.assertEqual(cursor.rowcount, -1) + def test_async_rowcount_on_completion_and_reset(self): + for count in (5, 0, None): + with self.subTest(count=count): + result_set = Mock(num_modified_rows=count, is_staging_operation=False) + backend = ThriftDatabricksClientMockFactory.new() + backend.get_execution_result.return_value = result_set + backend.get_query_state.return_value = CommandState.SUCCEEDED + cursor = client.Cursor(connection=Mock(), backend=backend) + cursor.rowcount = 9 + + cursor.execute_async("UPDATE t SET x = 1") + self.assertEqual(cursor.rowcount, -1) + cursor.active_command_id = Mock(spec=CommandId) + self.assertIs(cursor.get_async_execution_result(), cursor) + self.assertEqual(cursor.rowcount, count if count is not None else -1) + result_set.fetchall.assert_not_called() + + result_set.num_modified_rows = None + cursor.execute_async("SELECT 1") + self.assertEqual(cursor.rowcount, -1) + cursor.get_async_execution_result() + self.assertEqual(cursor.rowcount, -1) + def test_closed_cursor_doesnt_allow_operations(self): cursor = client.Cursor(Mock(), Mock()) cursor.close() diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index b285ce6e1..d8dacdbc9 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -1233,6 +1233,48 @@ def test_get_query_state_uses_reader_after_result_stream_started(): reader.status.assert_called_once_with() +@pytest.mark.parametrize("count", [7, 0, None, "unavailable"]) +@pytest.mark.parametrize("source", ["owner", "reader", "attached"]) +def test_async_result_count_after_await(source, count): + c = _make_client() + c._kernel_session = MagicMock() + cid = CommandId.from_sea_statement_id("async-count") + owner = MagicMock() + handle = owner + if source == "attached": + c._kernel_session.attach_async_statement.return_value = handle + else: + c._async_handles[cid.guid] = owner + if source == "reader": + c._async_result_stream_started.add(cid.guid) + handle = owner.reader.return_value + + del handle.num_modified_rows + stream = MagicMock() + stream.arrow_schema.return_value = pa.schema([("num_affected_rows", pa.int64())]) + stream.fetch_next_batch.side_effect = [ + pa.record_batch( + [[count if isinstance(count, int) else 1]], names=["num_affected_rows"] + ), + None, + ] + + def await_result(): + if count != "unavailable": + handle.num_modified_rows = count + return stream + + handle.await_result.side_effect = await_result + cursor = MagicMock(arraysize=100, buffer_size_bytes=1024) + rs = c.get_execution_result(cid, cursor) + + assert rs.num_modified_rows == (None if count == "unavailable" else count) + stream.fetch_next_batch.assert_not_called() + assert [list(row) for row in rs.fetchall()] == [ + [count if isinstance(count, int) else 1] + ] + + def test_get_execution_result_uses_retained_owning_handle_first(): """The first in-process result fetch uses the retained submitting handle so the kernel finalizes the original async statement telemetry.""" diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index dedadaeb2..a51c0a9d5 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -3,7 +3,7 @@ import itertools import unittest import pytest -from unittest.mock import patch, MagicMock, Mock +from unittest.mock import patch, MagicMock, Mock, sentinel from ssl import CERT_NONE, CERT_REQUIRED from urllib3 import HTTPSConnectionPool @@ -134,6 +134,51 @@ def test_metadata_catalog_escaping( assert getattr(request, thrift_field) == value +@pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") +@pytest.mark.parametrize("count", [7, 0, None]) +def test_async_result_reads_status_count_and_preserves_rows(count): + with ( + patch( + "databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True + ), + patch("databricks.sql.backend.thrift_backend.ThriftResultSet") as result_set, + ): + backend = ThriftDatabricksClient( + "foobar", + 443, + "path", + [], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + ) + operation_status = ttypes.TGetOperationStatusResp( + status=ThriftBackendTestSuite.okay_status, + operationState=ttypes.TOperationState.FINISHED_STATE, + numModifiedRows=count, + ) + backend._poll_for_status = Mock(return_value=operation_status) + metadata = ttypes.TGetResultSetMetadataResp( + schema=ttypes.TTableSchema(columns=[]), + arrowSchema=pyarrow.schema([]).serialize().to_pybytes(), + resultFormat=ttypes.TSparkRowSetType.ARROW_BASED_SET, + ) + backend.make_request = Mock( + return_value=Mock( + resultSetMetadata=metadata, + results=sentinel.counter_rows, + hasMoreRows=False, + ) + ) + cid = CommandId.from_thrift_handle(ThriftBackendTestSuite.operation_handle) + backend.get_execution_result(cid, Mock()) + + backend._poll_for_status.assert_called_once_with(cid.to_thrift_handle()) + response = result_set.call_args.kwargs["execute_response"] + assert response.num_modified_rows == count + assert result_set.call_args.kwargs["t_row_set"] is sentinel.counter_rows + + @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") class ThriftBackendTestSuite(unittest.TestCase): okay_status = ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS) From de947bf901d60815d3700c7ef765315f7a594135 Mon Sep 17 00:00:00 2001 From: Aakash Saravanan Date: Tue, 10 Nov 2026 23:59:22 +0000 Subject: [PATCH 2/8] fix: harden async result rowcount retrieval Signed-off-by: Aakash Saravanan --- src/databricks/sql/backend/thrift_backend.py | 5 +++ src/databricks/sql/client.py | 4 ++- tests/unit/test_thrift_backend.py | 38 ++++++++++++++++++++ 3 files changed, 46 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 9439861e1..d37f59e03 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -927,6 +927,11 @@ def _results_message_to_execute_response( def get_execution_result( self, command_id: CommandId, cursor: Cursor ) -> "ResultSet": + """Fetch a completed async result and its latest operation metadata. + + FetchResults does not carry ``numModifiedRows``, so Thrift refreshes + operation status before constructing the result set. + """ thrift_handle = command_id.to_thrift_handle() if not thrift_handle: raise ValueError("Not a valid Thrift command ID") diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 0ae0afea1..18b1d7838 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1701,7 +1701,9 @@ def get_async_execution_result(self): self.active_command_id, self ) - num_modified_rows = self.active_result_set.num_modified_rows + num_modified_rows = getattr( + self.active_result_set, "num_modified_rows", None + ) self.rowcount = num_modified_rows if num_modified_rows is not None else -1 if self.active_result_set and self.active_result_set.is_staging_operation: diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index a51c0a9d5..f48a2cf9b 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -179,6 +179,44 @@ def test_async_result_reads_status_count_and_preserves_rows(count): assert result_set.call_args.kwargs["t_row_set"] is sentinel.counter_rows +@pytest.mark.parametrize( + "operation_state,exception_type", + [ + (ttypes.TOperationState.ERROR_STATE, ServerOperationError), + (ttypes.TOperationState.CLOSED_STATE, DatabaseError), + ], +) +def test_async_result_rejects_terminal_error_state_before_fetch( + operation_state, exception_type +): + with patch( + "databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True + ): + backend = ThriftDatabricksClient( + "foobar", + 443, + "path", + [], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + ) + backend._poll_for_status = Mock( + return_value=ttypes.TGetOperationStatusResp( + status=ThriftBackendTestSuite.okay_status, + operationState=operation_state, + errorMessage="async operation failed", + ) + ) + backend.make_request = Mock() + cid = CommandId.from_thrift_handle(ThriftBackendTestSuite.operation_handle) + + with pytest.raises(exception_type): + backend.get_execution_result(cid, Mock()) + + backend.make_request.assert_not_called() + + @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") class ThriftBackendTestSuite(unittest.TestCase): okay_status = ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS) From fac2560888c3c33bdd8e80e291f2485b4202b280 Mon Sep 17 00:00:00 2001 From: Aakash Saravanan Date: Wed, 11 Nov 2026 00:23:18 +0000 Subject: [PATCH 3/8] docs: clarify async kernel rowcount requirement Signed-off-by: Aakash Saravanan --- CHANGELOG.md | 1 + src/databricks/sql/backend/kernel/client.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4fbb64b68..fb6037ddd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ # Unreleased - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) +- Report `cursor.rowcount` after asynchronous DML completes. With `use_kernel=True`, this requires a kernel release containing async affected-row count support; kernel 1.4.0 continues to report `-1`. # 4.9.0 (2026-11-04) - Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes. diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 012710540..98d0baf11 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -967,6 +967,8 @@ def get_execution_result( # result-set construction fails after a successful await. try: result_set = self._make_result_set(stream, cursor, command_id) + # Older kernel wheels do not expose this value. Keep them compatible + # while allowing a newer binding to populate async rowcount. result_set.num_modified_rows = getattr(handle, "num_modified_rows", None) return result_set except Exception as exc: From ac1ada2ea8f64fde3c472650853ecf807d3b93cd Mon Sep 17 00:00:00 2001 From: Aakash Saravanan Date: Wed, 11 Nov 2026 01:41:55 +0000 Subject: [PATCH 4/8] test: pin async rowcount kernel revision Signed-off-by: Aakash Saravanan --- KERNEL_REV | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/KERNEL_REV b/KERNEL_REV index 7538a248a..8c6899428 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -c58f62c95f3651488f6270c0b45686614486d351 +57ec9a6328114a2ff2b42ae465664498d185b013 From 2736b2d1d944a1ef578aba8d15a2d122d9d86a2d Mon Sep 17 00:00:00 2001 From: Aakash Saravanan Date: Wed, 11 Nov 2026 01:47:49 +0000 Subject: [PATCH 5/8] fix: reuse terminal async Thrift status Signed-off-by: Aakash Saravanan --- src/databricks/sql/backend/thrift_backend.py | 21 ++++++++++++++------ tests/unit/test_thrift_backend.py | 5 ++++- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index d37f59e03..f99bc5779 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -6,6 +6,7 @@ import re import time import threading +import weakref from typing import Dict, List, Optional, Union, Any, TYPE_CHECKING from uuid import UUID @@ -267,6 +268,7 @@ def __init__( raise self._request_lock = threading.RLock() + self._terminal_operation_statuses = weakref.WeakKeyDictionary() self._session_id_hex = None @property @@ -929,17 +931,21 @@ def get_execution_result( ) -> "ResultSet": """Fetch a completed async result and its latest operation metadata. - FetchResults does not carry ``numModifiedRows``, so Thrift refreshes - operation status before constructing the result set. + FetchResults does not carry ``numModifiedRows``, so reuse a terminal + status observed by ``get_query_state`` or refresh it when the caller + resumes directly by command ID. """ thrift_handle = command_id.to_thrift_handle() if not thrift_handle: raise ValueError("Not a valid Thrift command ID") - operation_status = self._poll_for_status(thrift_handle) - self._check_command_not_in_error_or_closed_state( - thrift_handle, operation_status - ) + with self._request_lock: + operation_status = self._terminal_operation_statuses.pop(command_id, None) + if operation_status is None: + operation_status = self._poll_for_status(thrift_handle) + self._check_command_not_in_error_or_closed_state( + thrift_handle, operation_status + ) req = ttypes.TFetchResultsReq( operationHandle=ttypes.TOperationHandle( @@ -1040,6 +1046,9 @@ def get_query_state(self, command_id: CommandId) -> CommandState: state = CommandState.from_thrift_state(operation_state) if state is None: raise ValueError(f"Unknown command state: {operation_state}") + if state == CommandState.SUCCEEDED: + with self._request_lock: + self._terminal_operation_statuses[command_id] = poll_resp return state @staticmethod diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index f48a2cf9b..ee016b235 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -136,7 +136,8 @@ def test_metadata_catalog_escaping( @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") @pytest.mark.parametrize("count", [7, 0, None]) -def test_async_result_reads_status_count_and_preserves_rows(count): +@pytest.mark.parametrize("cache_status", [True, False]) +def test_async_result_reads_status_count_and_preserves_rows(count, cache_status): with ( patch( "databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True @@ -171,6 +172,8 @@ def test_async_result_reads_status_count_and_preserves_rows(count): ) ) cid = CommandId.from_thrift_handle(ThriftBackendTestSuite.operation_handle) + if cache_status: + assert backend.get_query_state(cid) == CommandState.SUCCEEDED backend.get_execution_result(cid, Mock()) backend._poll_for_status.assert_called_once_with(cid.to_thrift_handle()) From d2b58f1315a534474bbf279c9e33136e7a921542 Mon Sep 17 00:00:00 2001 From: Aakash Saravanan Date: Wed, 11 Nov 2026 01:57:51 +0000 Subject: [PATCH 6/8] fix: annotate async status cache Signed-off-by: Aakash Saravanan --- src/databricks/sql/backend/thrift_backend.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index f99bc5779..8216f49e3 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -268,7 +268,9 @@ def __init__( raise self._request_lock = threading.RLock() - self._terminal_operation_statuses = weakref.WeakKeyDictionary() + self._terminal_operation_statuses: weakref.WeakKeyDictionary[ + CommandId, ttypes.TGetOperationStatusResp + ] = weakref.WeakKeyDictionary() self._session_id_hex = None @property From 0b93e035e79362026641af43d94a3766edf96d43 Mon Sep 17 00:00:00 2001 From: Aakash Saravanan Date: Wed, 11 Nov 2026 02:10:05 +0000 Subject: [PATCH 7/8] fix: normalize kernel rowcount access Signed-off-by: Aakash Saravanan --- src/databricks/sql/backend/kernel/client.py | 17 ++++++++++------- tests/unit/test_kernel_client.py | 19 +++++++++++++++++++ 2 files changed, 29 insertions(+), 7 deletions(-) diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 98d0baf11..20545dafc 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -138,6 +138,14 @@ def _async_statement_reader(owner: Any) -> Any: return reader() +def _read_num_modified_rows(handle: Any) -> Optional[int]: + try: + modified = getattr(handle, "num_modified_rows", None) + return modified() if callable(modified) else modified + except Exception: + return None + + def _catalog_or_none(value: Optional[str]) -> Optional[str]: """Map supported all-catalog wildcards to the kernel's unset filter.""" if value is None or value in ("%", "*"): @@ -691,12 +699,7 @@ def execute_command( # default. ``getattr`` guards against an older kernel wheel that # predates the pyo3 getter. NB the Thrift backend also hardcodes # ``-1`` here, so this makes the kernel path *exceed* Thrift. - try: - modified = getattr(executed, "num_modified_rows", None) - if callable(modified): - modified = modified() - except Exception: - modified = None + modified = _read_num_modified_rows(executed) if modified is not None: cursor.rowcount = modified # ``KernelResultSet.__init__`` calls ``arrow_schema()`` which @@ -969,7 +972,7 @@ def get_execution_result( result_set = self._make_result_set(stream, cursor, command_id) # Older kernel wheels do not expose this value. Keep them compatible # while allowing a newer binding to populate async rowcount. - result_set.num_modified_rows = getattr(handle, "num_modified_rows", None) + result_set.num_modified_rows = _read_num_modified_rows(handle) return result_set except Exception as exc: raise _wrap_kernel_exception("get_execution_result", exc) from exc diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index d8dacdbc9..958c35103 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -1233,6 +1233,25 @@ def test_get_query_state_uses_reader_after_result_stream_started(): reader.status.assert_called_once_with() +class _RaisingModifiedRows: + @property + def num_modified_rows(self): + raise RuntimeError("unavailable") + + +@pytest.mark.parametrize( + ("handle", "expected"), + [ + (types.SimpleNamespace(num_modified_rows=7), 7), + (types.SimpleNamespace(num_modified_rows=lambda: 8), 8), + (types.SimpleNamespace(), None), + (_RaisingModifiedRows(), None), + ], +) +def test_read_num_modified_rows_compatibility_shapes(handle, expected): + assert kernel_client._read_num_modified_rows(handle) == expected + + @pytest.mark.parametrize("count", [7, 0, None, "unavailable"]) @pytest.mark.parametrize("source", ["owner", "reader", "attached"]) def test_async_result_count_after_await(source, count): From 5f76129478fd3eee077829ffff022711c493c310 Mon Sep 17 00:00:00 2001 From: Aakash Saravanan Date: Wed, 11 Nov 2026 02:16:42 +0000 Subject: [PATCH 8/8] fix: trust terminal async row counts Signed-off-by: Aakash Saravanan --- src/databricks/sql/backend/thrift_backend.py | 13 ++++++-- tests/unit/test_thrift_backend.py | 33 ++++++++++++++++---- 2 files changed, 37 insertions(+), 9 deletions(-) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 8216f49e3..e3475ed0b 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -934,8 +934,9 @@ def get_execution_result( """Fetch a completed async result and its latest operation metadata. FetchResults does not carry ``numModifiedRows``, so reuse a terminal - status observed by ``get_query_state`` or refresh it when the caller - resumes directly by command ID. + status observed for the same ``CommandId`` instance by + ``get_query_state``, or refresh the status when the caller resumes + directly by command ID. """ thrift_handle = command_id.to_thrift_handle() if not thrift_handle: @@ -949,6 +950,12 @@ def get_execution_result( thrift_handle, operation_status ) + num_modified_rows = ( + operation_status.numModifiedRows + if operation_status.operationState == ttypes.TOperationState.FINISHED_STATE + else None + ) + req = ttypes.TFetchResultsReq( operationHandle=ttypes.TOperationHandle( thrift_handle.operationId, @@ -999,7 +1006,7 @@ def get_execution_result( is_staging_operation=is_staging_operation, arrow_schema_bytes=schema_bytes, result_format=t_result_set_metadata_resp.resultFormat, - num_modified_rows=operation_status.numModifiedRows, + num_modified_rows=num_modified_rows, ) return ThriftResultSet( diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index ee016b235..47ae325e9 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -135,9 +135,19 @@ def test_metadata_catalog_escaping( @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") -@pytest.mark.parametrize("count", [7, 0, None]) +@pytest.mark.parametrize( + "operation_state,count,expected_count", + [ + (ttypes.TOperationState.FINISHED_STATE, 7, 7), + (ttypes.TOperationState.FINISHED_STATE, 0, 0), + (ttypes.TOperationState.FINISHED_STATE, None, None), + (ttypes.TOperationState.RUNNING_STATE, 7, None), + ], +) @pytest.mark.parametrize("cache_status", [True, False]) -def test_async_result_reads_status_count_and_preserves_rows(count, cache_status): +def test_async_result_reads_only_terminal_status_count_and_preserves_rows( + operation_state, count, expected_count, cache_status +): with ( patch( "databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True @@ -155,7 +165,7 @@ def test_async_result_reads_status_count_and_preserves_rows(count, cache_status) ) operation_status = ttypes.TGetOperationStatusResp( status=ThriftBackendTestSuite.okay_status, - operationState=ttypes.TOperationState.FINISHED_STATE, + operationState=operation_state, numModifiedRows=count, ) backend._poll_for_status = Mock(return_value=operation_status) @@ -173,12 +183,23 @@ def test_async_result_reads_status_count_and_preserves_rows(count, cache_status) ) cid = CommandId.from_thrift_handle(ThriftBackendTestSuite.operation_handle) if cache_status: - assert backend.get_query_state(cid) == CommandState.SUCCEEDED + expected_state = ( + CommandState.SUCCEEDED + if operation_state == ttypes.TOperationState.FINISHED_STATE + else CommandState.RUNNING + ) + assert backend.get_query_state(cid) == expected_state backend.get_execution_result(cid, Mock()) - backend._poll_for_status.assert_called_once_with(cid.to_thrift_handle()) + expected_polls = ( + 2 + if cache_status and operation_state == ttypes.TOperationState.RUNNING_STATE + else 1 + ) + assert backend._poll_for_status.call_count == expected_polls + backend._poll_for_status.assert_called_with(cid.to_thrift_handle()) response = result_set.call_args.kwargs["execute_response"] - assert response.num_modified_rows == count + assert response.num_modified_rows == expected_count assert result_set.call_args.kwargs["t_row_set"] is sentinel.counter_rows