From 182d44885b72179f65834a06418e2456be29144e Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Tue, 10 Nov 2026 03:47:21 +0000
Subject: [PATCH 01/10] 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 <aakash.saravanan@databricks.com>
---
 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 df52f34a5..6b027d26f 100644
--- a/src/databricks/sql/backend/kernel/client.py
+++ b/src/databricks/sql/backend/kernel/client.py
@@ -1024,7 +1024,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 ee904c368..7cf60de36 100644
--- a/src/databricks/sql/backend/thrift_backend.py
+++ b/src/databricks/sql/backend/thrift_backend.py
@@ -939,6 +939,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,
@@ -989,6 +994,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 613b62386..b3e4bbbaf 100755
--- a/src/databricks/sql/client.py
+++ b/src/databricks/sql/client.py
@@ -1714,6 +1714,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 f94aef0c5..0e7b2c6ed 100644
--- a/tests/unit/test_kernel_client.py
+++ b/tests/unit/test_kernel_client.py
@@ -1432,6 +1432,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 f127ad8e2..dad887bfd 100644
--- a/tests/unit/test_thrift_backend.py
+++ b/tests/unit/test_thrift_backend.py
@@ -5,7 +5,7 @@
 import logging
 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
 
@@ -143,6 +143,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 5ef8f772eccaa93d0a6fc23386f355c9eed9031b Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Tue, 10 Nov 2026 23:59:22 +0000
Subject: [PATCH 02/10] fix: harden async result rowcount retrieval

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 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 7cf60de36..93f7c4f77 100644
--- a/src/databricks/sql/backend/thrift_backend.py
+++ b/src/databricks/sql/backend/thrift_backend.py
@@ -935,6 +935,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 b3e4bbbaf..1154aa1d1 100755
--- a/src/databricks/sql/client.py
+++ b/src/databricks/sql/client.py
@@ -1714,7 +1714,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 dad887bfd..45162d964 100644
--- a/tests/unit/test_thrift_backend.py
+++ b/tests/unit/test_thrift_backend.py
@@ -188,6 +188,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 22b6979c66030205c1b96a90f0344994a5b426af Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Wed, 11 Nov 2026 00:23:18 +0000
Subject: [PATCH 03/10] docs: clarify async kernel rowcount requirement

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 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 2bae4c54e..43168b97b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,7 @@
 - Fix the kernel backend ignoring `use_cloud_fetch=False`, so complete results can be fetched without cloud-storage downloads.
 - 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))
 - Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most.
+- 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 6b027d26f..ce8f687bd 100644
--- a/src/databricks/sql/backend/kernel/client.py
+++ b/src/databricks/sql/backend/kernel/client.py
@@ -1025,6 +1025,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 8e044f0d1ef4fe784e4245d3ff74b66b4d780735 Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Wed, 11 Nov 2026 01:47:49 +0000
Subject: [PATCH 04/10] fix: reuse terminal async Thrift status

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 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 93f7c4f77..283888128 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
 
