fix: read async affected-row counts directly

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>

Author: aakash-saravanan-db

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.

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

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
 

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