From 508709b40e45aa9f42be8e17f53345434bf9a37b Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 18:53:41 +0000 Subject: [PATCH] perf(kernel): skip the Python telemetry client on use_kernel connections The kernel exports telemetry for use_kernel connections, but every kernel connect still registered a no-op client with TelemetryClientFactory under its global lock (and the first one started the flush thread). In a 200-thread cold burst that lock convoys behind the GIL: initialize_telemetry_client took p50 692 ms / max 1.7 s, and connect p50 was ~3.0 s vs ~1.7 s without it. Kernel connections now use a NoopTelemetryClient directly and skip the factory on close. This also stops a kernel connection opened first from leaving a no-op client registered for the host, which silently disabled Python telemetry for later Thrift connections to that host. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 1 + src/databricks/sql/client.py | 14 +++++++- tests/unit/test_session.py | 11 +++---- tests/unit/test_telemetry.py | 64 ++++++++++++++++++++++++++++++++++++ 4 files changed, 82 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2bae4c54e..e4a24ceb8 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. +- Speed up concurrent `use_kernel=True` connects by skipping the Python telemetry client setup, which the kernel does not use; this also stops a kernel connection from disabling Python telemetry for later Thrift connections to the same host. # 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..00b531969 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -88,6 +88,7 @@ TSparkParameter, ) from databricks.sql.telemetry.telemetry_client import ( + NoopTelemetryClient, TelemetryHelper, TelemetryClientFactory, ) @@ -473,6 +474,13 @@ def read(self) -> Optional[OAuthToken]: self.enable_telemetry = kwargs.get("enable_telemetry", True) self.telemetry_enabled = TelemetryHelper.is_telemetry_enabled(self) + if self._kernel_owns_telemetry(): + # The kernel exports this connection's telemetry. Registering with the + # Python factory would only serialize concurrent connects on its lock + # and start an idle flush thread. + self._telemetry_client = NoopTelemetryClient() + return + TelemetryClientFactory.initialize_telemetry_client( telemetry_enabled=self.telemetry_enabled, session_id_hex=self.get_session_id_hex(), @@ -543,6 +551,9 @@ def read(self) -> Optional[OAuthToken]: session_id=self.get_session_id_hex(), ) + def _kernel_owns_telemetry(self) -> bool: + return bool(getattr(self.session, "use_kernel", False)) + def _open_session_with_reyden_fallback( self, server_hostname: str, @@ -772,7 +783,8 @@ def _close(self, close_cursors=True) -> None: except Exception as e: logger.error(f"Attempt to close session raised a local exception: {e}") - TelemetryClientFactory.close(host_url=self.session.host) + if not self._kernel_owns_telemetry(): + TelemetryClientFactory.close(host_url=self.session.host) # Close HTTP client that was created by this connection if self.http_client: diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 2a555a968..8a40acb01 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -16,6 +16,7 @@ OperationalError, ) from databricks.sql.session import Session +from databricks.sql.telemetry.telemetry_client import NoopTelemetryClient import databricks.sql @@ -974,13 +975,9 @@ def test_kernel_http_client_lifecycle(self, mock_thrift, route): ) assert conn.session.http_client is conn.http_client assert kernel.call_args.kwargs["http_client"] is conn.http_client - params = telemetry.return_value.export_initial_telemetry_log.call_args.kwargs[ - "driver_connection_params" - ] - assert params.use_proxy is (not explicit) - assert params.use_system_proxy is (not explicit) - if explicit: - assert params.proxy_host_info is None + # The kernel owns telemetry: no Python telemetry client is used. + telemetry.assert_not_called() + assert isinstance(conn._telemetry_client, NoopTelemetryClient) with conn.cursor() as cursor: cursor.execute("SELECT 1") assert cursor.fetchall() == [(1,)] diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py index f000eaa43..198fc6efd 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -530,6 +530,70 @@ def test_connection_failure_does_not_send_telemetry_for_kernel( mock_export_failure_log.assert_not_called() + @staticmethod + def _kernel_connection(mock_session): + session = mock_session.return_value + session.use_kernel = True + session.host = "test-host" + return sql.connect( + server_hostname="test-host", http_path="/test-path", use_kernel=True + ) + + @staticmethod + def _init_python_client(host_url): + with patch( + "databricks.sql.common.unified_http_client.UnifiedHttpClient._setup_pool_managers" + ): + TelemetryClientFactory.initialize_telemetry_client( + telemetry_enabled=True, + session_id_hex="thrift-session", + auth_provider=AccessTokenAuthProvider("token"), + host_url=host_url, + batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE, + client_context=MagicMock(), + ) + + @patch("databricks.sql.client.Session") + def test_kernel_connection_skips_python_telemetry_pipeline(self, mock_session): + """The kernel owns telemetry for use_kernel connections, so connect and + close must not start the Python factory (lock, flush thread, executor).""" + conn = self._kernel_connection(mock_session) + + assert isinstance(conn._telemetry_client, NoopTelemetryClient) + assert TelemetryClientFactory._clients == {} + assert TelemetryClientFactory._initialized is False + + conn.close() + assert TelemetryClientFactory._clients == {} + assert TelemetryClientFactory._initialized is False + + @patch("databricks.sql.client.Session") + def test_kernel_connection_leaves_shared_host_client_alone(self, mock_session): + """A kernel connection must not take or drop a reference on the Python + client that a Thrift connection to the same host is using.""" + self._init_python_client("test-host") + holder = TelemetryClientFactory._clients["test-host"] + + conn = self._kernel_connection(mock_session) + assert holder.refcount == 1 + conn.close() + + assert holder.refcount == 1 + assert TelemetryClientFactory._clients["test-host"] is holder + + @patch("databricks.sql.client.Session") + def test_thrift_telemetry_after_kernel_connection_is_not_noop(self, mock_session): + """A kernel connection opened first must not leave a no-op client + registered for the host that later Thrift connections would reuse.""" + conn = self._kernel_connection(mock_session) + self._init_python_client("test-host") + + assert isinstance( + TelemetryClientFactory.get_telemetry_client("test-host"), TelemetryClient + ) + conn.close() + TelemetryClientFactory.close(host_url="test-host") + @patch("databricks.sql.client.Session") class TestTelemetryFeatureFlag: