FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

fix: close unfetched async commands by nanookclaw · Pull Request #794 · databricks/databricks-sql-python · GitHub

Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension .py  (2) All 1 file type selected
Viewed files
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Unified
Split
Hide whitespace
Diff view
Unified
Split
Hide whitespace
10 changes: 7 additions & 3 deletions src/databricks/sql/client.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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

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.

Choose a reason Spam Abuse Off Topic Outdated Duplicate Resolved Low Quality
  • 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.

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.

Choose a reason Spam Abuse Off Topic Outdated Duplicate Resolved Low Quality

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

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.

Choose a reason Spam Abuse Off Topic Outdated Duplicate Resolved Low Quality

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.

This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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

Back | FazBrowse Home | New Git URL