@@ -273,6 +274,7 @@ def __init__(
             raise
 
         self._request_lock = threading.RLock()
+        self._terminal_operation_statuses = weakref.WeakKeyDictionary()
         self._session_id_hex = None
 
     @property
@@ -937,17 +939,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(
@@ -1048,6 +1054,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 45162d964..06dc07eb4 100644
--- a/tests/unit/test_thrift_backend.py
+++ b/tests/unit/test_thrift_backend.py
@@ -145,7 +145,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
@@ -180,6 +181,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 bf1894f80fbb98aab20d9ddfcbeb2e94b78c6ad0 Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Wed, 11 Nov 2026 01:57:51 +0000
Subject: [PATCH 05/10] fix: annotate async status cache

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 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 283888128..3005fb6a8 100644
--- a/src/databricks/sql/backend/thrift_backend.py
+++ b/src/databricks/sql/backend/thrift_backend.py
@@ -274,7 +274,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 7ca4462d4e43a7cd64cf8eb75f1a9cdd029b7016 Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Wed, 11 Nov 2026 02:10:05 +0000
Subject: [PATCH 06/10] fix: normalize kernel rowcount access

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 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 ce8f687bd..8eb0b8099 100644
--- a/src/databricks/sql/backend/kernel/client.py
+++ b/src/databricks/sql/backend/kernel/client.py
@@ -148,6 +148,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 ("%", "*"):
@@ -749,12 +757,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
@@ -1027,7 +1030,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 0e7b2c6ed..036fd269f 100644
--- a/tests/unit/test_kernel_client.py
+++ b/tests/unit/test_kernel_client.py
@@ -1432,6 +1432,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 ace5bab5619e67f6049cca22c92929df258371cd Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Wed, 11 Nov 2026 02:16:42 +0000
Subject: [PATCH 07/10] fix: trust terminal async row counts

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 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 3005fb6a8..f771b9399 100644
--- a/src/databricks/sql/backend/thrift_backend.py
+++ b/src/databricks/sql/backend/thrift_backend.py
@@ -942,8 +942,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:
@@ -957,6 +958,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,
@@ -1007,7 +1014,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 06dc07eb4..5bfe51779 100644
--- a/tests/unit/test_thrift_backend.py
+++ b/tests/unit/test_thrift_backend.py
@@ -144,9 +144,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
@@ -164,7 +174,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)
@@ -182,12 +192,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
 
 

From e86a5d552f04984bdce54b26163359eba326208a Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Wed, 11 Nov 2026 19:59:21 +0000
Subject: [PATCH 08/10] fix: read async affected-row counts directly

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 CHANGELOG.md                                |  2 +-
 src/databricks/sql/backend/kernel/client.py | 19 +++++-------
 src/databricks/sql/client.py                |  6 ++--
 tests/unit/test_kernel_client.py            | 32 ++++-----------------
 4 files changed, 17 insertions(+), 42 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 43168b97b..d29a5f86e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,7 +4,7 @@
 - Fix the kernel backend ignoring `use_cloud_fetch=False`, so complete results can be fetched without cloud-storage downloads.
 - 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))
 - Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most.
