Support cursor.rowcount for DML on the Thrift backend (#784)

cursor.rowcount was hardcoded to -1 and never updated for the Thrift
backend (the default). For DML (INSERT/UPDATE/DELETE/MERGE) the
Databricks Thrift server reports the affected-row count in
TGetOperationStatusResp.numModifiedRows, but the connector discarded it —
_wait_until_command_done kept only operationState.

Thread numModifiedRows through: _wait_until_command_done now returns the
terminal status response, _handle_execute_response reads numModifiedRows
from it, ExecuteResponse and ResultSet carry a num_modified_rows field,
and Cursor.execute sets self.rowcount from it. rowcount resets to -1
before each statement so a DML count never leaks into a later SELECT.

SELECT (and statements the server does not report a count for) leave
rowcount at its -1 default. This brings the Thrift path in line with the
kernel backend, which already surfaces num_modified_rows.

Closes #784

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Release History
 
+# Unreleased
+- Report `cursor.rowcount` for DML on the Thrift backend: INSERT/UPDATE/DELETE/MERGE now set `rowcount` to the server's affected-row count instead of the hardcoded `-1`; SELECT (and statements the server does not report a count for) still return `-1` ([#784](https://github.com/databricks/databricks-sql-python/issues/784))
+
 # 4.3.0 (2026-06-12)
 - **New: optional Rust kernel backend (`use_kernel=True`).** Adds an alternative connection path backed by the native [`databricks-sql-kernel`](https://pypi.org/project/databricks-sql-kernel/) client (a Rust core exposed via PyO3), installable with the new `databricks-sql-connector[kernel]` extra. The kernel talks to Databricks over the **SEA (Statement Execution API) HTTP transport** — not Thrift — with CloudFetch and inline-Arrow result fetching, so `use_kernel=True` gives you a modern SEA-native client through the same DB-API surface. Supports PAT, OAuth M2M, and OAuth U2M auth. Requires Python >= 3.10 (the kernel wheel is `cp310-abi3`); on older interpreters the extra is a no-op and `use_kernel=True` raises a clear `ImportError`. The default backend remains Thrift — opt in per connection.
 - Kernel backend behavior is aligned with the Thrift backend so application code works the same either way: consistent cursor-state tracking (`query_id` / `get_query_state`), metadata (catalogs/schemas/tables/columns with JDBC-style filter semantics and case-insensitive `table_types`), DML `rowcount`, server-sourced async execution state, sync `cancel()`, fail-loud staging/volume operations, and structured error context (SQLSTATE, diagnostic info). Kernel logs surface through Python `logging` under the `databricks.sql.kernel` logger (databricks/databricks-sql-python#824, #825, #830, #838, #839 by @vikrantpuppala)

src/databricks/sql/backend/thrift_backend.py

@@ -801,7 +801,9 @@ def _hive_schema_to_description(t_table_schema, schema_bytes=None, host_url=None
             for col in t_table_schema.columns
         ]
 
-    def _results_message_to_execute_response(self, resp, operation_state):
+    def _results_message_to_execute_response(
+        self, resp, operation_state, num_modified_rows=None
+    ):
         if resp.directResults and resp.directResults.resultSetMetadata:
             t_result_set_metadata_resp = resp.directResults.resultSetMetadata
         else:
@@ -864,6 +866,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
             is_staging_operation=t_result_set_metadata_resp.isStagingOperation,
             arrow_schema_bytes=schema_bytes,
             result_format=t_result_set_metadata_resp.resultFormat,
+            num_modified_rows=num_modified_rows,
         )
 
         return execute_response, has_more_rows
@@ -945,6 +948,7 @@ def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
             self._check_command_not_in_error_or_closed_state(
                 op_handle, initial_operation_status_resp
             )
+        final_status_resp = initial_operation_status_resp
         operation_state = (
             initial_operation_status_resp
             and initial_operation_status_resp.operationState
@@ -956,7 +960,10 @@ def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
             poll_resp = self._poll_for_status(op_handle)
             operation_state = poll_resp.operationState
             self._check_command_not_in_error_or_closed_state(op_handle, poll_resp)
-        return operation_state
+            final_status_resp = poll_resp
+        # Return the terminal status response (not just the state) so callers
+        # can read ``numModifiedRows`` — the DML affected-row count — from it.
+        return operation_state, final_status_resp
 
     def get_query_state(self, command_id: CommandId) -> CommandState:
         thrift_handle = command_id.to_thrift_handle()
@@ -1274,12 +1281,21 @@ def _handle_execute_response(self, resp, cursor):
         cursor.active_command_id = command_id
         self._check_direct_results_for_error(resp.directResults, self._host)
 
-        final_operation_state = self._wait_until_command_done(
+        final_operation_state, final_status_resp = self._wait_until_command_done(
             resp.operationHandle,
             resp.directResults and resp.directResults.operationStatus,
         )
 
-        return self._results_message_to_execute_response(resp, final_operation_state)
+        # ``numModifiedRows`` is populated by the server for DML statements
+        # (INSERT/UPDATE/DELETE/MERGE) and is None for SELECT. Surface it so it
+        # can flow to ``cursor.rowcount``.
+        num_modified_rows = (
+            final_status_resp.numModifiedRows if final_status_resp else None
+        )
+
+        return self._results_message_to_execute_response(
+            resp, final_operation_state, num_modified_rows
+        )
 
     def _handle_execute_response_async(self, resp, cursor):
         command_id = CommandId.from_thrift_handle(resp.operationHandle)

src/databricks/sql/backend/types.py

@@ -425,3 +425,7 @@ class ExecuteResponse:
     is_staging_operation: bool = False
     arrow_schema_bytes: Optional[bytes] = None
     result_format: Optional[Any] = None
+    # Number of rows modified by a DML statement (INSERT/UPDATE/DELETE/MERGE),
+    # surfaced as ``cursor.rowcount``. ``None`` for SELECT and any statement
+    # for which the server does not report a count → ``rowcount`` stays at -1.
+    num_modified_rows: Optional[int] = None

src/databricks/sql/client.py

@@ -874,7 +874,10 @@ def __init__(
 
         self.connection: Connection = connection
 
-        self.rowcount: int = -1  # Return -1 as this is not supported
+        # -1 until a statement runs. Set to the affected-row count after a DML
+        # statement (INSERT/UPDATE/DELETE/MERGE); stays -1 for SELECT and any
+        # statement the server does not report a modified-row count for.
+        self.rowcount: int = -1
         self.buffer_size_bytes: int = result_buffer_size_bytes
         self.active_result_set: Union[ResultSet, None] = None
         self.arraysize: int = arraysize
@@ -1039,6 +1042,9 @@ def _close_and_clear_active_result_set(self):
                 self.active_result_set.close()
         finally:
             self.active_result_set = None
+            # Reset rowcount to its -1 default so a prior DML's count never
+            # leaks into a subsequent SELECT (or unreported) statement.
+            self.rowcount = -1
 
     def _check_not_closed(self):
         if not self.open:
@@ -1367,6 +1373,14 @@ def execute(
             query_tags=query_tags,
         )
 
+        # Surface the affected-row count for DML (INSERT/UPDATE/DELETE/MERGE) as
+        # cursor.rowcount instead of the hardcoded -1. num_modified_rows is None
+        # for SELECT (and statements the server does not report a count for) →
+        # leave rowcount at its -1 default.
+        num_modified_rows = getattr(self.active_result_set, "num_modified_rows", None)
+        if num_modified_rows is not None:
+            self.rowcount = num_modified_rows
+
         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,

src/databricks/sql/result_set.py

@@ -50,6 +50,7 @@ def __init__(
         is_staging_operation: bool = False,
         lz4_compressed: bool = False,
         arrow_schema_bytes: Optional[bytes] = None,
+        num_modified_rows: Optional[int] = None,
     ):
         """
         A ResultSet manages the results of a single command.
@@ -82,6 +83,8 @@ def __init__(
         self._is_staging_operation = is_staging_operation
         self.lz4_compressed = lz4_compressed
         self._arrow_schema_bytes = arrow_schema_bytes
+        # Affected-row count for DML; None for SELECT / unreported.
+        self.num_modified_rows = num_modified_rows
 
     def __iter__(self):
         while True:
@@ -264,6 +267,7 @@ def __init__(
             is_staging_operation=execute_response.is_staging_operation,
             lz4_compressed=execute_response.lz4_compressed,
             arrow_schema_bytes=execute_response.arrow_schema_bytes,
+            num_modified_rows=execute_response.num_modified_rows,
         )
 
         # Initialize results queue if not provided

tests/unit/test_client.py

@@ -253,6 +253,56 @@ def test_executing_multiple_commands_uses_the_most_recent_command(self):
         mock_result_sets[0].fetchall.assert_not_called()
         mock_result_sets[1].fetchall.assert_called_once_with()
 
+    def test_rowcount_reports_num_modified_rows_for_dml(self):
+        """DML sets cursor.rowcount from the result set's num_modified_rows
+        instead of the hardcoded -1 (GH #784)."""
+        mock_result_set = Mock()
+        mock_result_set.is_staging_operation = False
+        mock_result_set.num_modified_rows = 5
+
+        mock_backend = ThriftDatabricksClientMockFactory.new()
+        mock_backend.execute_command.return_value = mock_result_set
+
+        cursor = client.Cursor(connection=Mock(), backend=mock_backend)
+        cursor.execute("UPDATE t SET x = 1 WHERE y = 2")
+
+        self.assertEqual(cursor.rowcount, 5)
+
+    def test_rowcount_stays_default_for_select(self):
+        """SELECT (num_modified_rows is None) leaves rowcount at -1."""
+        mock_result_set = Mock()
+        mock_result_set.is_staging_operation = False
+        mock_result_set.num_modified_rows = None
+
+        mock_backend = ThriftDatabricksClientMockFactory.new()
+        mock_backend.execute_command.return_value = mock_result_set
+
+        cursor = client.Cursor(connection=Mock(), backend=mock_backend)
+        cursor.execute("SELECT 1")
+
+        self.assertEqual(cursor.rowcount, -1)
+
+    def test_rowcount_resets_between_statements(self):
+        """A DML count must not leak into a subsequent SELECT on the same
+        cursor — rowcount resets to -1 before each execute."""
+        dml_rs = Mock()
+        dml_rs.is_staging_operation = False
+        dml_rs.num_modified_rows = 7
+
+        select_rs = Mock()
+        select_rs.is_staging_operation = False
+        select_rs.num_modified_rows = None
+
+        mock_backend = ThriftDatabricksClientMockFactory.new()
+        mock_backend.execute_command.side_effect = [dml_rs, select_rs]
+
+        cursor = client.Cursor(connection=Mock(), backend=mock_backend)
+        cursor.execute("DELETE FROM t WHERE y = 2")
+        self.assertEqual(cursor.rowcount, 7)
+
+        cursor.execute("SELECT 1")
+        self.assertEqual(cursor.rowcount, -1)
+
     def test_closed_cursor_doesnt_allow_operations(self):
         cursor = client.Cursor(Mock(), Mock())
         cursor.close()