fix: close unfetched async commands#794
Conversation
Signed-off-by: Nanook <nanookclaw@users.noreply.github.com>
vikrantpuppala
left a comment
There was a problem hiding this comment.
Thanks for the contribution, would be great if we can address the comments and get this in
| elif self.active_command_id is not None: | ||
| self.backend.close_command(self.active_command_id) |
There was a problem hiding this comment.
- 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.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
Summary
Cursor.close()now frees an active async command id when no result set has been fetched yet. This covers theexecute_async()followed byclose()path from #791, whereactive_command_idis set butactive_result_setis stillNone.When a result set exists, close still delegates to
active_result_set.close()so that path can perform its existing backend cleanup without a duplicateclose_commandcall. The cursor clearsactive_command_idin afinallyblock 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.pygit diff --checkCloses #791