From 58560644fde78f62f28ce3b30723666eb1672bf0 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" <3815206+peco-engineer-bot[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 23:18:14 +0000 Subject: [PATCH 1/2] Cursor.close() leaks server-side handle for async commands that were never fetched (#791) Signed-off-by: peco-engineer-bot[bot] <3815206+peco-engineer-bot[bot]@users.noreply.github.com> --- src/databricks/sql/client.py | 13 ++++++++- tests/e2e/test_driver.py | 53 ++++++++++++++++++++++++++++++++++++ tests/unit/test_client.py | 34 +++++++++++++++++++++++ 3 files changed, 99 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 00a19154a..75871d387 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1757,9 +1757,20 @@ 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: + # Async submission that never had get_async_execution_result called: + # the result-set close path didn't fire, so the server-side handle + # would leak until the session closes. Issue an explicit + # close_command to free it. Safe across all backends (Thrift sends + # TCloseOperation, SEA sends a close/cancel, kernel no-ops on + # unknown ids). + try: + self.backend.close_command(self.active_command_id) + except Exception as exc: + logger.warning("close_command on cursor close failed: %s", exc) + self.active_command_id = None @property def query_id(self) -> Optional[str]: diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py index 5fe3db037..3158dcec5 100644 --- a/tests/e2e/test_driver.py +++ b/tests/e2e/test_driver.py @@ -342,6 +342,59 @@ def test_execute_async__large_result(self, extra_params): assert len(result) == x_dimension * y_dimension + @pytest.mark.parametrize( + "extra_params", + [ + {}, + ], + ) + def test_close_frees_async_handle_without_fetch(self, extra_params): + """Regression test for #791. + + Cursor.close() must free the server-side statement handle for an + async-submitted command even when get_async_execution_result() was + never called. Before the fix, close() only nulled active_command_id + and (for the never-fetched case) never issued close_command, so the + handle leaked until the session closed. + """ + long_running_query = "SELECT COUNT(*) FROM RANGE(10000 * 16) x JOIN RANGE(10000) y ON FROM_UNIXTIME(x.id * y.id, 'yyyy-MM-dd') LIKE '%not%a%date%'" + with self.cursor(extra_params) as cursor: + cursor.execute_async(long_running_query) + + command_id = cursor.active_command_id + backend = cursor.backend + assert command_id is not None + + # The handle is live server-side right after async submission. + state = backend.get_query_state(command_id) + assert state in ( + CommandState.PENDING, + CommandState.RUNNING, + CommandState.SUCCEEDED, + ) + + # Close without ever fetching the result. + cursor.close() + + # The server must no longer report the handle as live. Polling the + # same command id should either report a terminal CLOSED/CANCELLED + # state, or raise signalling the command was closed / no longer + # exists server side (DatabaseError "closed server side" or a + # RESOURCE_DOES_NOT_EXIST RequestError). Any of these proves the + # server-side handle was freed by close(). + try: + post_close_state = backend.get_query_state(command_id) + assert post_close_state in ( + CommandState.CLOSED, + CommandState.CANCELLED, + ), f"Handle still live after close(): {post_close_state}" + except (DatabaseError, RequestError) as e: + assert ( + "closed server side" in str(e) + or "does not exist" in str(e) + or "RESOURCE_DOES_NOT_EXIST" in str(e) + ) + # Exclude Retry tests because they require specific setups, and LargeQueries too slow for core # tests diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 66a3722ec..36189e251 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -330,6 +330,40 @@ def test_context_manager_closes_cursor(self): cursor.close = mock_close mock_close.assert_called_once_with() + def test_close_frees_async_command_without_result_set(self): + """Regression test for #791. + + When a command was submitted asynchronously and never fetched, + active_result_set is None but active_command_id is set. close() must + still issue backend.close_command() to free the server-side handle. + """ + mock_backend = Mock() + cursor = client.Cursor(Mock(), mock_backend) + mock_command_id = Mock() + cursor.active_command_id = mock_command_id + cursor.active_result_set = None + + cursor.close() + + mock_backend.close_command.assert_called_once_with(mock_command_id) + self.assertIsNone(cursor.active_command_id) + self.assertFalse(cursor.open) + + def test_close_swallows_close_command_errors(self): + """close() must not raise if backend.close_command fails for an + unfetched async command (best-effort cleanup).""" + mock_backend = Mock() + mock_backend.close_command.side_effect = Exception("boom") + cursor = client.Cursor(Mock(), mock_backend) + cursor.active_command_id = Mock() + cursor.active_result_set = None + + # Should not raise. + cursor.close() + + self.assertIsNone(cursor.active_command_id) + self.assertFalse(cursor.open) + def dict_product(self, dicts): """ Generate cartesion product of values in input dictionary, outputting a dictionary From b63ea428242fd3ba601ab53e1c917f63696c634d Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 17 Jul 2026 23:26:00 +0000 Subject: [PATCH 2/2] ai: apply changes for #873 (1 review thread) Addresses: - #3606905116 at tests/e2e/test_driver.py:346 Signed-off-by: peco-engineer-bot[bot] --- tests/e2e/test_driver.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py index 3158dcec5..d5e170e99 100644 --- a/tests/e2e/test_driver.py +++ b/tests/e2e/test_driver.py @@ -346,6 +346,9 @@ def test_execute_async__large_result(self, extra_params): "extra_params", [ {}, + { + "use_sea": True, + }, ], ) def test_close_frees_async_handle_without_fetch(self, extra_params):