Skip to content

fix: populate rowcount after async DML completion - #1016

Open
aakash-saravanan-db wants to merge 10 commits into
mainfrom
aakash-saravanan-db/fix-async-dml-rowcount
Open

aakash-saravanan-db wants to merge 10 commits into
mainfrom
aakash-saravanan-db/fix-async-dml-rowcount

Conversation

@aakash-saravanan-db

@aakash-saravanan-db aakash-saravanan-db commented Nov 10, 2026 •

Copy link
Copy Markdown
Collaborator

What type of PR is this?

  • Bug Fix

Description

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

The kernel path reads handle.num_modified_rows after await_result(). Thrift reuses the successful terminal status observed during polling, or requests status when resuming directly by command ID. Failed or closed operations raise before result rows are fetched.

How is this tested?

  • Unit tests

  • Manually

  • Cursor, kernel-client, and Thrift unit suites: 603 passed, two skipped.

  • Eight focused connector regression cases passed. The reduced regressions fail against unmodified main at the cursor, kernel, and Thrift layers.

  • Existing kernel handle-routing tests verify count propagation after completion through owning, reader, and attached handles. Thrift tests cover status reuse, direct resume, non-terminal counts, preserved counter rows, and errors before fetch. Cursor tests cover positive, zero, and unknown counts plus reset for the next statement.

  • Eight connector tests passed with the rebuilt native kernel. Its source tree matches the merged binding commit.

  • Live SQL warehouse validation: 28/28 cases and 256 assertions passed on both kernel and Thrift backends. Covers async INSERT, UPDATE, DELETE, MERGE, CTE and bound writes, synchronous controls, counts before and after fetch, table contents, SELECT resets, and kernel reader re-fetches. Fixture schema removed.

  • Black on changed code, focused mypy checks, and diff whitespace checks passed.

Related Tickets & Documents

Requires a kernel release containing the Python binding merged in https://github.com/databricks/databricks-sql-kernel/pull/432. Source-built tests use that binding via KERNEL_REV. The connector release must set its minimum kernel version to the release containing the binding.

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread src/databricks/sql/client.py Resolved

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verdict: 1 High · 1 Low

The Thrift and cursor changes look correct, but there's one high-severity compatibility issue: on any 1.4.x wheel, which the ^1.4.0 dependency range still allows, the kernel path's direct handle.num_modified_rows read makes every async get_async_execution_result() call fail. The Thrift WeakKeyDictionary cache keys on CommandId instances, which use the default identity hash, so caching per instance works as intended.

try:
return self._make_result_set(stream, cursor, command_id)
result_set = self._make_result_set(stream, cursor, command_id)
result_set.num_modified_rows = handle.num_modified_rows

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 High — This reads handle.num_modified_rows directly. According to the PR description, that attribute on ExecutedAsyncStatement is added by kernel PR #432, which hasn't been released yet. pyproject.toml still allows databricks-sql-kernel ^1.4.0, and KERNEL_REV only affects source-built CI.

On a 1.4.x wheel, the attribute lookup raises AttributeError. The except Exception below then turns it into a DatabaseError. The result is that every get_async_execution_result() on use_kernel=True fails, including plain async SELECTs that worked before this change. The _make_result_set result is also discarded without being closed.

The sync path a few hundred lines up guards against this with getattr(executed, "num_modified_rows", None). Use the same pattern here, which degrades to rowcount=-1 on older wheels. Alternatively, bump the dependency floor in the same PR so the requirement is enforced at install time, not only in the PR text.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Kernel #432 is merged and KERNEL_REV points to its merge commit. The remaining dependency-version update requires publication of that kernel release and must land before this connector is merged or released. The source-built and native tests use that binding. Leaving this release prerequisite open.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, merging #432 and pinning KERNEL_REV to its merge commit means the source-built and native CI jobs exercise the new attribute. But pyproject.toml still allows a 1.4.x wheel, and line 970 still reads handle.num_modified_rows directly. Until the dependency floor goes up, any install that resolves to 1.4.x will fail every kernel get_async_execution_result(), including plain async SELECTs. I'm keeping this thread open as the merge gate you described, so it closes when the floor bump lands. If the release slips, the getattr(handle, "num_modified_rows", None) fallback the sync path already uses would remove the gate.

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verdict: 1 High · 1 Low

