Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion src/databricks/sql/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -1759,9 +1759,18 @@ def cancel(self) -> None:
def close(self) -> None:
"""Close cursor"""
self.open = False
self.active_command_id = None
if self.active_result_set:
self._close_and_clear_active_result_set()
elif self.active_command_id is not None:
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Outdated
# Async submission whose result was never fetched (no
# get_execution_result call), so the result-set close path never
# fired. Issue an explicit close_command to free the server-side
# statement handle instead of leaking it until session close.
try:
self.backend.close_command(self.active_command_id)
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
except Exception as exc:
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
logger.warning("close_command on cursor close failed: %s", exc)
self.active_command_id = None

@property
def query_id(self) -> Optional[str]:
Expand Down
32 changes: 32 additions & 0 deletions tests/e2e/test_driver.py
Original file line number Diff line number Diff line change
Expand Up @@ -342,6 +342,38 @@ def test_execute_async__large_result(self, extra_params):

assert len(result) == x_dimension * y_dimension

@pytest.mark.parametrize(
"extra_params",
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
[
{},
],
)
def test_execute_async__close_without_fetch_frees_handle(self, extra_params):
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
"""Closing a cursor whose async command result was never fetched must free
the server-side statement handle (issue #791). Otherwise the handle leaks
until the session closes."""
with self.cursor(extra_params) as cursor:
cursor.execute_async("SELECT 1")

# Capture the server-side command id before we close the cursor.
command_id = cursor.active_command_id
assert command_id is not None

backend = cursor.backend

# Sanity: the handle is live and pollable before close.
backend.get_query_state(command_id)

# User decides not to wait for the result and closes the cursor
# without ever calling get_async_execution_result().
cursor.close()

# After close, the server-side handle must have been freed, so a
# re-poll of the saved command id should raise a server error.
# Pre-fix (leak): this poll succeeds. Post-fix: it raises.
with pytest.raises((RequestError, OperationalError, DatabaseError)):
backend.get_query_state(command_id)

Comment thread
peco-review-bot[bot] marked this conversation as resolved.

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 — The SEA branch of this test may be brittle. After close_command issues a DELETE on the statement, the subsequent get_query_state does a plain GET (_poll_query) and returns whatever status.state the server reports. The test only tolerates a raised error or state in (CommandState.CLOSED, CommandState.CANCELLED). If SEA responds to a GET on a just-deleted statement with any other terminal/transient state (e.g. it still echoes SUCCEEDED briefly, or a state not in that tuple), the else assert fails and the test flakes — even though the handle was correctly freed. Since get_query_state intentionally does not call _check_command_not_in_failed_or_closed_state, there's no guarantee of a CLOSED/CANCELLED signal. Consider confirming empirically what SEA returns post-delete, or broadening the accepted signal (e.g. also accept a 404-derived error) so the test isn't tied to an unverified state assumption.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

NEEDS HUMAN DECISION — the bots can't resolve this thread; a maintainer's input is required.

The comment targets the SEA branch of an E2E test whose correct accepted-state set depends on what a live SEA warehouse actually returns from get_query_state() after close_command issues the statement DELETE. This job has no live-warehouse connection and must not run/add e2e tests, so I cannot verify SEA's real post-delete signal here. The reviewer's two options both hinge on that empirical fact: (a) confirm empirically — impossible in this job; (b) broaden the accepted states — unsafe to do blind, because the pre-fix leak "returns a live/terminal-success state," so broadening (e.g. accepting SUCCEEDED) without knowing SEA's freed-handle behavior would let the test pass in the leak case and destroy the regression it guards. A human needs to run this against a live SEA warehouse to observe the actual post-DELETE state (or 404), then either narrow the assertion to that confirmed signal or, if SEA gives no reliable freed signal, restructure the SEA branch (e.g. skip the re-poll assertion for SEA). Flagging for human judgment rather than guessing at a broadening that could mask the leak.


# Exclude Retry tests because they require specific setups, and LargeQueries too slow for core
# tests
Expand Down
Loading