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/KERNEL_REV b/KERNEL_REV index 7538a248a..8c6899428 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -c58f62c95f3651488f6270c0b45686614486d351 +57ec9a6328114a2ff2b42ae465664498d185b013 diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 0ca7ba6a8..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 @@ -966,7 +969,11 @@ 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) + # 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 = _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/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 19e378292..e3475ed0b 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,9 @@ def __init__( raise self._request_lock = threading.RLock() + self._terminal_operation_statuses: weakref.WeakKeyDictionary[ + CommandId, ttypes.TGetOperationStatusResp + ] = weakref.WeakKeyDictionary() self._session_id_hex = None @property @@ -927,10 +931,31 @@ 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 reuse a terminal + 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: raise ValueError("Not a valid Thrift command ID") + 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 + ) + + 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, @@ -981,6 +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=num_modified_rows, ) return ThriftResultSet( @@ -1029,6 +1055,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/src/databricks/sql/client.py b/src/databricks/sql/client.py index 6e03ed227..18b1d7838 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1701,6 +1701,11 @@ def get_async_execution_result(self): self.active_command_id, self ) + 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: 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..958c35103 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -1233,6 +1233,67 @@ 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): + 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..47ae325e9 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,113 @@ 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( + "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_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 + ), + 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=operation_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) + if cache_status: + 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()) + + 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 == expected_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)