The approach is sound, but there is one high-severity compatibility problem: on the kernel path, an unguarded attribute read breaks every async result fetch on kernel wheels the current dependency floor still allows. The Thrift path (reusing the terminal status from polling, plus a refresh on resume) and the cursor changes look correct and are well tested.

try:
return self._make_result_set(stream, cursor, command_id)
result_set = self._make_result_set(stream, cursor, command_id)
result_set.num_modified_rows = handle.num_modified_rows

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 High — handle.num_modified_rows is read directly. The PR description says this attribute only exists in a kernel release that ships databricks-sql-kernel#432, but pyproject.toml still lets the resolver install ^1.4.0. On a 1.4.x wheel this line raises AttributeError. The surrounding except wraps it as a DatabaseError, so every get_async_execution_result() on the kernel path fails, including async SELECTs that worked before. The result set built just above it is also left unclosed. The sync path at line 695 already reads this attribute with getattr(..., None).

Fix: use getattr(handle, "num_modified_rows", None) so older wheels fall back to rowcount=-1, or bump the kernel floor in this same PR. Add a unit test where the handle has no num_modified_rows attribute.


with self._request_lock:
operation_status = self._terminal_operation_statuses.pop(command_id, None)
if operation_status is None:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Low — When resuming directly by command ID (no cached terminal status), get_execution_result now always sends an extra GetOperationStatus RPC before FetchResults. That changes behavior: a failed or closed operation now raises from the status check, where before the error came from the fetch. That is probably the intended fix. But the error type and message for those cases have no test beyond the single ERROR_STATE case, and nothing covers the CLOSED/CANCELED states. Consider adding a parametrized case for CLOSED_STATE and CANCELED_STATE so the new pre-fetch guard is pinned.

Carry affected-row counts from kernel handles and Thrift operation status into async result sets, then publish them on the cursor. Preserve zero, unknown counts, and fetched counter rows.

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verdict: 1 High

The Thrift and cursor changes look correct, but there is one high-severity regression on the kernel path. The async kernel result path reads handle.num_modified_rows without a guard, while the kernel version floor still allows 1.4.x wheels that lack that getter. On those wheels every async result fetch would break.

try:
return self._make_result_set(stream, cursor, command_id)
result_set = self._make_result_set(stream, cursor, command_id)
result_set.num_modified_rows = handle.num_modified_rows

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 High — handle.num_modified_rows is read directly from the kernel handle. The PR description says this getter only exists in a kernel release that contains databricks-sql-kernel#432. But pyproject.toml:73 still allows databricks-sql-kernel ^1.4.0, and this PR doesn't change it.

On a 1.4.x wheel the attribute lookup raises AttributeError. That happens inside this try, so it becomes a DatabaseError from _wrap_kernel_exception. The result is that every use_kernel=True get_async_execution_result() breaks, including SELECTs. Before this change those calls worked.

It also fails after await_result() has already consumed the owning handle and _make_result_set has built a result set. That is the 'construction fails after a successful await' path the comment above calls a narrow gap: the new result set is dropped without being closed, and the owning-handle telemetry is lost.

The sync path already handles older wheels this way (L749-753: getattr(executed, "num_modified_rows", None), commented as guarding against "an older kernel wheel that predates the pyo3 getter"). Use the same getattr fallback here, or bump the kernel floor in pyproject.toml/poetry.lock in this PR. Either way, add a unit test where the handle has no num_modified_rows attribute.

This branch was previously deployed

1 inactive deployment
azure-prod — 3c3a3034 Deployed Nov 15, 2026 by peco-review-bot[bot] via followup #2376
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants