From bec22d81ab5a25f6348dd91c415631f52a0186df Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 18:55:26 +0000 Subject: [PATCH] fix: stop a connection from keeping every cursor it created Connection._cursors held every cursor ever created, so a long-lived connection grew by one closed cursor per query and close() walked all of them. Cursor.close() now removes its entry; Connection.close() still closes every cursor that is open. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- CHANGELOG.md | 1 + src/databricks/sql/client.py | 9 ++++++--- tests/unit/test_client.py | 33 +++++++++++++++++++++++++++++++++ 3 files changed, 40 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2bae4c54e..aa5b4ad3d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ - Fix the kernel backend ignoring `use_cloud_fetch=False`, so complete results can be fetched without cloud-storage downloads. - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) - Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most. +- Fix a connection keeping every cursor it ever created, so memory no longer grows with each query on a long-lived connection and `close()` no longer slows down with the number of queries run. # 4.9.0 (2026-11-04) - Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes. diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 613b62386..cba1ebd38 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -407,7 +407,9 @@ def read(self) -> Optional[OAuthToken]: self.disable_pandas = kwargs.get("_disable_pandas", False) self.lz4_compression = kwargs.get("enable_query_result_lz4_compression", True) self.use_cloud_fetch = kwargs.get("use_cloud_fetch", True) - self._cursors = [] # type: List[Cursor] + # Open cursors, in creation order. Cursor.close() removes its entry, so a + # long-lived connection does not keep every cursor it ever created. + self._cursors: Dict[Cursor, None] = {} self.telemetry_batch_size = kwargs.get( "telemetry_batch_size", TelemetryClientFactory.DEFAULT_BATCH_SIZE ) @@ -755,7 +757,7 @@ def cursor( result_buffer_size_bytes=buffer_size_bytes, row_limit=row_limit, ) - self._cursors.append(cursor) + self._cursors[cursor] = None return cursor def close(self) -> None: @@ -764,7 +766,7 @@ def close(self) -> None: def _close(self, close_cursors=True) -> None: if close_cursors: - for cursor in self._cursors: + for cursor in list(self._cursors): cursor.close() try: @@ -1978,6 +1980,7 @@ def cancel(self) -> None: def close(self) -> None: """Close cursor""" + self.connection._cursors.pop(self, None) self.open = False self.active_command_id = None if self.active_result_set: diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 047252488..09a55c262 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -6,6 +6,7 @@ import itertools from decimal import Decimal from datetime import datetime, date +from types import SimpleNamespace from uuid import UUID from databricks.sql.thrift_api.TCLIService.ttypes import ( @@ -171,6 +172,38 @@ def test_cant_open_cursor_on_closed_connection(self, mock_client_class): connection.cursor() self.assertIn("closed", str(cm.exception)) + def test_connection_tracks_only_open_cursors(self): + thrift = patch("%s.session.ThriftDatabricksClient" % self.PACKAGE_NAME) + # session.py imports the kernel client lazily; a fake module keeps the + # test independent of the kernel wheel. + kernel = patch.dict( + sys.modules, + { + "%s.backend.kernel.client" + % self.PACKAGE_NAME: SimpleNamespace(KernelDatabricksClient=MagicMock()) + }, + ) + for backend_patch, extra in ((thrift, {}), (kernel, {"use_kernel": True})): + with self.subTest(**extra), backend_patch: + connection = databricks.sql.connect( + **self.DUMMY_CONNECTION_ARGS, **extra + ) + for _ in range(3): + with connection.cursor(): + pass + open_cursors = [connection.cursor(), connection.cursor()] + result_set = Mock() + open_cursors[0].active_result_set = result_set + closed_twice = connection.cursor() + closed_twice.close() + closed_twice.close() + self.assertEqual(list(connection._cursors), open_cursors) + + connection.close() + self.assertFalse(any(cursor.open for cursor in open_cursors)) + result_set.close.assert_called_once() + self.assertEqual(list(connection._cursors), []) + @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME) @patch("%s.client.Cursor" % PACKAGE_NAME) def test_arraysize_buffer_size_passthrough(