From c481be89b9056c86f52a4585bcb13a3cf97bc8a2 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/11] 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 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 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 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 268a890f251e35974392af0185c8af312f969f64 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/11] 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 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 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 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 a70462eecdfc71e628e2be9efcd50f6a162f87c4 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/11] 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 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 9056c1f5aa8f1d83b3e8bd998669628ef042ef4f Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Wed, 11 Nov 2026 01:41:55 +0000
Subject: [PATCH 04/11] test: pin async rowcount kernel revision

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 KERNEL_REV | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/KERNEL_REV b/KERNEL_REV
index d025be44a..74d6f39fb 100644
--- a/KERNEL_REV
+++ b/KERNEL_REV
@@ -1 +1 @@
-6f4b6d607a61bdde61dcda2cf82bc8e80fd32cb9
+426993a6838aebb7d37b8fc89fa606b4897dd3f2

From c4ab0fd0ff6db24d4f28f4bc7cc1bb94694b92d5 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 05/11] 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 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 59879fc99257e93949edd804c418ad66f54ea9e6 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 06/11] 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 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 e0212c3112b6956fdbbed78c7919620bb05db9bc 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 07/11] 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 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 1dee6833167ec88cf4fbb03d0d119d54bfa1a646 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 08/11] 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 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
 
 

From f06979fae598812fd0d60ae4a8bb78a9303a7240 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 09/11] 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 fb6037ddd..5c9cae5b0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,7 +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`.
+- 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 20545dafc..5581f0325 100644
--- a/src/databricks/sql/backend/kernel/client.py
+++ b/src/databricks/sql/backend/kernel/client.py
@@ -138,14 +138,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 ("%", "*"):
@@ -699,7 +691,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
@@ -970,9 +967,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 958c35103..8903557ff 100644
--- a/tests/unit/test_kernel_client.py
+++ b/tests/unit/test_kernel_client.py
@@ -1233,26 +1233,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()
@@ -1273,25 +1254,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 14466cf89c633b2864008fa1f55df2ad578516a8 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 10/11] test: focus async rowcount regression coverage

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

diff --git a/KERNEL_REV b/KERNEL_REV
index 74d6f39fb..bacdc519c 100644
--- a/KERNEL_REV
+++ b/KERNEL_REV
@@ -1 +1 @@
-426993a6838aebb7d37b8fc89fa606b4897dd3f2
+453017f8c305e083e14b9b85183bc00920523ad5
diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py
index 8903557ff..ff4bc3f2d 100644
--- a/tests/unit/test_kernel_client.py
+++ b/tests/unit/test_kernel_client.py
@@ -1234,53 +1234,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
@@ -1291,6 +1259,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
@@ -1303,6 +1273,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
@@ -1311,6 +1282,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()
 
@@ -1365,6 +1337,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
@@ -1373,6 +1346,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 47ae325e9..10e37c6ae 100644
--- a/tests/unit/test_thrift_backend.py
+++ b/tests/unit/test_thrift_backend.py
@@ -134,113 +134,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)
@@ -320,6 +213,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 6464c30f3e8b8bb8be92e8924137001bfda5cb3f 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 11/11] 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 e3475ed0b..40ffdb797 100644
--- a/src/databricks/sql/backend/thrift_backend.py
+++ b/src/databricks/sql/backend/thrift_backend.py
@@ -937,6 +937,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: