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

* 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>

* test(e2e): verify DML cursor.rowcount against a live warehouse (#784)

End-to-end test in TestPySQLCoreSuite exercising INSERT/UPDATE/DELETE on
a real Thrift warehouse and asserting cursor.rowcount reports the exact
affected-row count (3/2/1), then that a following SELECT resets it to -1.
Verified live on dogfood; fails with the fix reverted (rowcount stays -1).

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

* Aggregate executemany rowcount across parameter sets (#784)

Address PR review: executemany looped execute() per parameter set and,
since each execute() resets rowcount, left rowcount equal to only the
final set's affected-row count. Per PEP 249, rowcount after executemany
should reflect the total across all operations. Accumulate the reported
per-statement counts; if no statement reports a count (all SELECT / the
server reports none), rowcount stays at its -1 default.

Adds unit tests for the sum, mixed reported/unreported, and all-unreported
cases, and extends the live e2e test to assert executemany aggregation.

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

---------

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

Author: vikrantpuppala

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 # Unreleased
 - Fix: `REMOVE` staging operations no longer require `staging_allowed_local_path` to be set, since removing a remote file does not touch the local filesystem (databricks/databricks-sql-python#726)
+- 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`. `executemany` aggregates the count across all parameter sets per PEP 249 ([#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.

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:
@@ -1373,6 +1379,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,
@@ -1511,8 +1525,22 @@ def executemany(
 
         :returns self
         """
+        # Per PEP 249, rowcount after executemany reflects the total rows
+        # affected across all parameter sets (or -1 when undeterminable). Each
+        # execute() resets self.rowcount and sets it from its own statement, so
+        # we accumulate the reported counts here. If no statement reports a
+        # count (e.g. all SELECT, or the server does not report one), rowcount
+        # stays at its -1 default.
+        total_rowcount = -1
         for parameters in seq_of_parameters:
             self.execute(operation, parameters, query_tags=query_tags)
+            if self.rowcount >= 0:
+                total_rowcount = (
+                    self.rowcount
+                    if total_rowcount < 0
+                    else total_rowcount + self.rowcount
+                )
+        self.rowcount = total_rowcount
         return self
 
     @log_latency(StatementType.METADATA)

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/e2e/test_driver.py

@@ -65,9 +65,9 @@
 for name in test_loader.getTestCaseNames(DecimalTestsMixin):
     if name.startswith("test_"):
         fn = getattr(DecimalTestsMixin, name)
-        decorated = skipUnless(pysql_supports_arrow(), "Decimal tests need arrow support")(
-            fn
-        )
+        decorated = skipUnless(
+            pysql_supports_arrow(), "Decimal tests need arrow support"
+        )(fn)
         setattr(DecimalTestsMixin, name, decorated)
 
 
@@ -152,9 +152,7 @@ def test_query_with_large_wide_result_set(self, extra_params, lz4_compression):
             cursor.connection.lz4_compression = lz4_compression
             uuids = ", ".join(["uuid() uuid{}".format(i) for i in range(cols)])
             cursor.execute(
-                "SELECT id, {uuids} FROM RANGE({rows})".format(
-                    uuids=uuids, rows=rows
-                )
+                "SELECT id, {uuids} FROM RANGE({rows})".format(uuids=uuids, rows=rows)
             )
             assert lz4_compression == cursor.active_result_set.lz4_compressed
             for row_id, row in enumerate(
@@ -455,6 +453,41 @@ def test_create_table_will_return_empty_result_set(self, extra_params):
             finally:
                 cursor.execute("DROP TABLE IF EXISTS {}".format(table_name))
 
+    def test_dml_rowcount(self):
+        """cursor.rowcount reports the affected-row count for DML
+        (INSERT/UPDATE/DELETE) instead of the hardcoded -1, and resets
+        to -1 for a subsequent SELECT (GH #784)."""
+        with self.cursor({}) as cursor:
+            table_name = "table_{uuid}".format(uuid=str(uuid4()).replace("-", "_"))
+            try:
+                cursor.execute("CREATE TABLE {} (n INT)".format(table_name))
+
+                cursor.execute("INSERT INTO {} VALUES (1), (2), (3)".format(table_name))
+                assert cursor.rowcount == 3, f"INSERT rowcount {cursor.rowcount!r}"
+
+                cursor.execute(
+                    "UPDATE {} SET n = n + 1 WHERE n >= 2".format(table_name)
+                )
+                assert cursor.rowcount == 2, f"UPDATE rowcount {cursor.rowcount!r}"
+
+                cursor.execute("DELETE FROM {} WHERE n = 1".format(table_name))
+                assert cursor.rowcount == 1, f"DELETE rowcount {cursor.rowcount!r}"
+
+                # SELECT must reset rowcount to -1 (not leak the DELETE count).
+                cursor.execute("SELECT * FROM {}".format(table_name))
+                assert cursor.rowcount == -1, f"SELECT rowcount {cursor.rowcount!r}"
+                assert len(cursor.fetchall()) == 2
+
+                # executemany aggregates the affected-row count across all
+                # parameter sets (PEP 249), not just the last one.
+                cursor.executemany(
+                    "INSERT INTO {} VALUES (%(v)s)".format(table_name),
+                    seq_of_parameters=[{"v": 10}, {"v": 20}, {"v": 30}],
+                )
+                assert cursor.rowcount == 3, f"executemany rowcount {cursor.rowcount!r}"
+            finally:
+                cursor.execute("DROP TABLE IF EXISTS {}".format(table_name))
+
     def test_get_tables(self):
         with self.cursor({}) as cursor:
             table_name = "table_{uuid}".format(uuid=str(uuid4()).replace("-", "_"))
@@ -966,7 +999,13 @@ def test_timestamps_arrow(self):
     )
     def test_multi_timestamps_arrow(self, extra_params):
         with self.cursor(
-            {"session_configuration": {"ansi_mode": False, "query_tags": "test:multi-timestamps,driver:python"}, **extra_params}
+            {
+                "session_configuration": {
+                    "ansi_mode": False,
+                    "query_tags": "test:multi-timestamps,driver:python",
+                },
+                **extra_params,
+            }
         ) as cursor:
             query, expected = self.multi_query()
             expected = [

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()
@@ -450,6 +500,8 @@ def test_executemany_parameter_passhthrough_and_uses_last_result_set(self):
         # Set is_staging_operation to False to avoid _handle_staging_operation being called
         for mock_rs in mock_result_set_instances:
             mock_rs.is_staging_operation = False
+            # SELECT statements: no modified-row count, so rowcount stays -1.
+            mock_rs.num_modified_rows = None
 
         mock_backend = ThriftDatabricksClientMockFactory.new()
         mock_backend.execute_command.side_effect = mock_result_set_instances
@@ -479,6 +531,63 @@ def test_executemany_parameter_passhthrough_and_uses_last_result_set(self):
             "last operation",
         )
 
+    def test_executemany_rowcount_sums_dml_across_parameter_sets(self):
+        """executemany aggregates rowcount across all parameter sets (PEP 249),
+        rather than reporting only the final statement's count (GH #784)."""
+        result_sets = []
+        for n in (2, 3, 5):
+            rs = Mock()
+            rs.is_staging_operation = False
+            rs.num_modified_rows = n
+            result_sets.append(rs)
+
+        mock_backend = ThriftDatabricksClientMockFactory.new()
+        mock_backend.execute_command.side_effect = result_sets
+
+        cursor = client.Cursor(Mock(), mock_backend)
+        cursor.executemany(
+            "INSERT INTO t VALUES (%(x)s)",
+            seq_of_parameters=[{"x": 1}, {"x": 2}, {"x": 3}],
+        )
+
+        self.assertEqual(cursor.rowcount, 10)
+
+    def test_executemany_rowcount_ignores_unreported_statements(self):
+        """Parameter sets the server doesn't report a count for (num_modified_rows
+        None) don't drag the aggregate down; only reported counts are summed."""
+        reported = Mock()
+        reported.is_staging_operation = False
+        reported.num_modified_rows = 4
+
+        unreported = Mock()
+        unreported.is_staging_operation = False
+        unreported.num_modified_rows = None
+
+        mock_backend = ThriftDatabricksClientMockFactory.new()
+        mock_backend.execute_command.side_effect = [reported, unreported]
+
+        cursor = client.Cursor(Mock(), mock_backend)
+        cursor.executemany("...", seq_of_parameters=[{"x": 1}, {"x": 2}])
+
+        self.assertEqual(cursor.rowcount, 4)
+
+    def test_executemany_rowcount_stays_default_when_nothing_reported(self):
+        """All-SELECT (or all-unreported) executemany leaves rowcount at -1."""
+        result_sets = []
+        for _ in range(2):
+            rs = Mock()
+            rs.is_staging_operation = False
+            rs.num_modified_rows = None
+            result_sets.append(rs)
+
+        mock_backend = ThriftDatabricksClientMockFactory.new()
+        mock_backend.execute_command.side_effect = result_sets
+
+        cursor = client.Cursor(Mock(), mock_backend)
+        cursor.executemany("SELECT %(x)s", seq_of_parameters=[{"x": 1}, {"x": 2}])
+
+        self.assertEqual(cursor.rowcount, -1)
+
     def test_setinputsizes_a_noop(self):
         cursor = client.Cursor(Mock(), Mock())
         cursor.setinputsizes(1)