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>

Author: aakash-saravanan-db

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
 

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(

src/databricks/sql/client.py

@@ -1701,6 +1701,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

tests/unit/test_client.py

@@ -303,6 +303,29 @@ def test_rowcount_resets_between_statements(self):
         cursor.execute("SELECT 1")
         self.assertEqual(cursor.rowcount, -1)
 
+    def test_async_rowcount_on_completion_and_reset(self):
+        for count in (5, 0, None):
+            with self.subTest(count=count):
+                result_set = Mock(num_modified_rows=count, is_staging_operation=False)
+                backend = ThriftDatabricksClientMockFactory.new()
+                backend.get_execution_result.return_value = result_set
+                backend.get_query_state.return_value = CommandState.SUCCEEDED
+                cursor = client.Cursor(connection=Mock(), backend=backend)
+                cursor.rowcount = 9
+
+                cursor.execute_async("UPDATE t SET x = 1")
+                self.assertEqual(cursor.rowcount, -1)
+                cursor.active_command_id = Mock(spec=CommandId)
+                self.assertIs(cursor.get_async_execution_result(), cursor)
+                self.assertEqual(cursor.rowcount, count if count is not None else -1)
+                result_set.fetchall.assert_not_called()
+
+                result_set.num_modified_rows = None
+                cursor.execute_async("SELECT 1")
+                self.assertEqual(cursor.rowcount, -1)
+                cursor.get_async_execution_result()
+                self.assertEqual(cursor.rowcount, -1)
+
     def test_closed_cursor_doesnt_allow_operations(self):
         cursor = client.Cursor(Mock(), Mock())
         cursor.close()

tests/unit/test_kernel_client.py

@@ -1233,6 +1233,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."""

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)