Skip to content

fix: close unfetched async commands#794

Open
nanookclaw wants to merge 1 commit into
databricks:mainfrom
nanookclaw:fix/close-unfetched-async-command
Open

fix: close unfetched async commands#794
nanookclaw wants to merge 1 commit into
databricks:mainfrom
nanookclaw:fix/close-unfetched-async-command

Conversation

@nanookclaw

Copy link
Copy Markdown

Summary

Cursor.close() now frees an active async command id when no result set has been fetched yet. This covers the execute_async() followed by close() path from #791, where active_command_id is set but active_result_set is still None.

When a result set exists, close still delegates to active_result_set.close() so that path can perform its existing backend cleanup without a duplicate close_command call. The cursor clears active_command_id in a finally block either way.

Tests

  • PYTHONPATH=src uv run --with pytest --with thrift --with pandas --with lz4 --with requests --with oauthlib --with openpyxl --with urllib3 --with python-dateutil --with pyjwt --with pybreaker --python 3.12 --env-file test.env.example python -m pytest tests/unit/test_client.py -k 'close_closes_active_command_without_result_set or close_delegates_to_active_result_set_without_double_closing_command or closing_result_set_hard_closes_commands or closing_connection_closes_commands'
  • uv run --with black==22.3.0 --python 3.12 black --check src/databricks/sql/client.py tests/unit/test_client.py
  • git diff --check

Closes #791

Signed-off-by: Nanook <nanookclaw@users.noreply.github.com>

@vikrantpuppala vikrantpuppala left a comment

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.

Thanks for the contribution, would be great if we can address the comments and get this in

Comment on lines +1719 to +1720
elif self.active_command_id is not None:
self.backend.close_command(self.active_command_id)

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.

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.

Comment thread 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.

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.

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.

Cursor.close() leaks server-side handle for async commands that were never fetched

2 participants