Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
14 changes: 13 additions & 1 deletion src/databricks/sql/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,7 @@
TSparkParameter,
)
from databricks.sql.telemetry.telemetry_client import (
NoopTelemetryClient,
TelemetryHelper,
TelemetryClientFactory,
)
Expand Down Expand Up @@ -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()
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Resolved
return

TelemetryClientFactory.initialize_telemetry_client(
telemetry_enabled=self.telemetry_enabled,
session_id_hex=self.get_session_id_hex(),
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand Down
11 changes: 4 additions & 7 deletions tests/unit/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
OperationalError,
)
from databricks.sql.session import Session
from databricks.sql.telemetry.telemetry_client import NoopTelemetryClient

import databricks.sql

Expand Down Expand Up @@ -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,)]
Expand Down
64 changes: 64 additions & 0 deletions tests/unit/test_telemetry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
Loading