GV

GeezerVison

The World Wide Web, restored to 800 pixels of glory.

Home|About|Gaming cabinet|Archives|How it works|Starter sites

On this site

  • Sign in
  • ActionsCI/CD and automation
  • CodespacesCloud development environments
  • All Features
  • Documentation
  • Training
  • Certifications
  • Build with Copilot
  • Customize Copilot
  • Integrations
  • Trending repos
  • GitHub Sponsors
  • Developer articles
  • Enterprises

This view

  • Original page
  • New address

GeezerVison is showing github.com — fix: populate rowcount after async DML completion by aakash-saravanan-db · Pull Request #1016 · databricks/databricks-sql-python

fix: populate rowcount after async DML completion #1016

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aakash-saravanan-db wants to merge 11 commits into main
base: main
Loading
Loading
from aakash-saravanan-db/fix-async-dml-rowcount

Uh oh!

There was an error while loading. Please reload this page.

Open

fix: populate rowcount after async DML completion #1016

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
c481be8
fix: populate rowcount after async DML completion
aakash-saravanan-db Nov 10, 2026
268a890
fix: harden async result rowcount retrieval
aakash-saravanan-db Nov 10, 2026
a70462e
docs: clarify async kernel rowcount requirement
aakash-saravanan-db Nov 11, 2026
9056c1f
test: pin async rowcount kernel revision
aakash-saravanan-db Nov 11, 2026
c4ab0fd
fix: reuse terminal async Thrift status
aakash-saravanan-db Nov 11, 2026
59879fc
fix: annotate async status cache
aakash-saravanan-db Nov 11, 2026
e0212c3
fix: normalize kernel rowcount access
aakash-saravanan-db Nov 11, 2026
1dee683
fix: trust terminal async row counts
aakash-saravanan-db Nov 11, 2026
f06979f
fix: read async affected-row counts directly
aakash-saravanan-db Nov 11, 2026
14466cf
test: focus async rowcount regression coverage
aakash-saravanan-db Nov 11, 2026
6464c30
docs: clarify async result status errors
aakash-saravanan-db Nov 11, 2026
File filter
Filter by extension

Clear filters
Conversations
Failed to load comments.
Loading

Uh oh!

There was an error while loading. Please reload this page.

Jump to
Jump to file
Failed to load files.
Loading

Uh oh!

There was an error while loading. Please reload this page.

Diff view
Diff view

There are no files selected for viewing

1 change: 1 addition & 0 deletions 1 CHANGELOG.md
Show comments
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

# Unreleased
- Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629))
- Report `cursor.rowcount` after asynchronous DML completes on both the Thrift and kernel backends, preserving zero counts and fetched result rows.

