Support cursor.rowcount for DML on the Thrift backend (#784)#847
Support cursor.rowcount for DML on the Thrift backend (#784)#847vikrantpuppala wants to merge 3 commits into
Conversation
|
📄 src/databricks/sql/client.py
▎ executemany reports only the last statement's rowcount instead of the aggregate. executemany (client.py:1522-1523) loops over seq_of_parameters calling execute() once per set. Each execute() resets self.rowcount to -1 via _close_and_clear_active_result_set() (client.py:1047) and then sets it from that single statement's num_modified_rows (client.py:1380-1382). Now that this PR wires rowcount for DML, executemany over N DML parameter sets leaves rowcount equal to only the final set's affected-row count — not the sum across all N. PEP 249 specifies rowcount after executemany should reflect the total rows affected across all operations (or -1 when undeterminable). Before this PR rowcount was always -1, so no caller could be misled; after it, callers summing DML impact through executemany silently get an undercount. |
| 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, |
|
Good catch on the
Covered by new unit tests (sum / mixed reported+unreported / all-unreported) and the live e2e test ( |
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>
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>
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>
d21c0df to
e1e004e
Compare
Summary
cursor.rowcountwas hardcoded to-1and never updated for the Thrift backend (the default), so applications had no way to learn how many rows a DML statement affected. This was tracked in #784.For DML (
INSERT/UPDATE/DELETE/MERGE), the Databricks Thrift server reports the affected-row count inTGetOperationStatusResp.numModifiedRows, but the connector discarded it —_wait_until_command_donekept onlyoperationStateand threw the rest of the status response away.This PR wires that value through to
cursor.rowcount, bringing the Thrift path in line with the kernel backend (which already surfacesnum_modified_rows, per the 4.3.0 changelog).What changed
_wait_until_command_donenow returns the terminal status response in addition to the state, sonumModifiedRowsis preserved (from either the direct-resultsoperationStatusor the final poll response)._handle_execute_responsereadsnumModifiedRowsand passes it into_results_message_to_execute_response.ExecuteResponseandResultSetcarry a new optionalnum_modified_rowsfield;ThriftResultSetthreads it through.Cursor.executesetsself.rowcountfromactive_result_set.num_modified_rowswhen present.rowcountresets to-1before each statement (in_close_and_clear_active_result_set) so a prior DML's count never leaks into a subsequentSELECT.Behavior
rowcount= affected-row count.SELECTand any statement the server doesn't report a count for →rowcountstays-1(unchanged).rowcountinside its ownexecute_command; theclient.pyread isNone-guarded so the two coexist.TOperationHandle.modifiedRowCountfield was intentionally not used — the reliable, server-populated signal for Databricks is the I64numModifiedRowson the operation status.Test plan
tests/unit/test_client.py:test_rowcount_reports_num_modified_rows_for_dml— DML setsrowcountto the affected count.test_rowcount_stays_default_for_select—SELECTleavesrowcountat-1.test_rowcount_resets_between_statements— a DML count does not leak into a laterSELECTon the same cursor.tests/unit/test_thrift_backend.py:test_handle_execute_response_reads_num_modified_rows_from_direct_results—numModifiedRowsreachesExecuteResponse.num_modified_rows.test_handle_execute_response_num_modified_rows_none_for_select—NonestaysNone.test_handle_execute_response_can_handle_with_direct_resultsandtest_handle_execute_response_sets_active_op_handlefor the new_wait_until_command_donereturn shape /_results_message_to_execute_responsesignature.rowcount==-1) and pass with it.872 passed, 4 skipped(plusrealkernelset skipped — requires the real kernel wheel).black --checkclean on all changed files.Closes #784