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>

Author: vikrantpuppala

CHANGELOG.md

@@ -2,7 +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` ([#784](https://github.com/databricks/databricks-sql-python/issues/784))
+- 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/client.py

@@ -1525,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)

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(
@@ -464,9 +462,7 @@ def test_dml_rowcount(self):
             try:
                 cursor.execute("CREATE TABLE {} (n INT)".format(table_name))
 
-                cursor.execute(
-                    "INSERT INTO {} VALUES (1), (2), (3)".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(
@@ -481,6 +477,14 @@ def test_dml_rowcount(self):
                 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))
 
@@ -995,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

@@ -500,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
@@ -529,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)