# 4.9.0 (2026-11-04)
- Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes.
Expand Down
2 changes: 1 addition & 1 deletion 2 KERNEL_REV
Show comments
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1 +1 @@
6f4b6d607a61bdde61dcda2cf82bc8e80fd32cb9
453017f8c305e083e14b9b85183bc00920523ad5
4 changes: 3 additions & 1 deletion 4 src/databricks/sql/backend/kernel/client.py
Show comments
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -966,7 +966,9 @@ def get_execution_result(
# telemetry — an accepted, narrow gap limited to the case where
# result-set construction fails after a successful await.
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
Comment thread
aakash-saravanan-db marked this conversation as resolved.
Resolved

Uh oh!

There was an error while loading. Please reload this page.

Comment thread

@peco-review-bot peco-review-bot Bot Nov 11, 2026

Copy link
Copy Markdown

There was a problem hiding this comment.

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, but that attribute only exists in the companion binding (databricks-sql-kernel#432). The dependency floor is still ^1.4.0, and the PR description says the floor will be raised later. On any wheel the resolver currently allows, this line raises AttributeError. The surrounding except Exception then turns it into a DatabaseError. The result is that every get_async_execution_result() on use_kernel=True fails, including plain async SELECTs, which work today. It also leaks the result set _make_result_set already built, since that result set (and its stream) is never closed.

The sync path in this file already handles this safely with getattr(executed, "num_modified_rows", None) (line 695). Either do the same here so older wheels just report rowcount=-1, or bump the pyproject floor in this same PR. Consider also adding a unit test where the handle has no num_modified_rows attribute. The new parametrized test always sets it, so the old-wheel case is never exercised.

Sorry, something went wrong.

Uh oh!

There was an error while loading. Please reload this page.

All reactions

@aakash-saravanan-db aakash-saravanan-db Nov 11, 2026

Copy link
Copy Markdown
Collaborator Author

There was a problem hiding this comment.

Choose a reason for hiding this comment

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

The connector release requires kernel #432 and a package minimum set to the version containing it, as stated in the PR description. That version update is pending publication of the companion binding; v1.4.0 is still the latest released kernel. The direct getter is verified against the pinned source build in the native tests and live matrix. Keeping this release prerequisite visible.

Sorry, something went wrong.

Uh oh!

There was an error while loading. Please reload this page.

All reactions

@peco-review-bot peco-review-bot Bot Nov 11, 2026

Copy link
Copy Markdown

There was a problem hiding this comment.

Choose a reason for hiding this comment

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

Understood. The kernel floor bump has to wait until the binding that includes #432 is published, and the native/live tests build against that source. Until then, though, line 970 is still a direct handle.num_modified_rows read, and pyproject still allows ^1.4.0. If this merges before the floor bump, every kernel async get_async_execution_result() raises AttributeError and leaks the result set. A one-line getattr(handle, "num_modified_rows", None), like the sync path uses, would make this safe whenever the two land. I'm leaving this open as a merge gate until the floor is raised in this PR or the guarded read goes in.

Sorry, something went wrong.

Uh oh!

There was an error while loading. Please reload this page.

All reactions
Comment thread

@peco-review-bot peco-review-bot Bot Nov 11, 2026

Copy link
Copy Markdown

There was a problem hiding this comment.

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.

Sorry, something went wrong.

Uh oh!

There was an error while loading. Please reload this page.

All reactions

@aakash-saravanan-db aakash-saravanan-db Nov 11, 2026

Copy link
Copy Markdown
Collaborator Author

There was a problem hiding this comment.

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.

Sorry, something went wrong.

Uh oh!

There was an error while loading. Please reload this page.

All reactions

@peco-review-bot peco-review-bot Bot Nov 11, 2026

Copy link
Copy Markdown

There was a problem hiding this comment.

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.

Sorry, something went wrong.

Uh oh!

There was an error while loading. Please reload this page.

All reactions
Comment thread

@peco-review-bot peco-review-bot Bot Nov 11, 2026

Copy link
Copy Markdown

There was a problem hiding this comment.

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.

Sorry, something went wrong.

Uh oh!

There was an error while loading. Please reload this page.

All reactions
return result_set
except Exception as exc:
raise _wrap_kernel_exception("get_execution_result", exc) from exc

Expand Down
30 changes: 30 additions & 0 deletions 30 src/databricks/sql/backend/thrift_backend.py
Show comments
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import re
import time
import threading
import weakref
from typing import Dict, List, Optional, Union, Any, TYPE_CHECKING
from uuid import UUID

Expand Down Expand Up @@ -267,6 +268,9 @@ def __init__(
raise

self._request_lock = threading.RLock()
self._terminal_operation_statuses: weakref.WeakKeyDictionary[
CommandId, ttypes.TGetOperationStatusResp
] = weakref.WeakKeyDictionary()
self._session_id_hex = None

@property
Expand Down Expand Up @@ -927,10 +931,32 @@ def _results_message_to_execute_response(
def get_execution_result(
self, command_id: CommandId, cursor: Cursor
) -> "ResultSet":
"""Fetch a completed async result and its latest operation metadata.

FetchResults does not carry ``numModifiedRows``, so reuse a terminal
status observed for the same ``CommandId`` instance by
``get_query_state``, or refresh the status when the caller resumes
directly by command ID.
Failed or closed operations raise before result rows are fetched.
"""
thrift_handle = command_id.to_thrift_handle()
if not thrift_handle:
raise ValueError("Not a valid Thrift command ID")

with self._request_lock:
operation_status = self._terminal_operation_statuses.pop(command_id, None)
Comment thread
aakash-saravanan-db marked this conversation as resolved.
Resolved

Uh oh!

There was an error while loading. Please reload this page.

if operation_status is None:
Comment thread
aakash-saravanan-db marked this conversation as resolved.
Resolved

Uh oh!

There was an error while loading. Please reload this page.

Comment thread
aakash-saravanan-db marked this conversation as resolved.
Resolved

Uh oh!

There was an error while loading. Please reload this page.

Comment thread
aakash-saravanan-db marked this conversation as resolved.
Resolved

Uh oh!

There was an error while loading. Please reload this page.

Comment thread

@peco-review-bot peco-review-bot Bot Nov 11, 2026

Copy link
Copy Markdown

There was a problem hiding this comment.

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.

Sorry, something went wrong.

Uh oh!

There was an error while loading. Please reload this page.

All reactions
operation_status = self._poll_for_status(thrift_handle)
self._check_command_not_in_error_or_closed_state(
Comment thread
aakash-saravanan-db marked this conversation as resolved.
Resolved

Uh oh!

There was an error while loading. Please reload this page.

thrift_handle, operation_status
)

num_modified_rows = (
operation_status.numModifiedRows
if operation_status.operationState == ttypes.TOperationState.FINISHED_STATE
else None
)

req = ttypes.TFetchResultsReq(
operationHandle=ttypes.TOperationHandle(
thrift_handle.operationId,
Expand Down Expand Up @@ -981,6 +1007,7 @@ def get_execution_result(
is_staging_operation=is_staging_operation,
arrow_schema_bytes=schema_bytes,
result_format=t_result_set_metadata_resp.resultFormat,
num_modified_rows=num_modified_rows,
)

return ThriftResultSet(
Expand Down Expand Up @@ -1029,6 +1056,9 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
state = CommandState.from_thrift_state(operation_state)
if state is None:
raise ValueError(f"Unknown command state: {operation_state}")
if state == CommandState.SUCCEEDED:
with self._request_lock:
self._terminal_operation_statuses[command_id] = poll_resp
return state

@staticmethod
Expand Down
7 changes: 7 additions & 0 deletions 7 src/databricks/sql/client.py
Show comments
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1714,6 +1714,13 @@ def get_async_execution_result(self):
self.active_command_id, self
)

num_modified_rows = (
self.active_result_set.num_modified_rows
if self.active_result_set
else None
)
self.rowcount = num_modified_rows if num_modified_rows is not None else -1

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
Expand Down
23 changes: 23 additions & 0 deletions 23 tests/unit/test_client.py
Show comments
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -303,6 +303,29 @@ def test_rowcount_resets_between_statements(self):
cursor.execute("SELECT 1")
self.assertEqual(cursor.rowcount, -1)

def test_async_rowcount_on_completion_and_reset(self):
for count in (5, 0, None):
with self.subTest(count=count):
result_set = Mock(num_modified_rows=count, is_staging_operation=False)
backend = ThriftDatabricksClientMockFactory.new()
backend.get_execution_result.return_value = result_set
backend.get_query_state.return_value = CommandState.SUCCEEDED
cursor = client.Cursor(connection=Mock(), backend=backend)
cursor.rowcount = 9

cursor.execute_async("UPDATE t SET x = 1")
self.assertEqual(cursor.rowcount, -1)
cursor.active_command_id = Mock(spec=CommandId)
self.assertIs(cursor.get_async_execution_result(), cursor)
self.assertEqual(cursor.rowcount, count if count is not None else -1)
result_set.fetchall.assert_not_called()

result_set.num_modified_rows = None
cursor.execute_async("SELECT 1")
self.assertEqual(cursor.rowcount, -1)
cursor.get_async_execution_result()
self.assertEqual(cursor.rowcount, -1)

def test_closed_cursor_doesnt_allow_operations(self):
cursor = client.Cursor(Mock(), Mock())
cursor.close()
Expand Down
17 changes: 15 additions & 2 deletions 17 tests/unit/test_kernel_client.py
Show comments
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1233,15 +1233,22 @@ def test_get_query_state_uses_reader_after_result_stream_started():
reader.status.assert_called_once_with()


def test_get_execution_result_uses_retained_owning_handle_first():
@pytest.mark.parametrize("count", [7, 0, None])
def test_get_execution_result_uses_retained_owning_handle_first(count):
"""The first in-process result fetch uses the retained submitting
handle so the kernel finalizes the original async statement telemetry."""
c = _make_client()
c._kernel_session = MagicMock()
fake_stream = MagicMock()
fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
handle = MagicMock()
handle.await_result.return_value = fake_stream
del handle.num_modified_rows

def await_result():
handle.num_modified_rows = count
return fake_stream

handle.await_result.side_effect = await_result
cursor = MagicMock()
cursor.arraysize = 100
cursor.buffer_size_bytes = 1024
Expand All @@ -1252,6 +1259,8 @@ def test_get_execution_result_uses_retained_owning_handle_first():
rs = c.get_execution_result(cid, cursor=cursor)

assert rs is not None
assert rs.num_modified_rows == count
fake_stream.fetch_next_batch.assert_not_called()
c._kernel_session.attach_async_statement.assert_not_called()
handle.await_result.assert_called_once_with()
assert cid.guid in c._async_result_stream_started
Expand All @@ -1264,6 +1273,7 @@ def test_get_execution_result_attaches_by_id_when_no_retained_handle():
fake_stream = MagicMock()
fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
handle = _attach_returns(c, await_result=fake_stream)
handle.num_modified_rows = 7
cursor = MagicMock()
cursor.arraysize = 100
cursor.buffer_size_bytes = 1024
Expand All @@ -1272,6 +1282,7 @@ def test_get_execution_result_attaches_by_id_when_no_retained_handle():
rs = c.get_execution_result(cid, cursor=cursor)

assert rs is not None
assert rs.num_modified_rows == 7
c._kernel_session.attach_async_statement.assert_called_once_with("async-1")
handle.await_result.assert_called_once_with()

Expand Down Expand Up @@ -1326,6 +1337,7 @@ def test_get_execution_result_construction_failure_retains_marker_and_uses_reade
retry_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
reader = MagicMock()
reader.await_result.return_value = retry_stream
reader.num_modified_rows = 0
owning_handle.reader.return_value = reader
cursor = MagicMock()
cursor.arraysize = 100
Expand All @@ -1334,6 +1346,7 @@ def test_get_execution_result_construction_failure_retains_marker_and_uses_reade
rs = c.get_execution_result(cid, cursor=cursor)

assert rs is not None
assert rs.num_modified_rows == 0
owning_handle.reader.assert_called_once_with()
c._kernel_session.attach_async_statement.assert_not_called()
reader.await_result.assert_called_once_with()
Expand Down
60 changes: 59 additions & 1 deletion 60 tests/unit/test_thrift_backend.py
Show comments
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import itertools
import unittest
import pytest
from unittest.mock import patch, MagicMock, Mock
from unittest.mock import patch, MagicMock, Mock, sentinel
from ssl import CERT_NONE, CERT_REQUIRED
from urllib3 import HTTPSConnectionPool

Expand Down Expand Up @@ -213,6 +213,64 @@ def _make_fake_thrift_backend(self):
thrift_backend._create_arrow_table.return_value = (MagicMock(), Mock())
return thrift_backend

@patch("databricks.sql.backend.thrift_backend.ThriftResultSet")
@patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
def test_async_result_rowcount(self, tcli_service_class, result_set):
for state, count, expected_count, polled in [
(ttypes.TOperationState.FINISHED_STATE, 7, 7, True),
(ttypes.TOperationState.FINISHED_STATE, 0, 0, True),
(ttypes.TOperationState.FINISHED_STATE, None, None, True),
(ttypes.TOperationState.FINISHED_STATE, 7, 7, False),
(ttypes.TOperationState.RUNNING_STATE, 7, None, False),
]:
with self.subTest(state=state, count=count, polled=polled):
backend = self._make_fake_thrift_backend()
backend._poll_for_status = Mock(
return_value=ttypes.TGetOperationStatusResp(
status=self.okay_status,
operationState=state,
numModifiedRows=count,
)
)
backend.make_request = Mock(
return_value=Mock(
resultSetMetadata=self.metadata_resp,
results=sentinel.counter_rows,
hasMoreRows=False,
)
)
cid = CommandId.from_thrift_handle(self.operation_handle)
if polled:
self.assertEqual(
backend.get_query_state(cid), CommandState.SUCCEEDED
)
backend.get_execution_result(cid, Mock())

backend._poll_for_status.assert_called_once_with(cid.to_thrift_handle())
response = result_set.call_args.kwargs["execute_response"]
self.assertEqual(response.num_modified_rows, expected_count)
self.assertIs(
result_set.call_args.kwargs["t_row_set"], sentinel.counter_rows
)

@patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
def test_async_result_checks_status_before_fetch(self, tcli_service_class):
backend = self._make_fake_thrift_backend()
backend._poll_for_status = Mock(
return_value=ttypes.TGetOperationStatusResp(
status=self.okay_status,
operationState=ttypes.TOperationState.ERROR_STATE,
errorMessage="async operation failed",
)
)
backend.make_request = Mock()
cid = CommandId.from_thrift_handle(self.operation_handle)

with self.assertRaises(ServerOperationError):
backend.get_execution_result(cid, Mock())

backend.make_request.assert_not_called()

def test_hive_schema_to_arrow_schema_preserves_column_names(self):
columns = [
ttypes.TColumnDesc(
Expand Down
Loading
Oops, something went wrong.

Uh oh!

There was an error while loading. Please reload this page.

Add this suggestion to a batch that can be applied as a single commit. This suggestion is invalid because no changes were made to the code. Suggestions cannot be applied while the pull request is closed. Suggestions cannot be applied while viewing a subset of changes. Only one suggestion per line can be applied in a batch. Add this suggestion to a batch that can be applied as a single commit. Applying suggestions on deleted lines is not supported. You must change the existing code in this line in order to create a valid suggestion. Outdated suggestions cannot be applied. This suggestion has been applied or marked resolved. Suggestions cannot be applied from pending reviews. Suggestions cannot be applied on multi-line comments. Suggestions cannot be applied while the pull request is queued to merge. Suggestion cannot be applied right now. Please check back later.
You are visitor 01402503 — GeezerVison does not store the pages you look up.
Best viewed with Netscape Navigator 4.7 or Internet Explorer 6 at 800×600. No cookies required.