From 8261c66bbef27db4ecbac3ff1d71e3fcc7afaacb Mon Sep 17 00:00:00 2001 From: Jay Xiao Date: Fri, 21 Aug 2026 19:04:47 +0000 Subject: [PATCH] Disable Python telemetry for kernel connections --- src/databricks/sql/backend/kernel/client.py | 31 ++++++++++ src/databricks/sql/session.py | 11 ++++ .../sql/telemetry/telemetry_client.py | 3 + tests/unit/test_kernel_client.py | 58 +++++++++++++++++ tests/unit/test_session.py | 61 +++++++++++++++++- tests/unit/test_telemetry.py | 62 +++++++++++++++++++ 6 files changed, 224 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index eeb496f73..83df73d1c 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -47,6 +47,7 @@ NotSupportedError, ProgrammingError, ) +from databricks.sql.telemetry.telemetry_client import TelemetryHelper from databricks.sql.thrift_api.TCLIService import ttypes if TYPE_CHECKING: @@ -165,6 +166,31 @@ def _is_staging_statement(operation: str) -> bool: return verb in _STAGING_VERBS +def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]: + """Build phase-7 telemetry/system kwargs for ``databricks_sql_kernel.Session``.""" + system = TelemetryHelper.get_driver_system_configuration() + out: Dict[str, Any] = { + "driver_name": system.driver_name, + "driver_version": system.driver_version, + "runtime_name": system.runtime_name, + "runtime_version": system.runtime_version, + "runtime_vendor": system.runtime_vendor, + "os_name": system.os_name, + "os_version": system.os_version, + "os_arch": system.os_arch, + "client_app_name": system.client_app_name, + "locale_name": system.locale_name, + "char_set_encoding": system.char_set_encoding, + # The Python telemetry model does not currently track process + # name; omit it and let the kernel fill what it can derive. + "process_name": None, + "telemetry_enabled": bool(options.get("enable_telemetry", True)), + } + if options.get("telemetry_batch_size") is not None: + out["telemetry_batch_size"] = options["telemetry_batch_size"] + return out + + # ─── Client ───────────────────────────────────────────────────────────────── @@ -217,6 +243,9 @@ def __init__( # to the kernel ``Session``'s ``retry_*`` kwargs in # ``open_session`` via ``_kernel_retry_kwargs``. self._retry_options = kwargs.get("retry_options") or {} + # Kernel telemetry phase 7 adds binding/runtime identity and + # telemetry config kwargs directly to ``databricks_sql_kernel.Session``. + self._telemetry_options = kwargs.get("telemetry_options") or {} self._catalog = catalog self._schema = schema # ``_use_arrow_native_complex_types`` is the connector-side @@ -316,6 +345,7 @@ def open_session( # Translate the connector's ``_retry_*`` kwargs into the # kernel's ``retry_*`` kwargs. Empty when at defaults. retry_kwargs = _kernel_retry_kwargs(self._retry_options) + telemetry_kwargs = _kernel_telemetry_kwargs(self._telemetry_options) # Forward caller / connector HTTP headers. The kernel applies # them on every request; a caller ``User-Agent`` is appended # to the kernel's base UA. Only pass the kwarg when there's @@ -358,6 +388,7 @@ def open_session( **auth_kwargs, **tls_kwargs, **retry_kwargs, + **telemetry_kwargs, **http_headers_kwargs, ) except Exception as exc: diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index a62b0d081..6eb9cf6a9 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -204,6 +204,16 @@ def _create_backend( "_retry_stop_after_attempts_duration" ), } + # Forward the binding/runtime identity and telemetry knobs + # added by kernel telemetry phase 7. Python-side telemetry + # still owns feature-flag evaluation and event export for the + # Thrift/SEA paths; the kernel path needs the same driver + # identity at Session construction time so kernel-owned + # telemetry can populate its system configuration. + kernel_telemetry_options = { + "enable_telemetry": kwargs.get("enable_telemetry", True), + "telemetry_batch_size": kwargs.get("telemetry_batch_size"), + } return KernelDatabricksClient( server_hostname=server_hostname, http_path=http_path, @@ -216,6 +226,7 @@ def _create_backend( _use_arrow_native_complex_types=_use_arrow_native_complex_types, auth_options=kernel_auth_options, retry_options=kernel_retry_options, + telemetry_options=kernel_telemetry_options, ) databricks_client_class: Type[DatabricksClient] diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py index a1ab4b5e0..7077b51f7 100644 --- a/src/databricks/sql/telemetry/telemetry_client.py +++ b/src/databricks/sql/telemetry/telemetry_client.py @@ -122,6 +122,9 @@ def get_auth_flow(auth_provider): @staticmethod def is_telemetry_enabled(connection: "Connection") -> bool: + if getattr(connection.session, "use_kernel", False) is True: + return False + # Fast path: force enabled - skip feature flag fetch entirely if connection.force_enable_telemetry: return True diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index 3253c6fff..c0918979f 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -344,6 +344,64 @@ def fake_session(**kw): assert captured.get("complex_types_as_json") is expected_flag +def test_open_session_passes_phase_7_telemetry_kwargs_to_kernel(monkeypatch): + """Kernel telemetry phase 7 added binding/runtime identity and + telemetry config kwargs to ``databricks_sql_kernel.Session``.""" + captured = {} + + def fake_session(**kw): + captured.update(kw) + sess = MagicMock() + sess.session_id = "sess-id" + return sess + + monkeypatch.setattr(kernel_client._kernel, "Session", fake_session) + monkeypatch.setattr( + kernel_client.TelemetryHelper, + "get_driver_system_configuration", + lambda: types.SimpleNamespace( + driver_name="Databricks SQL Python Connector", + driver_version="1.2.3", + runtime_name="Python 3.12.0", + runtime_version="3.12.0", + runtime_vendor="CPython", + os_name="Linux", + os_version="6.1", + os_arch="x86_64", + client_app_name=None, + locale_name="en_US", + char_set_encoding="utf-8", + ), + ) + + c = kernel_client.KernelDatabricksClient( + server_hostname="example.cloud.databricks.com", + http_path="/sql/1.0/warehouses/abc", + auth_provider=AccessTokenAuthProvider("dapi-test"), + ssl_options=None, + telemetry_options={ + "enable_telemetry": True, + "telemetry_batch_size": 17, + }, + ) + c.open_session(session_configuration=None, catalog=None, schema=None) + + assert captured["driver_name"] == "Databricks SQL Python Connector" + assert captured["driver_version"] == "1.2.3" + assert captured["runtime_name"] == "Python 3.12.0" + assert captured["runtime_version"] == "3.12.0" + assert captured["runtime_vendor"] == "CPython" + assert captured["os_name"] == "Linux" + assert captured["os_version"] == "6.1" + assert captured["os_arch"] == "x86_64" + assert captured["client_app_name"] is None + assert captured["locale_name"] == "en_US" + assert captured["char_set_encoding"] == "utf-8" + assert captured["process_name"] is None + assert captured["telemetry_enabled"] is True + assert captured["telemetry_batch_size"] == 17 + + def test_execute_command_forwards_parameters_to_bind_param(): """``execute_command(parameters=[...])`` routes each parameter through ``bind_tspark_params`` onto the kernel statement before diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index ba008b103..77c17670e 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -1,4 +1,5 @@ import pytest +import sys from unittest.mock import patch, MagicMock, Mock, PropertyMock import gc @@ -14,6 +15,14 @@ import databricks.sql +def _forget_kernel_client_module(): + sys.modules.pop("databricks.sql.backend.kernel.client", None) + import databricks.sql.backend.kernel as kernel_pkg + + if hasattr(kernel_pkg, "client"): + delattr(kernel_pkg, "client") + + class TestSession: """ Unit tests for Session functionality @@ -427,7 +436,6 @@ class TestKernelRetryOptionsThreading: PACKAGE = "databricks.sql" def test_retry_kwargs_threaded_into_kernel_client(self): - import sys import types pytest.importorskip( @@ -442,6 +450,7 @@ def test_retry_kwargs_threaded_into_kernel_client(self): fake = types.ModuleType("databricks_sql_kernel") fake.KernelError = type("KernelError", (Exception,), {}) fake.Session = MagicMock() + _forget_kernel_client_module() # Patch the kernel client class (imported lazily inside # _create_backend) and the provider builder; capture the kwargs @@ -478,6 +487,54 @@ def test_retry_kwargs_threaded_into_kernel_client(self): conn.close() +class TestKernelTelemetryOptionsThreading: + """The kernel path must forward telemetry options from connect() + into ``KernelDatabricksClient`` so phase-7 PyO3 Session kwargs can + be populated before the kernel opens its session.""" + + PACKAGE = "databricks.sql" + + def test_telemetry_kwargs_threaded_into_kernel_client(self): + import types + + pytest.importorskip( + "pyarrow", + reason="kernel client module imports pyarrow at load", + ) + + fake = types.ModuleType("databricks_sql_kernel") + fake.KernelError = type("KernelError", (Exception,), {}) + fake.Session = MagicMock() + _forget_kernel_client_module() + + with patch.dict(sys.modules, {"databricks_sql_kernel": fake}), patch( + "databricks.sql.backend.kernel.client.KernelDatabricksClient" + ) as mock_kernel_client, patch( + "%s.session.get_python_sql_connector_auth_provider" % self.PACKAGE + ): + instance = mock_kernel_client.return_value + instance.open_session.return_value = SessionId( + BackendType.SEA, "sess-id", None + ) + + conn = databricks.sql.connect( + server_hostname="foo", + http_path="/sql/1.0/warehouses/abc", + use_kernel=True, + access_token="dapi-xyz", + enable_telemetry=True, + force_enable_telemetry=False, + telemetry_batch_size=17, + ) + try: + _, kwargs = mock_kernel_client.call_args + opts = kwargs["telemetry_options"] + assert opts["enable_telemetry"] is True + assert opts["telemetry_batch_size"] == 17 + finally: + conn.close() + + class TestKernelUserAgentForwarding: """user_agent_entry must reach the kernel on the use_kernel path — session.py folds it into the composed User-Agent and includes it in @@ -488,7 +545,6 @@ class TestKernelUserAgentForwarding: PACKAGE = "databricks.sql" def test_user_agent_entry_reaches_kernel_client_http_headers(self): - import sys import types pytest.importorskip( @@ -498,6 +554,7 @@ def test_user_agent_entry_reaches_kernel_client_http_headers(self): fake = types.ModuleType("databricks_sql_kernel") fake.KernelError = type("KernelError", (Exception,), {}) fake.Session = MagicMock() + _forget_kernel_client_module() with patch.dict(sys.modules, {"databricks_sql_kernel": fake}), patch( "databricks.sql.backend.kernel.client.KernelDatabricksClient" diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py index 4f62fb833..9c6caa1c8 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -270,6 +270,68 @@ def test_token_federation_with_no_inner_provider(self): assert TelemetryHelper.get_auth_mechanism(fed) is None assert TelemetryHelper.get_auth_flow(fed) is None + @staticmethod + def _kernel_telemetry_kwargs_for_test(options): + import importlib + import sys + import types + + pytest.importorskip( + "pyarrow", + reason="kernel client module imports pyarrow at load", + ) + + fake = types.ModuleType("databricks_sql_kernel") + fake.KernelError = type("KernelError", (Exception,), {}) + fake.Session = MagicMock() + + sys.modules.pop("databricks.sql.backend.kernel.client", None) + import databricks.sql.backend.kernel as kernel_pkg + + if hasattr(kernel_pkg, "client"): + delattr(kernel_pkg, "client") + + try: + with patch.dict(sys.modules, {"databricks_sql_kernel": fake}): + kernel_client = importlib.import_module( + "databricks.sql.backend.kernel.client" + ) + return kernel_client._kernel_telemetry_kwargs(options) + finally: + sys.modules.pop("databricks.sql.backend.kernel.client", None) + if hasattr(kernel_pkg, "client"): + delattr(kernel_pkg, "client") + + @pytest.mark.parametrize( + ("enable_telemetry", "expected_kernel_telemetry_enabled"), + [ + (True, True), + (False, False), + ], + ) + def test_is_telemetry_enabled_returns_false_for_kernel( + self, + enable_telemetry, + expected_kernel_telemetry_enabled, + ): + connection = MagicMock() + connection.session.use_kernel = True + connection.force_enable_telemetry = True + connection.enable_telemetry = enable_telemetry + + assert TelemetryHelper.is_telemetry_enabled(connection) is False + + kernel_kwargs = self._kernel_telemetry_kwargs_for_test( + { + "enable_telemetry": enable_telemetry, + "force_enable_telemetry": True, + } + ) + assert ( + kernel_kwargs["telemetry_enabled"] + is expected_kernel_telemetry_enabled + ) + class TestTelemetryFactory: """Tests for TelemetryClientFactory lifecycle and management."""