-- 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`.
+- Report `cursor.rowcount` after asynchronous DML completes on both the Thrift and kernel backends, preserving zero counts and fetched result rows.
 
 # 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 8eb0b8099..821a63a56 100644
--- a/src/databricks/sql/backend/kernel/client.py
+++ b/src/databricks/sql/backend/kernel/client.py
@@ -148,14 +148,6 @@ 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 ("%", "*"):
@@ -757,7 +749,12 @@ 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.
-        modified = _read_num_modified_rows(executed)
+        try:
+            modified = getattr(executed, "num_modified_rows", None)
+            if callable(modified):
+                modified = modified()
+        except Exception:
+            modified = None
         if modified is not None:
             cursor.rowcount = modified
         # ``KernelResultSet.__init__`` calls ``arrow_schema()`` which
@@ -1028,9 +1025,7 @@ 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 = _read_num_modified_rows(handle)
+            result_set.num_modified_rows = handle.num_modified_rows
             return result_set
         except Exception as exc:
             raise _wrap_kernel_exception("get_execution_result", exc) from exc
diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
index 1154aa1d1..1d624d251 100755
--- a/src/databricks/sql/client.py
+++ b/src/databricks/sql/client.py
@@ -1714,8 +1714,10 @@ def get_async_execution_result(self):
                 self.active_command_id, self
             )
 
-            num_modified_rows = getattr(
-                self.active_result_set, "num_modified_rows", None
+            num_modified_rows = (
+                self.active_result_set.num_modified_rows
+                if self.active_result_set
+                else None
             )
             self.rowcount = num_modified_rows if num_modified_rows is not None else -1
 
diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py
index 036fd269f..aee87bed3 100644
--- a/tests/unit/test_kernel_client.py
+++ b/tests/unit/test_kernel_client.py
@@ -1432,26 +1432,7 @@ 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("count", [7, 0, None])
 @pytest.mark.parametrize("source", ["owner", "reader", "attached"])
 def test_async_result_count_after_await(source, count):
     c = _make_client()
@@ -1472,25 +1453,22 @@ def test_async_result_count_after_await(source, count):
     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"]
+            [[count if count is not None else 1]], names=["num_affected_rows"]
         ),
         None,
     ]
 
     def await_result():
-        if count != "unavailable":
-            handle.num_modified_rows = count
+        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)
+    assert rs.num_modified_rows == count
     stream.fetch_next_batch.assert_not_called()
-    assert [list(row) for row in rs.fetchall()] == [
-        [count if isinstance(count, int) else 1]
-    ]
+    assert [list(row) for row in rs.fetchall()] == [[count if count is not None else 1]]
 
 
 def test_get_execution_result_uses_retained_owning_handle_first():

From c9b24001fd8a043b1f3135a935ddb3397d7e417d Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Wed, 11 Nov 2026 22:16:45 +0000
Subject: [PATCH 09/10] test: focus async rowcount regression coverage

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 tests/unit/test_kernel_client.py  |  52 +++-------
 tests/unit/test_thrift_backend.py | 165 +++++++++++-------------------
 2 files changed, 71 insertions(+), 146 deletions(-)

diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py
index aee87bed3..8165fb9cb 100644
--- a/tests/unit/test_kernel_client.py
+++ b/tests/unit/test_kernel_client.py
@@ -1433,53 +1433,21 @@ def test_get_query_state_uses_reader_after_result_stream_started():
 
 
 @pytest.mark.parametrize("count", [7, 0, None])
-@pytest.mark.parametrize("source", ["owner", "reader", "attached"])
-def test_async_result_count_after_await(source, count):
+def test_get_execution_result_uses_retained_owning_handle_first(count):
+    """The first in-process result fetch uses the retained submitting
+    handle so the kernel finalizes the original async statement telemetry."""
     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
-
+    fake_stream = MagicMock()
+    fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
+    handle = MagicMock()
     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 count is not None else 1]], names=["num_affected_rows"]
-        ),
-        None,
-    ]
 
     def await_result():
         handle.num_modified_rows = count
-        return stream
+        return fake_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 == count
-    stream.fetch_next_batch.assert_not_called()
-    assert [list(row) for row in rs.fetchall()] == [[count if count is not None 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."""
-    c = _make_client()
-    c._kernel_session = MagicMock()
-    fake_stream = MagicMock()
-    fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
-    handle = MagicMock()
-    handle.await_result.return_value = fake_stream
     cursor = MagicMock()
     cursor.arraysize = 100
     cursor.buffer_size_bytes = 1024
@@ -1490,6 +1458,8 @@ def test_get_execution_result_uses_retained_owning_handle_first():
     rs = c.get_execution_result(cid, cursor=cursor)
 
     assert rs is not None
+    assert rs.num_modified_rows == count
+    fake_stream.fetch_next_batch.assert_not_called()
     c._kernel_session.attach_async_statement.assert_not_called()
     handle.await_result.assert_called_once_with()
     assert cid.guid in c._async_result_stream_started
@@ -1502,6 +1472,7 @@ def test_get_execution_result_attaches_by_id_when_no_retained_handle():
     fake_stream = MagicMock()
     fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
     handle = _attach_returns(c, await_result=fake_stream)
+    handle.num_modified_rows = 7
     cursor = MagicMock()
     cursor.arraysize = 100
     cursor.buffer_size_bytes = 1024
@@ -1510,6 +1481,7 @@ def test_get_execution_result_attaches_by_id_when_no_retained_handle():
     rs = c.get_execution_result(cid, cursor=cursor)
 
     assert rs is not None
+    assert rs.num_modified_rows == 7
     c._kernel_session.attach_async_statement.assert_called_once_with("async-1")
     handle.await_result.assert_called_once_with()
 
@@ -1564,6 +1536,7 @@ def test_get_execution_result_construction_failure_retains_marker_and_uses_reade
     retry_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
     reader = MagicMock()
     reader.await_result.return_value = retry_stream
+    reader.num_modified_rows = 0
     owning_handle.reader.return_value = reader
     cursor = MagicMock()
     cursor.arraysize = 100
@@ -1572,6 +1545,7 @@ def test_get_execution_result_construction_failure_retains_marker_and_uses_reade
     rs = c.get_execution_result(cid, cursor=cursor)
 
     assert rs is not None
+    assert rs.num_modified_rows == 0
     owning_handle.reader.assert_called_once_with()
     c._kernel_session.attach_async_statement.assert_not_called()
     reader.await_result.assert_called_once_with()
diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py
index 5bfe51779..c9ffc12b9 100644
--- a/tests/unit/test_thrift_backend.py
+++ b/tests/unit/test_thrift_backend.py
@@ -143,113 +143,6 @@ 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)
@@ -329,6 +222,64 @@ def _make_fake_thrift_backend(self):
         thrift_backend._create_arrow_table.return_value = (MagicMock(), Mock())
         return thrift_backend
 
+    @patch("databricks.sql.backend.thrift_backend.ThriftResultSet")
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
+    def test_async_result_rowcount(self, tcli_service_class, result_set):
+        for state, count, expected_count, polled in [
+            (ttypes.TOperationState.FINISHED_STATE, 7, 7, True),
+            (ttypes.TOperationState.FINISHED_STATE, 0, 0, True),
+            (ttypes.TOperationState.FINISHED_STATE, None, None, True),
+            (ttypes.TOperationState.FINISHED_STATE, 7, 7, False),
+            (ttypes.TOperationState.RUNNING_STATE, 7, None, False),
+        ]:
+            with self.subTest(state=state, count=count, polled=polled):
+                backend = self._make_fake_thrift_backend()
+                backend._poll_for_status = Mock(
+                    return_value=ttypes.TGetOperationStatusResp(
+                        status=self.okay_status,
+                        operationState=state,
+                        numModifiedRows=count,
+                    )
+                )
+                backend.make_request = Mock(
+                    return_value=Mock(
+                        resultSetMetadata=self.metadata_resp,
+                        results=sentinel.counter_rows,
+                        hasMoreRows=False,
+                    )
+                )
+                cid = CommandId.from_thrift_handle(self.operation_handle)
+                if polled:
+                    self.assertEqual(
+                        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())
+                response = result_set.call_args.kwargs["execute_response"]
+                self.assertEqual(response.num_modified_rows, expected_count)
+                self.assertIs(
+                    result_set.call_args.kwargs["t_row_set"], sentinel.counter_rows
+                )
+
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
+    def test_async_result_checks_status_before_fetch(self, tcli_service_class):
+        backend = self._make_fake_thrift_backend()
+        backend._poll_for_status = Mock(
+            return_value=ttypes.TGetOperationStatusResp(
+                status=self.okay_status,
+                operationState=ttypes.TOperationState.ERROR_STATE,
+                errorMessage="async operation failed",
+            )
+        )
+        backend.make_request = Mock()
+        cid = CommandId.from_thrift_handle(self.operation_handle)
+
+        with self.assertRaises(ServerOperationError):
+            backend.get_execution_result(cid, Mock())
+
+        backend.make_request.assert_not_called()
+
     def test_hive_schema_to_arrow_schema_preserves_column_names(self):
         columns = [
             ttypes.TColumnDesc(

From 3c3a3034804c0a878f1c40726f8eed9489c70179 Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Wed, 11 Nov 2026 22:21:07 +0000
Subject: [PATCH 10/10] docs: clarify async result status errors

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 src/databricks/sql/backend/thrift_backend.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py
index f771b9399..e6ca68d6d 100644
--- a/src/databricks/sql/backend/thrift_backend.py
+++ b/src/databricks/sql/backend/thrift_backend.py
@@ -945,6 +945,7 @@ def get_execution_result(
         status observed for the same ``CommandId`` instance by
         ``get_query_state``, or refresh the status when the caller resumes
         directly by command ID.
+        Failed or closed operations raise before result rows are fetched.
         """
         thrift_handle = command_id.to_thrift_handle()
         if not thrift_handle: