From e3658da2d6b0d854a5f6eff95815726edf80a1a7 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" <3815206+peco-engineer-bot[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 05:44:00 +0000 Subject: [PATCH 1/6] 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 | 11 ++++++++++- tests/e2e/test_driver.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index d45d51181..76711ee16 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1759,9 +1759,18 @@ 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 whose result was never fetched (no + # get_execution_result call), so the result-set close path never + # fired. Issue an explicit close_command to free the server-side + # statement handle instead of leaking it until session close. + 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..fd4e01935 100644 --- a/tests/e2e/test_driver.py +++ b/tests/e2e/test_driver.py @@ -342,6 +342,38 @@ def test_execute_async__large_result(self, extra_params): assert len(result) == x_dimension * y_dimension + @pytest.mark.parametrize( + "extra_params", + [ + {}, + ], + ) + def test_execute_async__close_without_fetch_frees_handle(self, extra_params): + """Closing a cursor whose async command result was never fetched must free + the server-side statement handle (issue #791). Otherwise the handle leaks + until the session closes.""" + with self.cursor(extra_params) as cursor: + cursor.execute_async("SELECT 1") + + # Capture the server-side command id before we close the cursor. + command_id = cursor.active_command_id + assert command_id is not None + + backend = cursor.backend + + # Sanity: the handle is live and pollable before close. + backend.get_query_state(command_id) + + # User decides not to wait for the result and closes the cursor + # without ever calling get_async_execution_result(). + cursor.close() + + # After close, the server-side handle must have been freed, so a + # re-poll of the saved command id should raise a server error. + # Pre-fix (leak): this poll succeeds. Post-fix: it raises. + with pytest.raises((RequestError, OperationalError, DatabaseError)): + backend.get_query_state(command_id) + # Exclude Retry tests because they require specific setups, and LargeQueries too slow for core # tests From c7901ab01f0bd5f00f75fcc0a6b3f8ec9439e603 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Thu, 23 Jul 2026 05:55:11 +0000 Subject: [PATCH 2/6] ai: apply changes for #873 (2 review threads) Addresses: - #3635836346 at src/databricks/sql/client.py:1764 - #3635836348 at src/databricks/sql/client.py:1770 Signed-off-by: peco-engineer-bot[bot] --- src/databricks/sql/client.py | 4 ++- tests/unit/test_client.py | 61 ++++++++++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 76711ee16..a8a409253 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1761,11 +1761,13 @@ def close(self) -> None: self.open = False if self.active_result_set: self._close_and_clear_active_result_set() - elif self.active_command_id is not None: + elif self.active_command_id is not None and self.connection.open: # Async submission whose result was never fetched (no # get_execution_result call), so the result-set close path never # fired. Issue an explicit close_command to free the server-side # statement handle instead of leaking it until session close. + # Gate on connection.open (mirroring ResultSet.close) so we don't + # attempt a network call on an already-torn-down session. try: self.backend.close_command(self.active_command_id) except Exception as exc: diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 66a3722ec..2c3090e3d 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -161,6 +161,67 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class): # Should NOT have called backend.close_command (already closed) mock_backend.close_command.assert_not_called() + def test_cursor_close_frees_command_when_no_result_set(self): + """Closing a cursor with an unfetched async command frees the server handle. + + When active_result_set is None but active_command_id is set (an async + submission whose result was never fetched), close() must issue an + explicit backend.close_command so the server-side statement handle is + not leaked. This branch is backend-agnostic (Thrift/SEA/kernel). + """ + mock_backend = Mock(spec=ThriftDatabricksClient) + mock_connection = Mock() + cursor = client.Cursor(connection=mock_connection, backend=mock_backend) + + command_id = Mock(spec=CommandId) + cursor.active_command_id = command_id + cursor.active_result_set = None + + cursor.close() + + mock_backend.close_command.assert_called_once_with(command_id) + self.assertIsNone(cursor.active_command_id) + + def test_cursor_close_skips_command_when_connection_closed(self): + """Closing a cursor after its session is gone must not call close_command. + + Mirrors ResultSet.close, which gates backend.close_command on + connection.open, to avoid a network call on a dead session (and a + spurious warning) during shutdown ordering. + """ + mock_backend = Mock(spec=ThriftDatabricksClient) + mock_connection = Mock() + mock_connection.open = False + cursor = client.Cursor(connection=mock_connection, backend=mock_backend) + + cursor.active_command_id = Mock(spec=CommandId) + cursor.active_result_set = None + + cursor.close() + + mock_backend.close_command.assert_not_called() + self.assertIsNone(cursor.active_command_id) + + def test_cursor_close_does_not_double_close_when_result_set_present(self): + """Closing a cursor with an active result set must NOT call close_command. + + The result-set close path is responsible for freeing the handle in that + case, so the elif branch must not fire (no double-close). + """ + mock_backend = Mock(spec=ThriftDatabricksClient) + mock_connection = Mock() + cursor = client.Cursor(connection=mock_connection, backend=mock_backend) + + cursor.active_command_id = Mock(spec=CommandId) + result_set = Mock(spec=ResultSet) + cursor.active_result_set = result_set + + cursor.close() + + result_set.close.assert_called_once() + mock_backend.close_command.assert_not_called() + self.assertIsNone(cursor.active_command_id) + @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME) def test_cant_open_cursor_on_closed_connection(self, mock_client_class): connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS) From a48060915c8db8192f997bc5cba941622e199a66 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Thu, 23 Jul 2026 06:02:34 +0000 Subject: [PATCH 3/6] ai: apply changes for #873 (1 review thread) Addresses: - #3635883523 at tests/e2e/test_driver.py:351 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 fd4e01935..4b9bb3671 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_execute_async__close_without_fetch_frees_handle(self, extra_params): From ee59c9c007e233ecc4cadf0089f4f602c77a04ec Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Thu, 23 Jul 2026 06:09:52 +0000 Subject: [PATCH 4/6] ai: apply changes for #873 (1 review thread) Addresses: - #3635909165 at src/databricks/sql/client.py:1773 Signed-off-by: peco-engineer-bot[bot] --- src/databricks/sql/client.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index a8a409253..7bf0803a9 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -16,6 +16,7 @@ from databricks.sql import * from databricks.sql.exc import ( OperationalError, + RequestError, SessionAlreadyClosedError, CursorAlreadyClosedError, InterfaceError, @@ -1770,6 +1771,14 @@ def close(self) -> None: # attempt a network call on an already-torn-down session. try: self.backend.close_command(self.active_command_id) + except RequestError as e: + if isinstance(e.args[1], CursorAlreadyClosedError): + # Already-closed handle (e.g. a prior cancel() or concurrent + # session teardown) is an expected, benign case — mirror + # ResultSet.close and log at info, not warning. + logger.info("Operation was canceled by a prior request") + else: + logger.warning("close_command on cursor close failed: %s", e) except Exception as exc: logger.warning("close_command on cursor close failed: %s", exc) self.active_command_id = None From 244563d0c6ea30744018632a3acfbf1e34ba8cfd Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Thu, 23 Jul 2026 06:20:38 +0000 Subject: [PATCH 5/6] ai: apply changes for #873 (2 review threads) Addresses: - #3635951484 at tests/e2e/test_driver.py:379 - #3635951489 at src/databricks/sql/client.py:1775 Signed-off-by: peco-engineer-bot[bot] --- src/databricks/sql/client.py | 2 +- tests/e2e/test_driver.py | 19 ++++++++++++++----- tests/unit/test_client.py | 26 ++++++++++++++++++++++++++ 3 files changed, 41 insertions(+), 6 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 7bf0803a9..e50b6e556 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1772,7 +1772,7 @@ def close(self) -> None: try: self.backend.close_command(self.active_command_id) except RequestError as e: - if isinstance(e.args[1], CursorAlreadyClosedError): + if len(e.args) > 1 and isinstance(e.args[1], CursorAlreadyClosedError): # Already-closed handle (e.g. a prior cancel() or concurrent # session teardown) is an expected, benign case — mirror # ResultSet.close and log at info, not warning. diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py index 4b9bb3671..5515cff7d 100644 --- a/tests/e2e/test_driver.py +++ b/tests/e2e/test_driver.py @@ -371,11 +371,20 @@ def test_execute_async__close_without_fetch_frees_handle(self, extra_params): # without ever calling get_async_execution_result(). cursor.close() - # After close, the server-side handle must have been freed, so a - # re-poll of the saved command id should raise a server error. - # Pre-fix (leak): this poll succeeds. Post-fix: it raises. - with pytest.raises((RequestError, OperationalError, DatabaseError)): - backend.get_query_state(command_id) + # After close, the server-side handle must have been freed. How a + # re-poll of the saved command id surfaces that is backend-specific: + # - Thrift raises a server error on the closed handle. + # - SEA's get_query_state() does a plain GET and returns the + # status.state, so a freed statement may come back as a terminal + # CLOSED/CANCELLED state instead of raising. + # Accept either signal as proof the handle was freed. Pre-fix (leak), + # the poll instead returns a live/terminal-success state. + try: + state = backend.get_query_state(command_id) + except (RequestError, OperationalError, DatabaseError): + pass + else: + assert state in (CommandState.CLOSED, CommandState.CANCELLED) # Exclude Retry tests because they require specific setups, and LargeQueries too slow for core diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 2c3090e3d..c21ddbc13 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -222,6 +222,32 @@ def test_cursor_close_does_not_double_close_when_result_set_present(self): mock_backend.close_command.assert_not_called() self.assertIsNone(cursor.active_command_id) + def test_cursor_close_tolerates_single_arg_request_error(self): + """A RequestError with only one positional arg must not escape close(). + + Some RequestErrors (e.g. from unified_http_client during shutdown) are + built with a single message arg, so e.args[1] would IndexError. close() + must guard the length check and swallow the error like any other + close_command failure. + """ + from databricks.sql.exc import RequestError + + mock_backend = Mock(spec=ThriftDatabricksClient) + mock_backend.close_command.side_effect = RequestError( + "HTTP client is closing or has been closed" + ) + mock_connection = Mock() + cursor = client.Cursor(connection=mock_connection, backend=mock_backend) + + cursor.active_command_id = Mock(spec=CommandId) + cursor.active_result_set = None + + # Should not raise despite the single-arg RequestError. + cursor.close() + + mock_backend.close_command.assert_called_once() + self.assertIsNone(cursor.active_command_id) + @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME) def test_cant_open_cursor_on_closed_connection(self, mock_client_class): connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS) From dce92689dfe3273c88c5ad0d0e94b6b569b0d8f9 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Thu, 23 Jul 2026 06:30:03 +0000 Subject: [PATCH 6/6] ai: apply changes for #873 (1 review thread) Addresses: - #3635998296 at src/databricks/sql/client.py:1782 Signed-off-by: peco-engineer-bot[bot] --- src/databricks/sql/client.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index e50b6e556..a775c4e45 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1780,7 +1780,16 @@ def close(self) -> None: else: logger.warning("close_command on cursor close failed: %s", e) except Exception as exc: - logger.warning("close_command on cursor close failed: %s", exc) + # Unlike the RequestError branch above (an expected, often + # transient network failure), reaching here means something we + # did not anticipate — surface the full traceback so it is + # diagnosable rather than indistinguishable from a benign + # network blip. + logger.warning( + "close_command on cursor close failed unexpectedly: %s", + exc, + exc_info=True, + ) self.active_command_id = None @property