Skip to content

fix: close unfetched async commands #794

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
nanookclaw wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: close unfetched async commands #794

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
10 changes: 7 additions & 3 deletions src/databricks/sql/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1713,9 +1713,13 @@ 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()
try:
if self.active_result_set:
self._close_and_clear_active_result_set()
elif self.active_command_id is not None:
self.backend.close_command(self.active_command_id)
Comment on lines +1719 to +1720

@vikrantpuppala vikrantpuppala Jul 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

  • ResultSet.close() (result_set.py:184-187) wraps the same call in try/except RequestError, swallowing CursorAlreadyClosedError. The new branch has neither the guard nor the except — only a finally that clears active_command_id but does not contain the raise.
  • Connection._close() (client.py:547-550) runs for cursor in self._cursors: cursor.close() with the try/except covering only the subsequent session.close(). A raise from any cursor aborts the loop → remaining cursors never closed, session.close() never runs.
  • Blast radius: conn.close() / with teardown when an async cursor's handle is already gone server-side (post-cancel, expired session) or a transient network blip → whole-session leak + raw stack trace out of close(). This is the exact scenario the PR targets, so it's a live regression, not hypothetical.
  • Fix: mirror ResultSet.close() — swallow RequestError/CursorAlreadyClosedError (+ a catch-all) and log. This is what issue Cursor.close() leaks server-side handle for async commands that were never fetched #791's suggested fix specified and the PR dropped.

@vikrantpuppala vikrantpuppala Jul 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

ResultSet.close() gates the RPC on status != CLOSED and not has_been_closed_server_side and connection.open (result_set.py:180-182). The new branch fires unconditionally, sending a pointless (and, per F1, unguarded) close_command RPC even when the session is already closed server-side. Gating on self.connection.open would avoid the round-trip and much of F1's trigger surface.

finally:
self.active_command_id = None

@property
def query_id(self) -> Optional[str]:
Expand Down
27 changes: 27 additions & 0 deletions tests/unit/test_client.py
View file
Open in desktop

@vikrantpuppala vikrantpuppala Jul 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Both added tests use bare Mock() backends whose close_command never raises, so they only prove the happy case. The one genuinely new invariant — active_command_id still cleared (via finally) and close() staying best-effort when close_command raises — is entirely unverified. Fix: add a test with mock_backend.close_command.side_effect = RequestError(...), asserting close() does not raise and active_command_id is None.

Original file line number Diff line number Diff line change
Expand Up @@ -280,6 +280,33 @@ def test_context_manager_closes_cursor(self):
cursor.close = mock_close
mock_close.assert_called_once_with()

def test_close_closes_active_command_without_result_set(self):
mock_backend = Mock()
mock_command_id = Mock(spec=CommandId)
cursor = client.Cursor(Mock(), mock_backend)
cursor.active_command_id = mock_command_id

cursor.close()

mock_backend.close_command.assert_called_once_with(mock_command_id)
self.assertIsNone(cursor.active_command_id)
self.assertIsNone(cursor.active_result_set)

def test_close_delegates_to_active_result_set_without_double_closing_command(self):
mock_backend = Mock()
mock_result_set = Mock()
mock_command_id = Mock(spec=CommandId)
cursor = client.Cursor(Mock(), mock_backend)
cursor.active_result_set = mock_result_set
cursor.active_command_id = mock_command_id

cursor.close()

mock_result_set.close.assert_called_once_with()
mock_backend.close_command.assert_not_called()
self.assertIsNone(cursor.active_command_id)
self.assertIsNone(cursor.active_result_set)

def dict_product(self, dicts):
"""
Generate cartesion product of values in input dictionary, outputting a dictionary
Expand Down