From ff0b26f54e3fcfb4b423b0990a9fa26ce00c9652 Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Wed, 8 Jul 2026 07:25:32 +0000 Subject: [PATCH 1/3] Support cursor.rowcount for DML on the Thrift backend (#784) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- CHANGELOG.md | 1 + src/databricks/sql/backend/thrift_backend.py | 24 +++- src/databricks/sql/backend/types.py | 4 + src/databricks/sql/client.py | 16 ++- src/databricks/sql/result_set.py | 4 + tests/unit/test_client.py | 50 ++++++++ tests/unit/test_thrift_backend.py | 125 +++++++++++++++++-- 7 files changed, 207 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 043a6569e..18a30ff67 100644 --- a/CHANGELOG.md +++ b/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` ([#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. diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 776a7b5c8..9947051d8 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/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) diff --git a/src/databricks/sql/backend/types.py b/src/databricks/sql/backend/types.py index 5708f5e54..d080d0187 100644 --- a/src/databricks/sql/backend/types.py +++ b/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 diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index a6c7b2997..d0e7bb428 100755 --- a/src/databricks/sql/client.py +++ b/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, diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py index 54af3e534..f24a07505 100644 --- a/src/databricks/sql/result_set.py +++ b/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 diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 16e705cec..bf00c0d84 100644 --- a/tests/unit/test_client.py +++ b/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() diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index 7254b66cb..4746b18ff 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -618,7 +618,7 @@ def test_handle_execute_response_checks_operation_state_in_direct_results(self): [], auth_provider=AuthProvider(), ssl_options=SSLOptions(), - http_client=MagicMock(), + http_client=MagicMock(), ) with self.assertRaises(DatabaseError) as cm: @@ -662,7 +662,7 @@ def test_handle_execute_response_sets_compression_in_direct_results( [], auth_provider=AuthProvider(), ssl_options=SSLOptions(), - http_client=MagicMock(), + http_client=MagicMock(), ) execute_response, _ = thrift_backend._handle_execute_response( @@ -670,6 +670,98 @@ def test_handle_execute_response_sets_compression_in_direct_results( ) self.assertEqual(execute_response.lz4_compressed, lz4Compressed) + @patch( + "databricks.sql.utils.ThriftResultSetQueueFactory.build_queue", + return_value=Mock(), + ) + def test_handle_execute_response_reads_num_modified_rows_from_direct_results( + self, build_queue + ): + """DML (INSERT/UPDATE/DELETE/MERGE) surfaces the server's + numModifiedRows on ExecuteResponse so it can flow to cursor.rowcount.""" + for resp_type in self.execute_response_types: + resultSet = MagicMock() + resultSet.results.startRowOffset = 0 + t_execute_resp = resp_type( + status=Mock(), + operationHandle=Mock(), + directResults=ttypes.TSparkDirectResults( + operationStatus=ttypes.TGetOperationStatusResp( + status=self.okay_status, + operationState=ttypes.TOperationState.FINISHED_STATE, + numModifiedRows=42, + ), + resultSetMetadata=ttypes.TGetResultSetMetadataResp( + status=self.okay_status, + resultFormat=ttypes.TSparkRowSetType.ARROW_BASED_SET, + schema=MagicMock(), + arrowSchema=MagicMock(), + ), + resultSet=resultSet, + closeOperation=None, + ), + ) + thrift_backend = ThriftDatabricksClient( + "foobar", + 443, + "path", + [], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + ) + + execute_response, _ = thrift_backend._handle_execute_response( + t_execute_resp, Mock() + ) + self.assertEqual(execute_response.num_modified_rows, 42) + + @patch( + "databricks.sql.utils.ThriftResultSetQueueFactory.build_queue", + return_value=Mock(), + ) + def test_handle_execute_response_num_modified_rows_none_for_select( + self, build_queue + ): + """SELECT (and statements the server doesn't report a count for) leave + num_modified_rows as None so cursor.rowcount stays at its -1 default.""" + for resp_type in self.execute_response_types: + resultSet = MagicMock() + resultSet.results.startRowOffset = 0 + t_execute_resp = resp_type( + status=Mock(), + operationHandle=Mock(), + directResults=ttypes.TSparkDirectResults( + operationStatus=ttypes.TGetOperationStatusResp( + status=self.okay_status, + operationState=ttypes.TOperationState.FINISHED_STATE, + numModifiedRows=None, + ), + resultSetMetadata=ttypes.TGetResultSetMetadataResp( + status=self.okay_status, + resultFormat=ttypes.TSparkRowSetType.ARROW_BASED_SET, + schema=MagicMock(), + arrowSchema=MagicMock(), + ), + resultSet=resultSet, + closeOperation=None, + ), + ) + thrift_backend = ThriftDatabricksClient( + "foobar", + 443, + "path", + [], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + ) + + execute_response, _ = thrift_backend._handle_execute_response( + t_execute_resp, Mock() + ) + self.assertIsNone(execute_response.num_modified_rows) + @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True) def test_handle_execute_response_checks_operation_state_in_polls( self, tcli_service_class @@ -707,7 +799,7 @@ def test_handle_execute_response_checks_operation_state_in_polls( [], auth_provider=AuthProvider(), ssl_options=SSLOptions(), - http_client=MagicMock(), + http_client=MagicMock(), ) with self.assertRaises(DatabaseError) as cm: @@ -859,7 +951,7 @@ def test_handle_execute_response_checks_direct_results_for_error_statuses(self): [], auth_provider=AuthProvider(), ssl_options=SSLOptions(), - http_client=MagicMock(), + http_client=MagicMock(), ) with self.assertRaises(DatabaseError) as cm: @@ -912,7 +1004,7 @@ def test_handle_execute_response_can_handle_without_direct_results( [], auth_provider=AuthProvider(), ssl_options=SSLOptions(), - http_client=MagicMock(), + http_client=MagicMock(), ) ( execute_response, @@ -951,7 +1043,7 @@ def test_handle_execute_response_can_handle_with_direct_results(self): [], auth_provider=AuthProvider(), ssl_options=SSLOptions(), - http_client=MagicMock(), + http_client=MagicMock(), ) thrift_backend._results_message_to_execute_response = Mock() @@ -960,6 +1052,7 @@ def test_handle_execute_response_can_handle_with_direct_results(self): thrift_backend._results_message_to_execute_response.assert_called_with( execute_resp, ttypes.TOperationState.FINISHED_STATE, + None, ) @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True) @@ -1723,7 +1816,9 @@ def test_cancel_command_uses_active_op_handle(self, tcli_service_class): def test_handle_execute_response_sets_active_op_handle(self): thrift_backend = self._make_fake_thrift_backend() thrift_backend._check_direct_results_for_error = Mock() - thrift_backend._wait_until_command_done = Mock() + thrift_backend._wait_until_command_done = Mock( + return_value=(ttypes.TOperationState.FINISHED_STATE, Mock()) + ) thrift_backend._results_message_to_execute_response = Mock() # Create a mock response with a real operation handle @@ -2108,7 +2203,7 @@ def test_retry_args_bounding(self, mock_http_client): [], auth_provider=AuthProvider(), ssl_options=SSLOptions(), - http_client=MagicMock(), + http_client=MagicMock(), **retry_delay_args, ) retry_delay_expected_vals = { @@ -2361,10 +2456,14 @@ def test_col_to_description(self): test_cases = [ ("variant_col", {b"Spark:DataType:SqlName": b"VARIANT"}, "variant"), ("normal_col", {}, "string"), - ("weird_field", {b"Spark:DataType:SqlName": b"Some unexpected value"}, "string"), + ( + "weird_field", + {b"Spark:DataType:SqlName": b"Some unexpected value"}, + "string", + ), ("missing_field", None, "string"), # None field case ] - + for column_name, field_metadata, expected_type in test_cases: with self.subTest(column_name=column_name, expected_type=expected_type): col = ttypes.TColumnDesc( @@ -2375,7 +2474,9 @@ def test_col_to_description(self): field = ( None if field_metadata is None - else pyarrow.field(column_name, pyarrow.string(), metadata=field_metadata) + else pyarrow.field( + column_name, pyarrow.string(), metadata=field_metadata + ) ) result = ThriftDatabricksClient._col_to_description(col, field) @@ -2408,7 +2509,7 @@ def test_hive_schema_to_description(self): [("regular_col", "string")], ), ] - + for columns, arrow_fields, expected_types in test_cases: with self.subTest(arrow_fields=arrow_fields is not None): t_table_schema = ttypes.TTableSchema( From 68d57d335c30977c28a8147f9956f82c5990bd88 Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Wed, 8 Jul 2026 07:34:56 +0000 Subject: [PATCH 2/3] 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 --- tests/e2e/test_driver.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py index 45b56ae08..63a5931b9 100644 --- a/tests/e2e/test_driver.py +++ b/tests/e2e/test_driver.py @@ -455,6 +455,35 @@ 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 + 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("-", "_")) From e1e004ef30c5b5e3dfd2e61b376fb4c1600d4676 Mon Sep 17 00:00:00 2001 From: Vikrant Puppala Date: Wed, 8 Jul 2026 15:41:49 +0000 Subject: [PATCH 3/3] 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 --- CHANGELOG.md | 2 +- src/databricks/sql/client.py | 14 +++++++++ tests/e2e/test_driver.py | 30 ++++++++++++------ tests/unit/test_client.py | 59 ++++++++++++++++++++++++++++++++++++ 4 files changed, 94 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 18a30ff67..ab49a440d 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index d0e7bb428..00a19154a 100755 --- a/src/databricks/sql/client.py +++ b/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) diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py index 63a5931b9..5fe3db037 100644 --- a/tests/e2e/test_driver.py +++ b/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 = [ diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index bf00c0d84..66a3722ec 100644 --- a/tests/unit/test_client.py +++ b/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)