diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 3d3a33b86..caa7b93f5 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -22,6 +22,7 @@ from __future__ import annotations +import inspect import logging import threading import uuid @@ -47,6 +48,7 @@ NotSupportedError, ProgrammingError, ) +from databricks.sql.telemetry.telemetry_client import TelemetryHelper if TYPE_CHECKING: from databricks.sql.client import Cursor @@ -172,6 +174,75 @@ def _is_staging_statement(operation: str) -> bool: return verb in _STAGING_VERBS +def _kernel_session_accepts_kwarg(name: str) -> bool: + """True iff the installed ``databricks_sql_kernel.Session`` constructor + declares keyword ``name``. + + The kernel ``Session`` is a PyO3 class with a **fixed** signature (no + ``**kwargs`` catch-all), so forwarding a kwarg it doesn't declare raises + ``TypeError`` at construction. The phase-7 identity/telemetry kwargs + (``driver_name`` etc.) only exist on wheels newer than the pinned + ``^0.2.0`` (whose ``Session`` accepts none of them), so we must gate them + on what the actually-installed wheel supports rather than pass them + unconditionally. Falls **closed** (returns ``False``) when the signature + can't be introspected: a PyO3 class only exposes ``__text_signature__`` + (and thus an introspectable signature) when built with + ``#[pyo3(signature=...)]``; otherwise ``inspect.signature`` raises + ``ValueError``. Since the pinned ``^0.2.0`` ``Session`` accepts none of + these kwargs, forwarding one it doesn't declare is a hard ``TypeError`` at + construction that breaks every ``use_kernel=True`` connection, whereas + omitting one the wheel *would* have accepted only loses telemetry + richness — so we omit the kwarg on introspection failure. + """ + try: + params = inspect.signature(_kernel.Session).parameters + except (TypeError, ValueError): + return False + if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()): + return True + return name in params + + +def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]: + """Build phase-7 telemetry/system kwargs for ``databricks_sql_kernel.Session``. + + Only kwargs the installed ``Session`` constructor actually accepts are + returned; on the pinned ``^0.2.0`` wheel (which predates phase 7) this is + empty, so ``open_session`` doesn't break with ``TypeError`` on a wheel + that doesn't yet know these kwargs. + """ + system = TelemetryHelper.get_driver_system_configuration() + candidates: 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, + } + if options.get("enable_telemetry") is not None: + candidates["telemetry_enabled"] = bool(options["enable_telemetry"]) + if options.get("telemetry_batch_size") is not None: + candidates["telemetry_batch_size"] = options["telemetry_batch_size"] + if options.get("telemetry_circuit_breaker_enabled") is not None: + candidates["telemetry_circuit_breaker_enabled"] = options[ + "telemetry_circuit_breaker_enabled" + ] + return { + name: value + for name, value in candidates.items() + if _kernel_session_accepts_kwarg(name) + } + + # ─── Client ───────────────────────────────────────────────────────────────── @@ -226,6 +297,9 @@ def __init__( self._retry_options = kwargs.get("retry_options") or {} # The kernel binding owns type and range validation. self._request_timeout_secs = kwargs.get("request_timeout_secs") + # 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 @@ -339,6 +413,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 @@ -382,6 +457,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/client.py b/src/databricks/sql/client.py index 1f3a8f69d..4d9ca0327 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -178,6 +178,16 @@ def __init__( ``ImportError``. Supports PAT, OAuth M2M, and OAuth U2M auth, and native (positional and named) parameter binding. Mutually exclusive with ``use_sea``. + + Telemetry note: on the kernel path telemetry is + strictly opt-in. An explicit ``enable_telemetry=True`` + turns telemetry on unconditionally; it is NOT gated by + the server-side ``enableTelemetryForPythonDriver`` + feature flag that applies on the Thrift/SEA paths. When + ``enable_telemetry`` is unset the kernel owns the enable + decision. This is an intentional divergence from the + Thrift/SEA paths, where an explicit ``True`` can still + be suppressed by the feature flag. :param use_hybrid_disposition: `bool`, optional (default is False) Use the hybrid disposition instead of the inline disposition. :param server_hostname: Databricks instance host name. @@ -402,8 +412,12 @@ def read(self) -> Optional[OAuthToken]: ) self.session.open() except Exception as e: - # Respect user's telemetry preference even during connection failure - enable_telemetry = kwargs.get("enable_telemetry", True) + # Respect user's telemetry preference even during connection failure. + # For use_kernel connections the kernel owns telemetry, so suppress + # the wrapper-side failure log to avoid wrapper-vs-kernel duplication. + enable_telemetry = kwargs.get("enable_telemetry", True) and not kwargs.get( + "use_kernel", False + ) TelemetryClientFactory.connection_failure_log( error_name="Exception", error_message=str(e), diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index bced2db97..a4166ef13 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -13,6 +13,7 @@ from databricks.sql.backend.types import SessionId, BackendType from databricks.sql.common.unified_http_client import UnifiedHttpClient from databricks.sql.common.agent import detect as detect_agent +from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory if TYPE_CHECKING: from databricks.sql.backend.thrift_backend import ThriftDatabricksClient @@ -250,6 +251,55 @@ 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 = { + # Preserve the caller's explicit telemetry choice. When unset, + # leave it as None so the kernel owns the enable decision -- + # this is an INTENTIONAL divergence from the Thrift/SEA path, + # not an oversight. There, an unset enable_telemetry defaults to + # True (client.py) but only actually emits when the + # `enableTelemetryForPythonDriver` server feature flag is on + # (telemetry_client.py). That feature-flag gate is Python-side + # and is bypassed on the kernel path (is_telemetry_enabled + # short-circuits to False for use_kernel), so forwarding the + # connector's True default here would force telemetry on without + # an equivalent gate. Passing None instead defers to the + # kernel's own default/gating, which is expected to mirror the + # feature-flag-gated wrapper behaviour; only an explicit caller + # opt-in/opt-out overrides it. + "enable_telemetry": kwargs.get("enable_telemetry"), + # Match the connector's default batch size (client.py forwards + # the same TelemetryClientFactory.DEFAULT_BATCH_SIZE fallback) + # so an unset telemetry_batch_size resolves to the same value + # on the kernel path as on the Thrift/SEA path, rather than + # letting the kernel silently pick its own internal default. + "telemetry_batch_size": kwargs.get( + "telemetry_batch_size", TelemetryClientFactory.DEFAULT_BATCH_SIZE + ), + # Preserve the caller's explicit circuit-breaker choice. When + # unset, leave it as None so the kernel owns the decision -- + # mirroring the enable_telemetry handling above rather than + # forcing a default. This is deliberately NOT defaulted to True: + # although ClientContext's signature default is True + # (auth/common.py:55), the Thrift/SEA path never reaches it -- + # build_client_context (utils.py:1018) always passes + # _telemetry_circuit_breaker_enabled explicitly (None when the + # caller leaves it unset), and ClientContext coerces it with + # bool(None) -> False (auth/common.py:89). So the *effective* + # Thrift/SEA default when unset is False, not True; forwarding + # True here would turn the circuit breaker on for an + # unconfigured connection while Thrift/SEA leaves it off. + # Passing None instead defers to the kernel's own default; + # only an explicit caller value overrides it. + "telemetry_circuit_breaker_enabled": kwargs.get( + "_telemetry_circuit_breaker_enabled" + ), + } return KernelDatabricksClient( server_hostname=server_hostname, http_path=http_path, @@ -263,6 +313,7 @@ def _create_backend( auth_options=kernel_auth_options, retry_options=kernel_retry_options, request_timeout_secs=kwargs.get("_socket_timeout"), + telemetry_options=kernel_telemetry_options, ) # These reference the lazily-resolved module attributes defined via diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py index a1ab4b5e0..2051fb2f8 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): + 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 7e8249553..5fbf81ae7 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -368,6 +368,208 @@ def fake_session(**kw): assert captured["request_timeout_secs"] == timeout +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, + "telemetry_circuit_breaker_enabled": False, + }, + ) + 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 + assert captured["telemetry_circuit_breaker_enabled"] is False + + +def test_open_session_omits_phase_7_kwargs_kernel_does_not_accept(monkeypatch): + """Phase-7 identity/telemetry kwargs must NOT be forwarded to a kernel + ``Session`` whose (fixed, no-``**kwargs``) constructor doesn't declare + them. + + The real ``databricks_sql_kernel.Session`` is a PyO3 class with a fixed + signature; the pinned ``^0.2.0`` wheel predates phase 7 and accepts none + of these kwargs, so forwarding them unconditionally raises ``TypeError`` + and breaks every ``use_kernel=True`` connection. The other tests here use + a ``**kwargs`` MagicMock that silently swallows the kwargs and hides the + break; this one uses a fixed-signature fake mirroring the real 0.2.0 + surface to prove the client gates on what the installed Session supports. + """ + captured = {} + + # Fixed signature mirroring the pinned 0.2.0 kernel Session: it accepts + # the base connection/tls/retry kwargs but NONE of the phase-7 identity + # or telemetry kwargs, and has no **kwargs catch-all. + def fake_session_v0_2_0( + host, + http_path, + *, + auth_type=None, + access_token=None, + client_id=None, + client_secret=None, + oauth_scopes=None, + token_url=None, + redirect_port=None, + oauth_callback_timeout_secs=None, + tls_ca_cert=None, + tls_skip_verify=False, + tls_skip_hostname_verify=False, + tls_client_cert=None, + tls_client_key=None, + retry_min_wait_secs=None, + retry_max_wait_secs=None, + retry_max_attempts=None, + retry_overall_timeout_secs=None, + http_headers=None, + catalog=None, + schema=None, + session_conf=None, + complex_types_as_json=False, + intervals_as_string=False, + request_timeout_secs=None, + ): + captured["host"] = host + sess = MagicMock() + sess.session_id = "sess-id" + return sess + + monkeypatch.setattr(kernel_client._kernel, "Session", fake_session_v0_2_0) + 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", + ), + ) + + # The kwargs the client builds must be filtered to what fake_session + # accepts, so open_session succeeds instead of raising TypeError. + kwargs = kernel_client._kernel_telemetry_kwargs( + {"enable_telemetry": True, "telemetry_batch_size": 17} + ) + assert kwargs == {}, f"expected no phase-7 kwargs on 0.2.0 Session, got {kwargs}" + + 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}, + ) + # Would raise TypeError: unexpected keyword argument if the client + # forwarded phase-7 kwargs the fixed-signature Session doesn't declare. + c.open_session(session_configuration=None, catalog=None, schema=None) + assert captured["host"] == "example.cloud.databricks.com" + + +def test_kernel_session_accepts_kwarg_falls_closed_when_not_introspectable(monkeypatch): + """When ``inspect.signature(_kernel.Session)`` raises (a PyO3 class built + without ``#[pyo3(signature=...)]`` exposes no ``__text_signature__``, so + ``inspect.signature`` raises ``ValueError``), the gate must fall + **closed** and omit every phase-7 kwarg. + + Forwarding a kwarg the installed ``Session`` doesn't declare is a hard + ``TypeError`` at construction that breaks every ``use_kernel=True`` + connection; omitting one it would have accepted only loses telemetry + richness. The fixed-signature fake in the sibling test always introspects, + so this test uses a stand-in whose signature genuinely can't be read to + cover the real non-introspectable PyO3 binding. + """ + + class NonIntrospectableSession: + # Mirrors a PyO3 class with no exposed __text_signature__: + # inspect.signature() raises ValueError on it. + def __init__(self, *args, **kwargs): # pragma: no cover - never called + pass + + def raise_value_error(_obj): + raise ValueError("no signature found for builtin type") + + monkeypatch.setattr(kernel_client._kernel, "Session", NonIntrospectableSession) + monkeypatch.setattr(kernel_client.inspect, "signature", raise_value_error) + + assert kernel_client._kernel_session_accepts_kwarg("driver_name") is False + + 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", + ), + ) + kwargs = kernel_client._kernel_telemetry_kwargs( + {"enable_telemetry": True, "telemetry_batch_size": 17} + ) + assert kwargs == {}, f"expected no phase-7 kwargs when signature unreadable, got {kwargs}" + + 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 2da8bd194..e5339d9ca 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -570,6 +570,92 @@ def test_oauth_token_cache_enabled_threaded_into_kernel_auth_options(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 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() + + 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, + _telemetry_circuit_breaker_enabled=False, + ) + try: + _, kwargs = mock_kernel_client.call_args + opts = kwargs["telemetry_options"] + assert opts["enable_telemetry"] is True + assert opts["telemetry_batch_size"] == 17 + assert opts["telemetry_circuit_breaker_enabled"] is False + finally: + conn.close() + + def test_telemetry_enabled_defaults_none_for_kernel_client(self): + 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() + + 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", + ) + try: + _, kwargs = mock_kernel_client.call_args + opts = kwargs["telemetry_options"] + assert opts["enable_telemetry"] is None + 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 diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py index 4f62fb833..d2d69b9f9 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -270,6 +270,80 @@ 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 + ) + + def test_kernel_telemetry_enabled_omitted_when_unset(self): + kernel_kwargs = self._kernel_telemetry_kwargs_for_test({}) + + assert "telemetry_enabled" not in kernel_kwargs + + def test_kernel_telemetry_enabled_omitted_when_none(self): + kernel_kwargs = self._kernel_telemetry_kwargs_for_test( + {"enable_telemetry": None} + ) + + assert "telemetry_enabled" not in kernel_kwargs + class TestTelemetryFactory: """Tests for TelemetryClientFactory lifecycle and management.""" @@ -421,6 +495,36 @@ def test_connection_failure_sends_correct_telemetry_payload( assert call_arguments[0][0] == "Exception" assert call_arguments[0][1] == error_message + @patch( + "databricks.sql.telemetry.telemetry_client.TelemetryClient.export_failure_log" + ) + @patch("databricks.sql.client.Session") + def test_connection_failure_does_not_send_telemetry_for_kernel( + self, mock_session, mock_export_failure_log + ): + """ + A use_kernel=True connection that fails to open must NOT emit a + wrapper-side failure log — the kernel owns telemetry, so emitting + here would duplicate the kernel's own failure reporting. + """ + + error_message = "Could not connect to host" + mock_session_instance = MagicMock() + mock_session_instance.is_open = False + mock_session_instance.open.side_effect = Exception(error_message) + mock_session.return_value = mock_session_instance + + try: + sql.connect( + server_hostname="test-host", + http_path="/test-path", + use_kernel=True, + ) + except Exception as e: + assert str(e) == error_message + + mock_export_failure_log.assert_not_called() + @patch("databricks.sql.client.Session") class TestTelemetryFeatureFlag: @@ -457,6 +561,7 @@ def test_telemetry_enabled_when_flag_is_true(self, mock_http_request, MockSessio self._mock_ff_response(mock_http_request, enabled=True) mock_session_instance = MockSession.return_value mock_session_instance.guid_hex = "test-session-ff-true" + mock_session_instance.use_kernel = False mock_session_instance.host = "test-host" # Set host for telemetry client lookup mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = ( @@ -488,6 +593,7 @@ def test_telemetry_disabled_when_flag_is_false( self._mock_ff_response(mock_http_request, enabled=False) mock_session_instance = MockSession.return_value mock_session_instance.guid_hex = "test-session-ff-false" + mock_session_instance.use_kernel = False mock_session_instance.host = "test-host" # Set host for telemetry client lookup mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = ( @@ -519,6 +625,7 @@ def test_telemetry_disabled_when_flag_request_fails( mock_http_request.side_effect = Exception("Network is down") mock_session_instance = MockSession.return_value mock_session_instance.guid_hex = "test-session-ff-fail" + mock_session_instance.use_kernel = False mock_session_instance.host = "test-host" # Set host for telemetry client lookup mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = ( @@ -771,6 +878,7 @@ def test_connection_with_proxy_populates_telemetry(self, mock_setup_pools, mock_ """Test that proxy configuration is captured in telemetry.""" mock_session_instance = MagicMock() mock_session_instance.guid_hex = "test-session-proxy" + mock_session_instance.use_kernel = False mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = False mock_session_instance.use_sea = True @@ -806,6 +914,7 @@ def test_connection_with_azure_params_populates_telemetry(self, mock_setup_pools """Test that Azure-specific parameters are captured in telemetry.""" mock_session_instance = MagicMock() mock_session_instance.guid_hex = "test-session-azure" + mock_session_instance.use_kernel = False mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = False mock_session_instance.use_sea = False @@ -835,6 +944,7 @@ def test_connection_populates_arrow_and_performance_params(self, mock_setup_pool """Test that Arrow and performance parameters are captured in telemetry.""" mock_session_instance = MagicMock() mock_session_instance.guid_hex = "test-session-perf" + mock_session_instance.use_kernel = False mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = False mock_session_instance.use_sea = True @@ -879,6 +989,7 @@ def test_federated_pat_populates_telemetry_as_pat(self, mock_setup_pools, mock_s ) mock_session_instance = MagicMock() mock_session_instance.guid_hex = "test-session-fed-pat" + mock_session_instance.use_kernel = False mock_session_instance.auth_provider = federated_pat mock_session_instance.is_open = False mock_session_instance.use_sea = False @@ -906,6 +1017,7 @@ def test_cf_proxy_fields_default_to_false_none(self, mock_setup_pools, mock_sess """Test that CloudFlare proxy fields default to False/None (not yet supported).""" mock_session_instance = MagicMock() mock_session_instance.guid_hex = "test-session-cfproxy" + mock_session_instance.use_kernel = False mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = False mock_session_instance.use_sea = True