[aakash-saravanan-db]

What type of PR is this?

• Bug Fix

Description

After async DML completes, Cursor.get_async_execution_result() leaves rowcount=-1 even when the server reports an affected-row count. Populate the result-set count from the kernel handle after await_result() or from Thrift operation status, then copy it to the cursor. Preserve zero counts, unknown counts as -1, and the fetched counter rows.

Thrift reuses the terminal GetOperationStatus response already observed during normal async polling; a direct resume-by-ID fetch performs one status request. Kernel mode uses the existing cached count exposed by the companion binding change. Older kernel wheels remain supported and retain an unknown async count until upgraded; select the dependency floor and kernel pin once the binding release is available.

How is this tested?

• Unit tests

• Manually

• Connector unit suite: 1,338 passed, four skipped. New regressions fail against the unmodified connector.

• Native connector tests: four passed on standard and realtime routes with the rebuilt kernel wheel.

• Actual changed connector on staging DBSQL: 28/28 DML cases, 256 assertions across kernel and Thrift. Covers 12 async write variants per backend and positive/zero synchronous UPDATE controls, counts before/after fetch, counter rows, metadata, table contents, async SELECT resets, and kernel reader re-fetches.

• Black and diff whitespace checks passed for changed code. Fixtures were removed and the warehouse restored to STOPPED.

Related Tickets & Documents

Companion kernel Python binding: https://github.com/databricks/databricks-sql-kernel/pull/432

Release ordering

KERNEL_REV pins kernel PR #432 commit 57ec9a6328114a2ff2b42ae465664498d185b013, so source-based kernel E2E tests exercise the exact companion implementation. This test pin does not change the published ^1.4.0 dependency range, so the compatibility fallback remains necessary for currently allowed installations.

Do not merge or release this connector change until kernel #432 has been released. Once its concrete version exists, update the production kernel dependency floor and the changelog minimum version before this PR merges.

[peco-review-bot]

Verdict: 1 Low · 2 Nit

Looks good overall, with 3 low-severity points: a direct attribute access next to an existing None guard in the cursor, an extra status RPC on every Thrift async result fetch, and no test for the new error-state check. As the PR description says, the kernel count depends on a binding release that isn't out yet, and the getattr fallback safely leaves the count unknown (-1) until then. Per the repo learning log, raise the databricks-sql-kernel floor once that binding ships so the feature isn't silently off on in-range wheels.

Other findings

• ⚪ Nit — This adds an unconditional GetOperationStatus round-trip on every get_execution_result call. In the normal async flow the caller has usually just polled status through get_query_state, which makes the same RPC. The PR description acknowledges the extra call, but that earlier status response is thrown away instead of reused, so each async retrieval now costs one more network round-trip. This also affects any caller that uses get_execution_result for re-fetch. Consider caching the last TGetOperationStatusResp per command id from get_query_state and reusing it when it is terminal. Alternatively, document the extra round-trip in the method's docstring.
• ⚪ Nit — _check_command_not_in_error_or_closed_state is new on this path. A command in ERROR, CLOSED, or otherwise invalid state now raises ServerOperationError/DatabaseError before FetchResults, where before the outcome depended on whatever FetchResults returned. That is a real behavior change, and the new test only covers FINISHED_STATE. Please add a parametrized case where _poll_for_status returns ERROR_STATE (and/or CLOSED_STATE). It should assert the expected exception and that make_request for FetchResults is never called, so this guard can't be silently dropped.

[peco-review-bot]

Verdict: 1 Low

Looks good, with one low note about kernel wheel version gating. The Thrift path refreshes operation status and checks for error/closed states before calling FetchResults. The kernel path reads the cached count after await_result(), and the cursor maps a None count to -1. Tests cover positive, zero and unknown counts on both backends.

[peco-review-bot]

Verdict: 1 Low

Looks good overall, with one low concern: the kernel dependency floor isn't bumped yet to a wheel that exposes the async count. The Thrift change reuses the existing status-poll and error-state checks the same way the sync path does, and the cursor now sets -1 explicitly when the count is unknown. Tests cover 0, None and an unavailable count on both backends.

[peco-review-bot]

Verdict: 1 Low

Looks good, with one low-severity efficiency note. Async DML rowcount is now set on both backends: on Thrift from a fresh GetOperationStatus read, on the kernel through a getattr fallback that works with older wheels. The tests cover positive, zero and unknown counts, and check that counter rows are preserved and the count resets on the next statement. As the PR description says, the kernel dependency floor and the CHANGELOG minimum version still need concrete values before merge.

[peco-review-bot]

✅ No issues identified by the review bot.

[peco-review-bot]

Verdict: 1 Medium

Mostly fine, with one medium concern. The Thrift change caches the status by object identity, and the normal cursor path always calls get_query_state with the same CommandId before fetching the result, so the cache is hit. The kernel path reads the count less defensively than the existing sync path. Release ordering is not in the code yet: pyproject and the CHANGELOG still need the concrete kernel version floor before merge.

[peco-review-bot]

Verdict: 2 Low

Looks good overall, with 2 low concerns about Thrift-side status caching and refresh. The cursor, kernel and Thrift changes are consistent, keep zero counts and unknown counts (-1) distinct, and have unit tests that would fail without the change. Merge-ordering note: as the PR description says, KERNEL_REV pins an unmerged kernel commit and the CHANGELOG doesn't name a minimum kernel version. Both need updating before merge.

[peco-review-bot]

✅ No issues identified by the review bot.