fix: populate rowcount after async DML completion

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.

KERNEL_REV

@@ -1 +1 @@
-c58f62c95f3651488f6270c0b45686614486d351
+57ec9a6328114a2ff2b42ae465664498d185b013

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
 

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

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

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()

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."""

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)