From 6e6eab598ed1733af42f47bd86ab35be4af8bb82 Mon Sep 17 00:00:00 2001 From: Jay Xiao Date: Fri, 21 Aug 2026 19:04:47 +0000 Subject: [PATCH 01/17] Disable Python telemetry for kernel connections Signed-off-by: Jay Xiao --- src/databricks/sql/backend/kernel/client.py | 45 +++++++++++++ src/databricks/sql/session.py | 23 +++++++ .../sql/telemetry/telemetry_client.py | 3 + tests/unit/test_kernel_client.py | 66 +++++++++++++++++++ tests/unit/test_session.py | 56 ++++++++++++++++ tests/unit/test_telemetry.py | 62 +++++++++++++++++ 6 files changed, 255 insertions(+) diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 3d3a33b86..ff72082ff 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 if TYPE_CHECKING: from databricks.sql.client import Cursor @@ -172,6 +173,45 @@ 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"] + if options.get("telemetry_flush_interval_ms") is not None: + out["telemetry_flush_interval_ms"] = options["telemetry_flush_interval_ms"] + if options.get("telemetry_circuit_breaker_enabled") is not None: + out["telemetry_circuit_breaker_enabled"] = options[ + "telemetry_circuit_breaker_enabled" + ] + if options.get("telemetry_circuit_breaker_threshold") is not None: + out["telemetry_circuit_breaker_threshold"] = options[ + "telemetry_circuit_breaker_threshold" + ] + if options.get("telemetry_circuit_breaker_timeout_ms") is not None: + out["telemetry_circuit_breaker_timeout_ms"] = options[ + "telemetry_circuit_breaker_timeout_ms" + ] + return out + + # ─── Client ───────────────────────────────────────────────────────────────── @@ -226,6 +266,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 +382,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 +426,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 bced2db97..828d04fe0 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -250,6 +250,28 @@ 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"), + "telemetry_flush_interval_ms": kwargs.get( + "telemetry_flush_interval_ms" + ), + "telemetry_circuit_breaker_enabled": kwargs.get( + "_telemetry_circuit_breaker_enabled" + ), + "telemetry_circuit_breaker_threshold": kwargs.get( + "telemetry_circuit_breaker_threshold" + ), + "telemetry_circuit_breaker_timeout_ms": kwargs.get( + "telemetry_circuit_breaker_timeout_ms" + ), + } return KernelDatabricksClient( server_hostname=server_hostname, http_path=http_path, @@ -263,6 +285,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..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 7e8249553..1da16bb4f 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -368,6 +368,72 @@ 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_flush_interval_ms": 250, + "telemetry_circuit_breaker_enabled": False, + "telemetry_circuit_breaker_threshold": 9, + "telemetry_circuit_breaker_timeout_ms": 60000, + }, + ) + 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_flush_interval_ms"] == 250 + assert captured["telemetry_circuit_breaker_enabled"] is False + assert captured["telemetry_circuit_breaker_threshold"] == 9 + assert captured["telemetry_circuit_breaker_timeout_ms"] == 60000 + + 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..748de7d07 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -570,6 +570,62 @@ 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_flush_interval_ms=250, + _telemetry_circuit_breaker_enabled=False, + telemetry_circuit_breaker_threshold=9, + telemetry_circuit_breaker_timeout_ms=60000, + ) + 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_flush_interval_ms"] == 250 + assert opts["telemetry_circuit_breaker_enabled"] is False + assert opts["telemetry_circuit_breaker_threshold"] == 9 + assert opts["telemetry_circuit_breaker_timeout_ms"] == 60000 + 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..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.""" From 17eaa1db91903698318c4f1479de1e8788b757a1 Mon Sep 17 00:00:00 2001 From: Jay Xiao Date: Thu, 27 Aug 2026 21:59:35 +0000 Subject: [PATCH 02/17] Align kernel telemetry circuit breaker options with Thrift Signed-off-by: Jay Xiao --- src/databricks/sql/backend/kernel/client.py | 10 ---------- src/databricks/sql/session.py | 9 --------- tests/unit/test_kernel_client.py | 6 ------ tests/unit/test_session.py | 6 ------ 4 files changed, 31 deletions(-) diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index ff72082ff..83d9715f4 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -195,20 +195,10 @@ def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]: } if options.get("telemetry_batch_size") is not None: out["telemetry_batch_size"] = options["telemetry_batch_size"] - if options.get("telemetry_flush_interval_ms") is not None: - out["telemetry_flush_interval_ms"] = options["telemetry_flush_interval_ms"] if options.get("telemetry_circuit_breaker_enabled") is not None: out["telemetry_circuit_breaker_enabled"] = options[ "telemetry_circuit_breaker_enabled" ] - if options.get("telemetry_circuit_breaker_threshold") is not None: - out["telemetry_circuit_breaker_threshold"] = options[ - "telemetry_circuit_breaker_threshold" - ] - if options.get("telemetry_circuit_breaker_timeout_ms") is not None: - out["telemetry_circuit_breaker_timeout_ms"] = options[ - "telemetry_circuit_breaker_timeout_ms" - ] return out diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index 828d04fe0..30e4f975e 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -259,18 +259,9 @@ def _create_backend( kernel_telemetry_options = { "enable_telemetry": kwargs.get("enable_telemetry", True), "telemetry_batch_size": kwargs.get("telemetry_batch_size"), - "telemetry_flush_interval_ms": kwargs.get( - "telemetry_flush_interval_ms" - ), "telemetry_circuit_breaker_enabled": kwargs.get( "_telemetry_circuit_breaker_enabled" ), - "telemetry_circuit_breaker_threshold": kwargs.get( - "telemetry_circuit_breaker_threshold" - ), - "telemetry_circuit_breaker_timeout_ms": kwargs.get( - "telemetry_circuit_breaker_timeout_ms" - ), } return KernelDatabricksClient( server_hostname=server_hostname, diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index 1da16bb4f..e707fd25c 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -406,10 +406,7 @@ def fake_session(**kw): telemetry_options={ "enable_telemetry": True, "telemetry_batch_size": 17, - "telemetry_flush_interval_ms": 250, "telemetry_circuit_breaker_enabled": False, - "telemetry_circuit_breaker_threshold": 9, - "telemetry_circuit_breaker_timeout_ms": 60000, }, ) c.open_session(session_configuration=None, catalog=None, schema=None) @@ -428,10 +425,7 @@ def fake_session(**kw): assert captured["process_name"] is None assert captured["telemetry_enabled"] is True assert captured["telemetry_batch_size"] == 17 - assert captured["telemetry_flush_interval_ms"] == 250 assert captured["telemetry_circuit_breaker_enabled"] is False - assert captured["telemetry_circuit_breaker_threshold"] == 9 - assert captured["telemetry_circuit_breaker_timeout_ms"] == 60000 def test_execute_command_forwards_parameters_to_bind_param(): diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 748de7d07..3ff194480 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -608,20 +608,14 @@ def test_telemetry_kwargs_threaded_into_kernel_client(self): enable_telemetry=True, force_enable_telemetry=False, telemetry_batch_size=17, - telemetry_flush_interval_ms=250, _telemetry_circuit_breaker_enabled=False, - telemetry_circuit_breaker_threshold=9, - telemetry_circuit_breaker_timeout_ms=60000, ) 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_flush_interval_ms"] == 250 assert opts["telemetry_circuit_breaker_enabled"] is False - assert opts["telemetry_circuit_breaker_threshold"] == 9 - assert opts["telemetry_circuit_breaker_timeout_ms"] == 60000 finally: conn.close() From c71dbf95844a34861e16b04dbce3aae6c6ca3fd9 Mon Sep 17 00:00:00 2001 From: Jay Xiao Date: Thu, 27 Aug 2026 22:04:21 +0000 Subject: [PATCH 03/17] Forward Python telemetry defaults to kernel Signed-off-by: Jay Xiao --- src/databricks/sql/backend/kernel/client.py | 14 +++++++++++++- tests/unit/test_kernel_client.py | 3 +++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 83d9715f4..0c35ed470 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -47,7 +47,14 @@ NotSupportedError, ProgrammingError, ) -from databricks.sql.telemetry.telemetry_client import TelemetryHelper +from databricks.sql.telemetry.circuit_breaker_manager import ( + MINIMUM_CALLS, + RESET_TIMEOUT, +) +from databricks.sql.telemetry.telemetry_client import ( + TelemetryClientFactory, + TelemetryHelper, +) if TYPE_CHECKING: from databricks.sql.client import Cursor @@ -192,6 +199,11 @@ def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]: # name; omit it and let the kernel fill what it can derive. "process_name": None, "telemetry_enabled": bool(options.get("enable_telemetry", True)), + "telemetry_flush_interval_ms": ( + TelemetryClientFactory._flush_interval_seconds * 1000 + ), + "telemetry_circuit_breaker_threshold": MINIMUM_CALLS, + "telemetry_circuit_breaker_timeout_ms": RESET_TIMEOUT * 1000, } if options.get("telemetry_batch_size") is not None: out["telemetry_batch_size"] = options["telemetry_batch_size"] diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index e707fd25c..281e1a6bb 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -425,7 +425,10 @@ def fake_session(**kw): assert captured["process_name"] is None assert captured["telemetry_enabled"] is True assert captured["telemetry_batch_size"] == 17 + assert captured["telemetry_flush_interval_ms"] == 300000 assert captured["telemetry_circuit_breaker_enabled"] is False + assert captured["telemetry_circuit_breaker_threshold"] == 20 + assert captured["telemetry_circuit_breaker_timeout_ms"] == 30000 def test_execute_command_forwards_parameters_to_bind_param(): From 989c2250c5bb06186d60dea6430aff620889ac37 Mon Sep 17 00:00:00 2001 From: Jay Xiao Date: Thu, 27 Aug 2026 22:07:02 +0000 Subject: [PATCH 04/17] Revert "Forward Python telemetry defaults to kernel" This reverts commit e615bbf664d5f8bed0b468dde7490c8002279826. Signed-off-by: Jay Xiao --- src/databricks/sql/backend/kernel/client.py | 14 +------------- tests/unit/test_kernel_client.py | 3 --- 2 files changed, 1 insertion(+), 16 deletions(-) diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 0c35ed470..83d9715f4 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -47,14 +47,7 @@ NotSupportedError, ProgrammingError, ) -from databricks.sql.telemetry.circuit_breaker_manager import ( - MINIMUM_CALLS, - RESET_TIMEOUT, -) -from databricks.sql.telemetry.telemetry_client import ( - TelemetryClientFactory, - TelemetryHelper, -) +from databricks.sql.telemetry.telemetry_client import TelemetryHelper if TYPE_CHECKING: from databricks.sql.client import Cursor @@ -199,11 +192,6 @@ def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]: # name; omit it and let the kernel fill what it can derive. "process_name": None, "telemetry_enabled": bool(options.get("enable_telemetry", True)), - "telemetry_flush_interval_ms": ( - TelemetryClientFactory._flush_interval_seconds * 1000 - ), - "telemetry_circuit_breaker_threshold": MINIMUM_CALLS, - "telemetry_circuit_breaker_timeout_ms": RESET_TIMEOUT * 1000, } if options.get("telemetry_batch_size") is not None: out["telemetry_batch_size"] = options["telemetry_batch_size"] diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index 281e1a6bb..e707fd25c 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -425,10 +425,7 @@ def fake_session(**kw): assert captured["process_name"] is None assert captured["telemetry_enabled"] is True assert captured["telemetry_batch_size"] == 17 - assert captured["telemetry_flush_interval_ms"] == 300000 assert captured["telemetry_circuit_breaker_enabled"] is False - assert captured["telemetry_circuit_breaker_threshold"] == 20 - assert captured["telemetry_circuit_breaker_timeout_ms"] == 30000 def test_execute_command_forwards_parameters_to_bind_param(): From 45a44cda0dc451540ae2c9cac1e3c744b53a3d26 Mon Sep 17 00:00:00 2001 From: Jay Xiao Date: Thu, 27 Aug 2026 22:07:45 +0000 Subject: [PATCH 05/17] Default kernel telemetry to disabled Signed-off-by: Jay Xiao --- src/databricks/sql/backend/kernel/client.py | 2 +- src/databricks/sql/session.py | 2 +- tests/unit/test_session.py | 36 +++++++++++++++++++++ tests/unit/test_telemetry.py | 5 +++ 4 files changed, 43 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 83d9715f4..17c6d7c95 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -191,7 +191,7 @@ def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]: # 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)), + "telemetry_enabled": bool(options.get("enable_telemetry", False)), } if options.get("telemetry_batch_size") is not None: out["telemetry_batch_size"] = options["telemetry_batch_size"] diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index 30e4f975e..73259e1bd 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -257,7 +257,7 @@ def _create_backend( # identity at Session construction time so kernel-owned # telemetry can populate its system configuration. kernel_telemetry_options = { - "enable_telemetry": kwargs.get("enable_telemetry", True), + "enable_telemetry": kwargs.get("enable_telemetry", False), "telemetry_batch_size": kwargs.get("telemetry_batch_size"), "telemetry_circuit_breaker_enabled": kwargs.get( "_telemetry_circuit_breaker_enabled" diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 3ff194480..6b3f28c12 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -619,6 +619,42 @@ def test_telemetry_kwargs_threaded_into_kernel_client(self): finally: conn.close() + def test_telemetry_enabled_defaults_false_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 False + finally: + conn.close() + class TestKernelUserAgentForwarding: """user_agent_entry must reach the kernel on the use_kernel path — diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py index 9c6caa1c8..48f71b044 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -332,6 +332,11 @@ def test_is_telemetry_enabled_returns_false_for_kernel( is expected_kernel_telemetry_enabled ) + def test_kernel_telemetry_enabled_defaults_false(self): + kernel_kwargs = self._kernel_telemetry_kwargs_for_test({}) + + assert kernel_kwargs["telemetry_enabled"] is False + class TestTelemetryFactory: """Tests for TelemetryClientFactory lifecycle and management.""" From 6cba80943130e9339f80085e754b81be06dc7bf5 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Thu, 27 Aug 2026 23:38:37 +0000 Subject: [PATCH 06/17] ai: apply changes for #925 (5 review threads) Addresses: - #3833615986 at src/databricks/sql/telemetry/telemetry_client.py:125 - #3833615990 at src/databricks/sql/backend/kernel/client.py:176 - #3833661667 at src/databricks/sql/backend/kernel/client.py:419 - #3858839800 at src/databricks/sql/session.py:228 - #3876277610 at src/databricks/sql/session.py:260 Signed-off-by: peco-engineer-bot[bot] --- src/databricks/sql/backend/kernel/client.py | 3 +++ src/databricks/sql/session.py | 16 +++++++++++++++- src/databricks/sql/telemetry/telemetry_client.py | 2 +- tests/unit/test_kernel_sig_probe.py | 0 tests/unit/test_telemetry.py | 8 ++++++++ tests/unit/test_zz_kernel_sig_probe.py | 1 + 6 files changed, 28 insertions(+), 2 deletions(-) create mode 100644 tests/unit/test_kernel_sig_probe.py create mode 100644 tests/unit/test_zz_kernel_sig_probe.py diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 17c6d7c95..f93ceb9de 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -191,6 +191,9 @@ def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]: # The Python telemetry model does not currently track process # name; omit it and let the kernel fill what it can derive. "process_name": None, + # Defaults to False by design: the kernel path deliberately diverges + # from the connector-wide True default (see session.py and client.py). + # Telemetry is off unless explicitly enabled on this backend. "telemetry_enabled": bool(options.get("enable_telemetry", False)), } if options.get("telemetry_batch_size") is not None: diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index 73259e1bd..ad23f1fcc 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 @@ -257,8 +258,21 @@ def _create_backend( # identity at Session construction time so kernel-owned # telemetry can populate its system configuration. kernel_telemetry_options = { + # Intentionally defaults to False, diverging from the + # connector-wide True default on the Thrift/SEA path + # (client.py). The kernel path opts out of telemetry unless + # explicitly enabled; this is asserted by + # test_telemetry_enabled_defaults_false_for_kernel_client. + # Do not "fix" this back to True to match the other backends. "enable_telemetry": kwargs.get("enable_telemetry", False), - "telemetry_batch_size": kwargs.get("telemetry_batch_size"), + # 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 + ), "telemetry_circuit_breaker_enabled": kwargs.get( "_telemetry_circuit_breaker_enabled" ), diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py index 7077b51f7..2051fb2f8 100644 --- a/src/databricks/sql/telemetry/telemetry_client.py +++ b/src/databricks/sql/telemetry/telemetry_client.py @@ -122,7 +122,7 @@ def get_auth_flow(auth_provider): @staticmethod def is_telemetry_enabled(connection: "Connection") -> bool: - if getattr(connection.session, "use_kernel", False) is True: + if getattr(connection.session, "use_kernel", False): return False # Fast path: force enabled - skip feature flag fetch entirely diff --git a/tests/unit/test_kernel_sig_probe.py b/tests/unit/test_kernel_sig_probe.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py index 48f71b044..abcdd6767 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -524,6 +524,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 = ( @@ -555,6 +556,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 = ( @@ -586,6 +588,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 = ( @@ -838,6 +841,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 @@ -873,6 +877,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 @@ -902,6 +907,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 @@ -946,6 +952,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 @@ -973,6 +980,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 diff --git a/tests/unit/test_zz_kernel_sig_probe.py b/tests/unit/test_zz_kernel_sig_probe.py new file mode 100644 index 000000000..80f73fb4e --- /dev/null +++ b/tests/unit/test_zz_kernel_sig_probe.py @@ -0,0 +1 @@ +# (temporary probe removed) From 78ef1a7321583911cc70e7723e85b174bdbf7927 Mon Sep 17 00:00:00 2001 From: Jay Xiao Date: Fri, 28 Aug 2026 00:09:24 +0000 Subject: [PATCH 07/17] Remove temporary kernel telemetry probe tests Signed-off-by: Jay Xiao --- tests/unit/test_telemetry.py | 1065 ------------------------ tests/unit/test_zz_kernel_sig_probe.py | 1 - 2 files changed, 1066 deletions(-) delete mode 100644 tests/unit/test_telemetry.py delete mode 100644 tests/unit/test_zz_kernel_sig_probe.py diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py deleted file mode 100644 index abcdd6767..000000000 --- a/tests/unit/test_telemetry.py +++ /dev/null @@ -1,1065 +0,0 @@ -import uuid -import pytest -from unittest.mock import patch, MagicMock -import json -from dataclasses import asdict - -from databricks.sql.telemetry.telemetry_client import ( - TelemetryClient, - NoopTelemetryClient, - TelemetryClientFactory, - TelemetryHelper, -) -from databricks.sql.common.feature_flag import ( - FeatureFlagsContextFactory, - FeatureFlagsContext, -) -from databricks.sql.telemetry.models.enums import AuthMech, AuthFlow, DatabricksClientType -from databricks.sql.telemetry.models.event import ( - TelemetryEvent, - DriverConnectionParameters, - DriverSystemConfiguration, - SqlExecutionEvent, - DriverErrorInfo, - DriverVolumeOperation, - HostDetails, -) -from databricks.sql.auth.authenticators import ( - AccessTokenAuthProvider, - DatabricksOAuthProvider, - ExternalAuthProvider, -) -from databricks.sql.auth.token_federation import TokenFederationProvider -from databricks import sql - - -@pytest.fixture -def mock_telemetry_client(): - """Create a mock telemetry client for testing.""" - session_id = str(uuid.uuid4()) - auth_provider = AccessTokenAuthProvider("test-token") - executor = MagicMock() - client_context = MagicMock() - - # Patch the _setup_pool_manager method to avoid SSL file loading - with patch( - "databricks.sql.common.unified_http_client.UnifiedHttpClient._setup_pool_managers" - ): - return TelemetryClient( - telemetry_enabled=True, - session_id_hex=session_id, - auth_provider=auth_provider, - host_url="test-host.com", - executor=executor, - batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE, - client_context=client_context, - ) - - -class TestNoopTelemetryClient: - """Tests for NoopTelemetryClient - should do nothing safely.""" - - def test_noop_client_behavior(self): - """Test that NoopTelemetryClient is a singleton and all methods are safe no-ops.""" - # Test singleton behavior - client1 = NoopTelemetryClient() - client2 = NoopTelemetryClient() - assert client1 is client2 - - # Test that all methods can be called without exceptions - client1.export_initial_telemetry_log(MagicMock(), "test-agent") - client1.export_failure_log("TestError", "Test message") - client1.export_latency_log(100, "EXECUTE_STATEMENT", "test-id") - client1.close() - - -class TestTelemetryClient: - """Tests for actual telemetry client functionality and flows.""" - - def test_event_batching_and_flushing_flow(self, mock_telemetry_client): - """Test the complete event batching and flushing flow.""" - client = mock_telemetry_client - client._batch_size = 3 # Small batch for testing - - # Mock the network call - with patch.object(client, "_send_telemetry") as mock_send: - # Add events one by one - should not flush yet - client._export_event("event1") - client._export_event("event2") - mock_send.assert_not_called() - assert client._events_queue.qsize() == 2 - - # Third event should trigger flush - client._export_event("event3") - mock_send.assert_called_once() - assert client._events_queue.qsize() == 0 # Queue cleared after flush - - @patch("databricks.sql.common.unified_http_client.UnifiedHttpClient.request") - def test_network_request_flow(self, mock_http_request, mock_telemetry_client): - """Test the complete network request flow with authentication.""" - # Mock response for unified HTTP client - mock_response = MagicMock() - mock_response.status = 200 - mock_response.status_code = 200 - mock_http_request.return_value = mock_response - - client = mock_telemetry_client - - # Create mock events - mock_events = [MagicMock() for _ in range(2)] - for i, event in enumerate(mock_events): - event.to_json.return_value = f'{{"event": "{i}"}}' - - # Send telemetry - client._send_telemetry(mock_events) - - # Verify request was submitted to executor - client._executor.submit.assert_called_once() - args, kwargs = client._executor.submit.call_args - - # Verify correct function and URL - assert args[0] == client._send_with_unified_client - assert args[1] == "https://test-host.com/telemetry-ext" - assert kwargs["headers"]["Authorization"] == "Bearer test-token" - - # Verify request body structure - request_data = kwargs["data"] - assert '"uploadTime"' in request_data - assert '"protoLogs"' in request_data - - def test_telemetry_logging_flows(self, mock_telemetry_client): - """Test all telemetry logging methods work end-to-end.""" - client = mock_telemetry_client - - with patch.object(client, "_export_event") as mock_export: - # Test initial log - client.export_initial_telemetry_log(MagicMock(), "test-agent") - assert mock_export.call_count == 1 - - # Test failure log - client.export_failure_log("TestError", "Error message") - assert mock_export.call_count == 2 - - # Test latency log - client.export_latency_log(150, "EXECUTE_STATEMENT", "stmt-123") - assert mock_export.call_count == 3 - - def test_error_handling_resilience(self, mock_telemetry_client): - """Test that telemetry errors don't break the client.""" - client = mock_telemetry_client - - # Test that exceptions in telemetry don't propagate - with patch.object(client, "_export_event", side_effect=Exception("Test error")): - # These should not raise exceptions - client.export_initial_telemetry_log(MagicMock(), "test-agent") - client.export_failure_log("TestError", "Error message") - client.export_latency_log(100, "EXECUTE_STATEMENT", "stmt-123") - - # Test executor submission failure - client._executor.submit.side_effect = Exception("Thread pool error") - client._send_telemetry([MagicMock()]) # Should not raise - - -class TestTelemetryHelper: - """Tests for TelemetryHelper utility functions.""" - - def test_system_configuration_caching(self): - """Test that system configuration is cached and contains expected data.""" - config1 = TelemetryHelper.get_driver_system_configuration() - config2 = TelemetryHelper.get_driver_system_configuration() - - # Should be cached (same instance) - assert config1 is config2 - - def test_auth_mechanism_detection(self): - """Test authentication mechanism detection for different providers.""" - test_cases = [ - (AccessTokenAuthProvider("token"), AuthMech.PAT), - (MagicMock(spec=DatabricksOAuthProvider), AuthMech.OAUTH), - (MagicMock(spec=ExternalAuthProvider), AuthMech.OTHER), - (MagicMock(), AuthMech.OTHER), # Unknown provider - (None, None), - ] - - for provider, expected in test_cases: - assert TelemetryHelper.get_auth_mechanism(provider) == expected - - def test_auth_flow_detection(self): - """Test authentication flow detection for OAuth providers.""" - # OAuth with existing tokens - oauth_with_tokens = MagicMock(spec=DatabricksOAuthProvider) - oauth_with_tokens._access_token = "test-access-token" - oauth_with_tokens._refresh_token = "test-refresh-token" - assert ( - TelemetryHelper.get_auth_flow(oauth_with_tokens) - == AuthFlow.TOKEN_PASSTHROUGH - ) - - # Test OAuth with browser-based auth - oauth_with_browser = MagicMock(spec=DatabricksOAuthProvider) - oauth_with_browser._access_token = None - oauth_with_browser._refresh_token = None - oauth_with_browser.oauth_manager = MagicMock() - assert ( - TelemetryHelper.get_auth_flow(oauth_with_browser) - == AuthFlow.BROWSER_BASED_AUTHENTICATION - ) - - # Test non-OAuth provider - pat_auth = AccessTokenAuthProvider("test-token") - assert TelemetryHelper.get_auth_flow(pat_auth) is None - - # Test None auth provider - assert TelemetryHelper.get_auth_flow(None) is None - - def _make_real_federation(self, inner): - """Build a real TokenFederationProvider so attribute renames break tests.""" - return TokenFederationProvider( - hostname="example.cloud.databricks.com", - external_provider=inner, - http_client=MagicMock(), - ) - - def test_token_federation_unwraps_pat(self): - fed = self._make_real_federation(AccessTokenAuthProvider("test-token")) - assert TelemetryHelper.get_auth_mechanism(fed) == AuthMech.PAT - assert TelemetryHelper.get_auth_flow(fed) is None - - def test_token_federation_unwraps_m2m(self): - fed = self._make_real_federation(MagicMock(spec=ExternalAuthProvider)) - assert TelemetryHelper.get_auth_mechanism(fed) == AuthMech.OTHER - assert TelemetryHelper.get_auth_flow(fed) == AuthFlow.CLIENT_CREDENTIALS - - def test_token_federation_unwraps_oauth_browser(self): - oauth = MagicMock(spec=DatabricksOAuthProvider) - oauth._access_token = None - oauth._refresh_token = None - fed = self._make_real_federation(oauth) - assert TelemetryHelper.get_auth_mechanism(fed) == AuthMech.OAUTH - assert TelemetryHelper.get_auth_flow(fed) == AuthFlow.BROWSER_BASED_AUTHENTICATION - - def test_token_federation_unwraps_oauth_passthrough(self): - oauth = MagicMock(spec=DatabricksOAuthProvider) - oauth._access_token = "a" - oauth._refresh_token = "r" - fed = self._make_real_federation(oauth) - assert TelemetryHelper.get_auth_mechanism(fed) == AuthMech.OAUTH - assert TelemetryHelper.get_auth_flow(fed) == AuthFlow.TOKEN_PASSTHROUGH - - def test_token_federation_payload_serialization(self): - """End-to-end: federated PAT must serialize as PAT in the connection-params payload.""" - fed = self._make_real_federation(AccessTokenAuthProvider("test-token")) - params = DriverConnectionParameters( - http_path="/sql/1.0/warehouses/abc", - mode=DatabricksClientType.THRIFT, - host_info=HostDetails(host_url="https://example.cloud.databricks.com", port=443), - auth_mech=TelemetryHelper.get_auth_mechanism(fed), - auth_flow=TelemetryHelper.get_auth_flow(fed), - ) - payload = json.loads(params.to_json()) - assert payload["auth_mech"] == "PAT" - assert "auth_flow" not in payload # None-valued fields are stripped - - def test_token_federation_with_no_inner_provider(self): - """Federation with a None inner provider should not crash; both helpers return None.""" - fed = TokenFederationProvider( - hostname="example.cloud.databricks.com", - external_provider=None, - http_client=MagicMock(), - ) - 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_defaults_false(self): - kernel_kwargs = self._kernel_telemetry_kwargs_for_test({}) - - assert kernel_kwargs["telemetry_enabled"] is False - - -class TestTelemetryFactory: - """Tests for TelemetryClientFactory lifecycle and management.""" - - @pytest.fixture(autouse=True) - def telemetry_system_reset(self): - """Reset telemetry system state before each test.""" - TelemetryClientFactory._clients.clear() - if TelemetryClientFactory._executor: - TelemetryClientFactory._executor.shutdown(wait=True) - TelemetryClientFactory._executor = None - TelemetryClientFactory._initialized = False - yield - TelemetryClientFactory._clients.clear() - if TelemetryClientFactory._executor: - TelemetryClientFactory._executor.shutdown(wait=True) - TelemetryClientFactory._executor = None - TelemetryClientFactory._initialized = False - - def test_client_lifecycle_flow(self): - """Test complete client lifecycle: initialize -> use -> close.""" - session_id_hex = "test-session" - auth_provider = AccessTokenAuthProvider("token") - client_context = MagicMock() - - # Initialize enabled client - with patch( - "databricks.sql.common.unified_http_client.UnifiedHttpClient._setup_pool_managers" - ): - TelemetryClientFactory.initialize_telemetry_client( - telemetry_enabled=True, - session_id_hex=session_id_hex, - auth_provider=auth_provider, - host_url="test-host.com", - batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE, - client_context=client_context, - ) - - client = TelemetryClientFactory.get_telemetry_client("test-host.com") - assert isinstance(client, TelemetryClient) - assert client._session_id_hex == session_id_hex - - # Close client - with patch.object(client, "close") as mock_close: - TelemetryClientFactory.close(host_url="test-host.com") - mock_close.assert_called_once() - - # Should get NoopTelemetryClient after close - - def test_disabled_telemetry_creates_noop_client(self): - """Test that disabled telemetry creates NoopTelemetryClient.""" - session_id_hex = "test-session" - client_context = MagicMock() - - TelemetryClientFactory.initialize_telemetry_client( - telemetry_enabled=False, - session_id_hex=session_id_hex, - auth_provider=None, - host_url="test-host.com", - batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE, - client_context=client_context, - ) - - client = TelemetryClientFactory.get_telemetry_client("test-host.com") - assert isinstance(client, NoopTelemetryClient) - - def test_factory_error_handling(self): - """Test that factory errors fall back to NoopTelemetryClient.""" - session_id = "test-session" - client_context = MagicMock() - - # Simulate initialization error - with patch( - "databricks.sql.telemetry.telemetry_client.TelemetryClient", - side_effect=Exception("Init error"), - ): - TelemetryClientFactory.initialize_telemetry_client( - telemetry_enabled=True, - session_id_hex=session_id, - auth_provider=AccessTokenAuthProvider("token"), - host_url="test-host.com", - batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE, - client_context=client_context, - ) - - # Should fall back to NoopTelemetryClient - client = TelemetryClientFactory.get_telemetry_client("test-host.com") - assert isinstance(client, NoopTelemetryClient) - - def test_factory_shutdown_flow(self): - """Test factory shutdown when last client is removed.""" - session1 = "session-1" - session2 = "session-2" - client_context = MagicMock() - - # Initialize multiple clients - with patch( - "databricks.sql.common.unified_http_client.UnifiedHttpClient._setup_pool_managers" - ): - for session in [session1, session2]: - TelemetryClientFactory.initialize_telemetry_client( - telemetry_enabled=True, - session_id_hex=session, - auth_provider=AccessTokenAuthProvider("token"), - host_url="test-host.com", - batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE, - client_context=client_context, - ) - - # Factory should be initialized - assert TelemetryClientFactory._initialized is True - assert TelemetryClientFactory._executor is not None - - # Close first client - factory should stay initialized - TelemetryClientFactory.close(host_url="test-host.com") - assert TelemetryClientFactory._initialized is True - - # Close second client - factory should shut down - TelemetryClientFactory.close(host_url="test-host.com") - assert TelemetryClientFactory._initialized is False - assert TelemetryClientFactory._executor is None - - @patch( - "databricks.sql.telemetry.telemetry_client.TelemetryClient.export_failure_log" - ) - @patch("databricks.sql.client.Session") - def test_connection_failure_sends_correct_telemetry_payload( - self, mock_session, mock_export_failure_log - ): - """ - Verify that a connection failure constructs and sends the correct - telemetry payload via _send_telemetry. - """ - - error_message = "Could not connect to host" - # Set up the mock to create a session instance first, then make open() fail - mock_session_instance = MagicMock() - mock_session_instance.is_open = False # Ensure cleanup is safe - 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") - except Exception as e: - assert str(e) == error_message - - mock_export_failure_log.assert_called_once() - call_arguments = mock_export_failure_log.call_args - assert call_arguments[0][0] == "Exception" - assert call_arguments[0][1] == error_message - - -@patch("databricks.sql.client.Session") -class TestTelemetryFeatureFlag: - """Tests the interaction between the telemetry feature flag and connection parameters.""" - - def teardown_method(self): - """Clean up telemetry factory state after each test to prevent test pollution.""" - from databricks.sql.common.feature_flag import FeatureFlagsContextFactory - - TelemetryClientFactory._clients.clear() - FeatureFlagsContextFactory._context_map.clear() - - def _mock_ff_response(self, mock_http_request, enabled: bool): - """Helper method to mock feature flag response for unified HTTP client.""" - mock_response = MagicMock() - mock_response.status = 200 - mock_response.status_code = 200 # Compatibility attribute - payload = { - "flags": [ - { - "name": "databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForPythonDriver", - "value": str(enabled).lower(), - } - ], - "ttl_seconds": 3600, - } - mock_response.json.return_value = payload - mock_response.data = json.dumps(payload).encode() - mock_http_request.return_value = mock_response - - @patch("databricks.sql.common.unified_http_client.UnifiedHttpClient.request") - def test_telemetry_enabled_when_flag_is_true(self, mock_http_request, MockSession): - """Telemetry should be ON when enable_telemetry=True and server flag is 'true'.""" - 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 = ( - False # Connection starts closed for test cleanup - ) - - # Set up mock HTTP client on the session - mock_http_client = MagicMock() - mock_http_client.request = mock_http_request - mock_session_instance.http_client = mock_http_client - - conn = sql.client.Connection( - server_hostname="test", - http_path="test", - access_token="test", - enable_telemetry=True, - ) - - assert conn.telemetry_enabled is True - mock_http_request.assert_called_once() - client = TelemetryClientFactory.get_telemetry_client("test-host") - assert isinstance(client, TelemetryClient) - - @patch("databricks.sql.common.unified_http_client.UnifiedHttpClient.request") - def test_telemetry_disabled_when_flag_is_false( - self, mock_http_request, MockSession - ): - """Telemetry should be OFF when enable_telemetry=True but server 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 = ( - False # Connection starts closed for test cleanup - ) - - # Set up mock HTTP client on the session - mock_http_client = MagicMock() - mock_http_client.request = mock_http_request - mock_session_instance.http_client = mock_http_client - - conn = sql.client.Connection( - server_hostname="test", - http_path="test", - access_token="test", - enable_telemetry=True, - ) - - assert conn.telemetry_enabled is False - mock_http_request.assert_called_once() - client = TelemetryClientFactory.get_telemetry_client("test-host") - assert isinstance(client, NoopTelemetryClient) - - @patch("databricks.sql.common.unified_http_client.UnifiedHttpClient.request") - def test_telemetry_disabled_when_flag_request_fails( - self, mock_http_request, MockSession - ): - """Telemetry should default to OFF if the feature flag network 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 = ( - False # Connection starts closed for test cleanup - ) - - # Set up mock HTTP client on the session - mock_http_client = MagicMock() - mock_http_client.request = mock_http_request - mock_session_instance.http_client = mock_http_client - - conn = sql.client.Connection( - server_hostname="test", - http_path="test", - access_token="test", - enable_telemetry=True, - ) - - assert conn.telemetry_enabled is False - mock_http_request.assert_called_once() - client = TelemetryClientFactory.get_telemetry_client("test-host") - assert isinstance(client, NoopTelemetryClient) - - -class TestTelemetryEventModels: - """Tests for telemetry event model data structures and JSON serialization.""" - - def test_host_details_serialization(self): - """Test HostDetails model serialization.""" - host = HostDetails(host_url="test-host.com", port=443) - - # Test JSON string generation - json_str = host.to_json() - assert isinstance(json_str, str) - parsed = json.loads(json_str) - assert parsed["host_url"] == "test-host.com" - assert parsed["port"] == 443 - - def test_driver_connection_parameters_all_fields(self): - """Test DriverConnectionParameters with all fields populated.""" - host_info = HostDetails(host_url="workspace.databricks.com", port=443) - proxy_info = HostDetails(host_url="proxy.company.com", port=8080) - cf_proxy_info = HostDetails(host_url="cf-proxy.company.com", port=8080) - - params = DriverConnectionParameters( - http_path="/sql/1.0/warehouses/abc123", - mode=DatabricksClientType.SEA, - host_info=host_info, - auth_mech=AuthMech.OAUTH, - auth_flow=AuthFlow.BROWSER_BASED_AUTHENTICATION, - socket_timeout=30000, - azure_workspace_resource_id="/subscriptions/test/resourceGroups/test", - azure_tenant_id="tenant-123", - use_proxy=True, - use_system_proxy=True, - proxy_host_info=proxy_info, - use_cf_proxy=False, - cf_proxy_host_info=cf_proxy_info, - non_proxy_hosts=["localhost", "127.0.0.1"], - allow_self_signed_support=False, - use_system_trust_store=True, - enable_arrow=True, - enable_direct_results=True, - enable_sea_hybrid_results=True, - http_connection_pool_size=100, - rows_fetched_per_block=100000, - async_poll_interval_millis=2000, - support_many_parameters=True, - enable_complex_datatype_support=True, - allowed_volume_ingestion_paths="/Volumes/catalog/schema/volume", - query_tags="team:engineering,project:telemetry", - ) - - # Serialize to JSON and parse back - json_str = params.to_json() - json_dict = json.loads(json_str) - - # Verify all new fields are in JSON - assert json_dict["http_path"] == "/sql/1.0/warehouses/abc123" - assert json_dict["mode"] == "SEA" - assert json_dict["host_info"]["host_url"] == "workspace.databricks.com" - assert json_dict["auth_mech"] == "OAUTH" - assert json_dict["auth_flow"] == "BROWSER_BASED_AUTHENTICATION" - assert json_dict["socket_timeout"] == 30000 - assert json_dict["azure_workspace_resource_id"] == "/subscriptions/test/resourceGroups/test" - assert json_dict["azure_tenant_id"] == "tenant-123" - assert json_dict["use_proxy"] is True - assert json_dict["use_system_proxy"] is True - assert json_dict["proxy_host_info"]["host_url"] == "proxy.company.com" - assert json_dict["use_cf_proxy"] is False - assert json_dict["cf_proxy_host_info"]["host_url"] == "cf-proxy.company.com" - assert json_dict["non_proxy_hosts"] == ["localhost", "127.0.0.1"] - assert json_dict["allow_self_signed_support"] is False - assert json_dict["use_system_trust_store"] is True - assert json_dict["enable_arrow"] is True - assert json_dict["enable_direct_results"] is True - assert json_dict["enable_sea_hybrid_results"] is True - assert json_dict["http_connection_pool_size"] == 100 - assert json_dict["rows_fetched_per_block"] == 100000 - assert json_dict["async_poll_interval_millis"] == 2000 - assert json_dict["support_many_parameters"] is True - assert json_dict["enable_complex_datatype_support"] is True - assert json_dict["allowed_volume_ingestion_paths"] == "/Volumes/catalog/schema/volume" - assert json_dict["query_tags"] == "team:engineering,project:telemetry" - - def test_driver_connection_parameters_minimal_fields(self): - """Test DriverConnectionParameters with only required fields.""" - host_info = HostDetails(host_url="workspace.databricks.com", port=443) - - params = DriverConnectionParameters( - http_path="/sql/1.0/warehouses/abc123", - mode=DatabricksClientType.THRIFT, - host_info=host_info, - ) - - # Note: to_json() filters out None values, so we need to check asdict for complete structure - json_str = params.to_json() - json_dict = json.loads(json_str) - - # Required fields should be present - assert json_dict["http_path"] == "/sql/1.0/warehouses/abc123" - assert json_dict["mode"] == "THRIFT" - assert json_dict["host_info"]["host_url"] == "workspace.databricks.com" - - # Optional fields with None are filtered out by to_json() - # This is expected behavior - None values are excluded from JSON output - - def test_driver_system_configuration_serialization(self): - """Test DriverSystemConfiguration model serialization.""" - sys_config = DriverSystemConfiguration( - driver_name="Databricks SQL Connector for Python", - driver_version="3.0.0", - runtime_name="CPython", - runtime_version="3.11.0", - runtime_vendor="Python Software Foundation", - os_name="Darwin", - os_version="23.0.0", - os_arch="arm64", - char_set_encoding="utf-8", - locale_name="en_US", - client_app_name="MyApp", - ) - - json_str = sys_config.to_json() - json_dict = json.loads(json_str) - - assert json_dict["driver_name"] == "Databricks SQL Connector for Python" - assert json_dict["driver_version"] == "3.0.0" - assert json_dict["runtime_name"] == "CPython" - assert json_dict["runtime_version"] == "3.11.0" - assert json_dict["runtime_vendor"] == "Python Software Foundation" - assert json_dict["os_name"] == "Darwin" - assert json_dict["os_version"] == "23.0.0" - assert json_dict["os_arch"] == "arm64" - assert json_dict["locale_name"] == "en_US" - assert json_dict["char_set_encoding"] == "utf-8" - assert json_dict["client_app_name"] == "MyApp" - - def test_telemetry_event_complete_serialization(self): - """Test complete TelemetryEvent serialization with all nested objects.""" - host_info = HostDetails(host_url="workspace.databricks.com", port=443) - proxy_info = HostDetails(host_url="proxy.company.com", port=8080) - - connection_params = DriverConnectionParameters( - http_path="/sql/1.0/warehouses/abc123", - mode=DatabricksClientType.SEA, - host_info=host_info, - auth_mech=AuthMech.OAUTH, - use_proxy=True, - proxy_host_info=proxy_info, - enable_arrow=True, - rows_fetched_per_block=100000, - ) - - sys_config = DriverSystemConfiguration( - driver_name="Databricks SQL Connector for Python", - driver_version="3.0.0", - runtime_name="CPython", - runtime_version="3.11.0", - runtime_vendor="Python Software Foundation", - os_name="Darwin", - os_version="23.0.0", - os_arch="arm64", - char_set_encoding="utf-8", - ) - - error_info = DriverErrorInfo( - error_name="ConnectionError", - stack_trace="Traceback...", - ) - - event = TelemetryEvent( - session_id="test-session-123", - sql_statement_id="test-stmt-456", - operation_latency_ms=1500, - auth_type="OAUTH", - system_configuration=sys_config, - driver_connection_params=connection_params, - error_info=error_info, - ) - - # Test JSON serialization - json_str = event.to_json() - assert isinstance(json_str, str) - - # Parse and verify structure - parsed = json.loads(json_str) - assert parsed["session_id"] == "test-session-123" - assert parsed["sql_statement_id"] == "test-stmt-456" - assert parsed["operation_latency_ms"] == 1500 - assert parsed["auth_type"] == "OAUTH" - - # Verify nested objects - assert parsed["system_configuration"]["driver_name"] == "Databricks SQL Connector for Python" - assert parsed["driver_connection_params"]["http_path"] == "/sql/1.0/warehouses/abc123" - assert parsed["driver_connection_params"]["use_proxy"] is True - assert parsed["driver_connection_params"]["proxy_host_info"]["host_url"] == "proxy.company.com" - assert parsed["error_info"]["error_name"] == "ConnectionError" - - def test_json_serialization_excludes_none_values(self): - """Test that JSON serialization properly excludes None values.""" - host_info = HostDetails(host_url="workspace.databricks.com", port=443) - - params = DriverConnectionParameters( - http_path="/sql/1.0/warehouses/abc123", - mode=DatabricksClientType.SEA, - host_info=host_info, - # All optional fields left as None - ) - - json_str = params.to_json() - parsed = json.loads(json_str) - - # Required fields present - assert parsed["http_path"] == "/sql/1.0/warehouses/abc123" - - # None values should be EXCLUDED from JSON (not included as null) - # This is the behavior of JsonSerializableMixin - assert "auth_mech" not in parsed - assert "azure_tenant_id" not in parsed - assert "proxy_host_info" not in parsed - - -@patch("databricks.sql.client.Session") -@patch("databricks.sql.common.unified_http_client.UnifiedHttpClient._setup_pool_managers") -class TestConnectionParameterTelemetry: - """Tests for connection parameter population in telemetry.""" - - def test_connection_with_proxy_populates_telemetry(self, mock_setup_pools, mock_session): - """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 - mock_session_instance.port = 443 - mock_session_instance.host = "workspace.databricks.com" - mock_session.return_value = mock_session_instance - - with patch("databricks.sql.telemetry.telemetry_client.TelemetryClient.export_initial_telemetry_log") as mock_export: - conn = sql.connect( - server_hostname="workspace.databricks.com", - http_path="/sql/1.0/warehouses/test", - access_token="test-token", - enable_telemetry=True, - force_enable_telemetry=True, - ) - - # Verify export was called - mock_export.assert_called_once() - call_args = mock_export.call_args - - # Extract driver_connection_params - driver_params = call_args.kwargs.get("driver_connection_params") - assert driver_params is not None - assert isinstance(driver_params, DriverConnectionParameters) - - # Verify fields are populated - assert driver_params.http_path == "/sql/1.0/warehouses/test" - assert driver_params.mode == DatabricksClientType.SEA - assert driver_params.host_info.host_url == "workspace.databricks.com" - assert driver_params.host_info.port == 443 - - def test_connection_with_azure_params_populates_telemetry(self, mock_setup_pools, mock_session): - """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 - mock_session_instance.port = 443 - mock_session_instance.host = "workspace.azuredatabricks.net" - mock_session.return_value = mock_session_instance - - with patch("databricks.sql.telemetry.telemetry_client.TelemetryClient.export_initial_telemetry_log") as mock_export: - conn = sql.connect( - server_hostname="workspace.azuredatabricks.net", - http_path="/sql/1.0/warehouses/test", - access_token="test-token", - azure_workspace_resource_id="/subscriptions/test/resourceGroups/test", - azure_tenant_id="tenant-123", - enable_telemetry=True, - force_enable_telemetry=True, - ) - - mock_export.assert_called_once() - driver_params = mock_export.call_args.kwargs.get("driver_connection_params") - - # Verify Azure fields - assert driver_params.azure_workspace_resource_id == "/subscriptions/test/resourceGroups/test" - assert driver_params.azure_tenant_id == "tenant-123" - - def test_connection_populates_arrow_and_performance_params(self, mock_setup_pools, mock_session): - """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 - mock_session_instance.port = 443 - mock_session_instance.host = "workspace.databricks.com" - mock_session.return_value = mock_session_instance - - with patch("databricks.sql.telemetry.telemetry_client.TelemetryClient.export_initial_telemetry_log") as mock_export: - # Import pyarrow availability check - try: - import pyarrow - arrow_available = True - except ImportError: - arrow_available = False - - conn = sql.connect( - server_hostname="workspace.databricks.com", - http_path="/sql/1.0/warehouses/test", - access_token="test-token", - pool_maxsize=200, - enable_telemetry=True, - force_enable_telemetry=True, - ) - - mock_export.assert_called_once() - driver_params = mock_export.call_args.kwargs.get("driver_connection_params") - - # Verify performance fields - assert driver_params.enable_arrow == arrow_available - assert driver_params.enable_direct_results is True - assert driver_params.http_connection_pool_size == 200 - assert driver_params.rows_fetched_per_block == 100000 # DEFAULT_ARRAY_SIZE - assert driver_params.async_poll_interval_millis == 2000 - assert driver_params.support_many_parameters is True - - def test_federated_pat_populates_telemetry_as_pat(self, mock_setup_pools, mock_session): - """End-to-end: a TokenFederationProvider wrapping a PAT should report mech=PAT in the captured telemetry payload.""" - federated_pat = TokenFederationProvider( - hostname="workspace.databricks.com", - external_provider=AccessTokenAuthProvider("token"), - http_client=MagicMock(), - ) - 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 - mock_session_instance.port = 443 - mock_session_instance.host = "workspace.databricks.com" - mock_session.return_value = mock_session_instance - - with patch( - "databricks.sql.telemetry.telemetry_client.TelemetryClient.export_initial_telemetry_log" - ) as mock_export: - sql.connect( - server_hostname="workspace.databricks.com", - http_path="/sql/1.0/warehouses/test", - access_token="test-token", - enable_telemetry=True, - force_enable_telemetry=True, - ) - - mock_export.assert_called_once() - driver_params = mock_export.call_args.kwargs.get("driver_connection_params") - assert driver_params.auth_mech == AuthMech.PAT - assert driver_params.auth_flow is None - - def test_cf_proxy_fields_default_to_false_none(self, mock_setup_pools, mock_session): - """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 - mock_session_instance.port = 443 - mock_session_instance.host = "workspace.databricks.com" - mock_session.return_value = mock_session_instance - - with patch("databricks.sql.telemetry.telemetry_client.TelemetryClient.export_initial_telemetry_log") as mock_export: - conn = sql.connect( - server_hostname="workspace.databricks.com", - http_path="/sql/1.0/warehouses/test", - access_token="test-token", - enable_telemetry=True, - force_enable_telemetry=True, - ) - - mock_export.assert_called_once() - driver_params = mock_export.call_args.kwargs.get("driver_connection_params") - - # CF proxy not yet supported - should be False/None - assert driver_params.use_cf_proxy is False - assert driver_params.cf_proxy_host_info is None - - -class TestFeatureFlagsContextFactory: - """Tests for FeatureFlagsContextFactory host-level caching.""" - - @pytest.fixture(autouse=True) - def reset_factory(self): - """Reset factory state before/after each test.""" - FeatureFlagsContextFactory._context_map.clear() - if FeatureFlagsContextFactory._executor: - FeatureFlagsContextFactory._executor.shutdown(wait=False) - FeatureFlagsContextFactory._executor = None - yield - FeatureFlagsContextFactory._context_map.clear() - if FeatureFlagsContextFactory._executor: - FeatureFlagsContextFactory._executor.shutdown(wait=False) - FeatureFlagsContextFactory._executor = None - - @pytest.mark.parametrize( - "hosts,expected_contexts", - [ - (["host1.com", "host1.com"], 1), # Same host shares context - (["host1.com", "host2.com"], 2), # Different hosts get separate contexts - (["host1.com", "host1.com", "host2.com"], 2), # Mixed scenario - ], - ) - def test_host_level_caching(self, hosts, expected_contexts): - """Test that contexts are cached by host correctly.""" - contexts = [] - for host in hosts: - conn = MagicMock() - conn.session.host = host - conn.session.http_client = MagicMock() - contexts.append(FeatureFlagsContextFactory.get_instance(conn)) - - assert len(FeatureFlagsContextFactory._context_map) == expected_contexts - if expected_contexts == 1: - assert all(ctx is contexts[0] for ctx in contexts) - - def test_remove_instance_and_executor_cleanup(self): - """Test removal uses host key and cleans up executor when empty.""" - conn1 = MagicMock() - conn1.session.host = "host1.com" - conn1.session.http_client = MagicMock() - - conn2 = MagicMock() - conn2.session.host = "host2.com" - conn2.session.http_client = MagicMock() - - FeatureFlagsContextFactory.get_instance(conn1) - FeatureFlagsContextFactory.get_instance(conn2) - assert FeatureFlagsContextFactory._executor is not None - - FeatureFlagsContextFactory.remove_instance(conn1) - assert len(FeatureFlagsContextFactory._context_map) == 1 - assert FeatureFlagsContextFactory._executor is not None - - FeatureFlagsContextFactory.remove_instance(conn2) - assert len(FeatureFlagsContextFactory._context_map) == 0 - assert FeatureFlagsContextFactory._executor is None diff --git a/tests/unit/test_zz_kernel_sig_probe.py b/tests/unit/test_zz_kernel_sig_probe.py deleted file mode 100644 index 80f73fb4e..000000000 --- a/tests/unit/test_zz_kernel_sig_probe.py +++ /dev/null @@ -1 +0,0 @@ -# (temporary probe removed) From ffc76ee6dce516192612b171a6f588b3380da2ca Mon Sep 17 00:00:00 2001 From: Jay Xiao Date: Fri, 28 Aug 2026 00:10:50 +0000 Subject: [PATCH 08/17] Correct temporary probe test cleanup Signed-off-by: Jay Xiao --- tests/unit/test_kernel_sig_probe.py | 0 tests/unit/test_telemetry.py | 1065 +++++++++++++++++++++++++++ 2 files changed, 1065 insertions(+) delete mode 100644 tests/unit/test_kernel_sig_probe.py create mode 100644 tests/unit/test_telemetry.py diff --git a/tests/unit/test_kernel_sig_probe.py b/tests/unit/test_kernel_sig_probe.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py new file mode 100644 index 000000000..abcdd6767 --- /dev/null +++ b/tests/unit/test_telemetry.py @@ -0,0 +1,1065 @@ +import uuid +import pytest +from unittest.mock import patch, MagicMock +import json +from dataclasses import asdict + +from databricks.sql.telemetry.telemetry_client import ( + TelemetryClient, + NoopTelemetryClient, + TelemetryClientFactory, + TelemetryHelper, +) +from databricks.sql.common.feature_flag import ( + FeatureFlagsContextFactory, + FeatureFlagsContext, +) +from databricks.sql.telemetry.models.enums import AuthMech, AuthFlow, DatabricksClientType +from databricks.sql.telemetry.models.event import ( + TelemetryEvent, + DriverConnectionParameters, + DriverSystemConfiguration, + SqlExecutionEvent, + DriverErrorInfo, + DriverVolumeOperation, + HostDetails, +) +from databricks.sql.auth.authenticators import ( + AccessTokenAuthProvider, + DatabricksOAuthProvider, + ExternalAuthProvider, +) +from databricks.sql.auth.token_federation import TokenFederationProvider +from databricks import sql + + +@pytest.fixture +def mock_telemetry_client(): + """Create a mock telemetry client for testing.""" + session_id = str(uuid.uuid4()) + auth_provider = AccessTokenAuthProvider("test-token") + executor = MagicMock() + client_context = MagicMock() + + # Patch the _setup_pool_manager method to avoid SSL file loading + with patch( + "databricks.sql.common.unified_http_client.UnifiedHttpClient._setup_pool_managers" + ): + return TelemetryClient( + telemetry_enabled=True, + session_id_hex=session_id, + auth_provider=auth_provider, + host_url="test-host.com", + executor=executor, + batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE, + client_context=client_context, + ) + + +class TestNoopTelemetryClient: + """Tests for NoopTelemetryClient - should do nothing safely.""" + + def test_noop_client_behavior(self): + """Test that NoopTelemetryClient is a singleton and all methods are safe no-ops.""" + # Test singleton behavior + client1 = NoopTelemetryClient() + client2 = NoopTelemetryClient() + assert client1 is client2 + + # Test that all methods can be called without exceptions + client1.export_initial_telemetry_log(MagicMock(), "test-agent") + client1.export_failure_log("TestError", "Test message") + client1.export_latency_log(100, "EXECUTE_STATEMENT", "test-id") + client1.close() + + +class TestTelemetryClient: + """Tests for actual telemetry client functionality and flows.""" + + def test_event_batching_and_flushing_flow(self, mock_telemetry_client): + """Test the complete event batching and flushing flow.""" + client = mock_telemetry_client + client._batch_size = 3 # Small batch for testing + + # Mock the network call + with patch.object(client, "_send_telemetry") as mock_send: + # Add events one by one - should not flush yet + client._export_event("event1") + client._export_event("event2") + mock_send.assert_not_called() + assert client._events_queue.qsize() == 2 + + # Third event should trigger flush + client._export_event("event3") + mock_send.assert_called_once() + assert client._events_queue.qsize() == 0 # Queue cleared after flush + + @patch("databricks.sql.common.unified_http_client.UnifiedHttpClient.request") + def test_network_request_flow(self, mock_http_request, mock_telemetry_client): + """Test the complete network request flow with authentication.""" + # Mock response for unified HTTP client + mock_response = MagicMock() + mock_response.status = 200 + mock_response.status_code = 200 + mock_http_request.return_value = mock_response + + client = mock_telemetry_client + + # Create mock events + mock_events = [MagicMock() for _ in range(2)] + for i, event in enumerate(mock_events): + event.to_json.return_value = f'{{"event": "{i}"}}' + + # Send telemetry + client._send_telemetry(mock_events) + + # Verify request was submitted to executor + client._executor.submit.assert_called_once() + args, kwargs = client._executor.submit.call_args + + # Verify correct function and URL + assert args[0] == client._send_with_unified_client + assert args[1] == "https://test-host.com/telemetry-ext" + assert kwargs["headers"]["Authorization"] == "Bearer test-token" + + # Verify request body structure + request_data = kwargs["data"] + assert '"uploadTime"' in request_data + assert '"protoLogs"' in request_data + + def test_telemetry_logging_flows(self, mock_telemetry_client): + """Test all telemetry logging methods work end-to-end.""" + client = mock_telemetry_client + + with patch.object(client, "_export_event") as mock_export: + # Test initial log + client.export_initial_telemetry_log(MagicMock(), "test-agent") + assert mock_export.call_count == 1 + + # Test failure log + client.export_failure_log("TestError", "Error message") + assert mock_export.call_count == 2 + + # Test latency log + client.export_latency_log(150, "EXECUTE_STATEMENT", "stmt-123") + assert mock_export.call_count == 3 + + def test_error_handling_resilience(self, mock_telemetry_client): + """Test that telemetry errors don't break the client.""" + client = mock_telemetry_client + + # Test that exceptions in telemetry don't propagate + with patch.object(client, "_export_event", side_effect=Exception("Test error")): + # These should not raise exceptions + client.export_initial_telemetry_log(MagicMock(), "test-agent") + client.export_failure_log("TestError", "Error message") + client.export_latency_log(100, "EXECUTE_STATEMENT", "stmt-123") + + # Test executor submission failure + client._executor.submit.side_effect = Exception("Thread pool error") + client._send_telemetry([MagicMock()]) # Should not raise + + +class TestTelemetryHelper: + """Tests for TelemetryHelper utility functions.""" + + def test_system_configuration_caching(self): + """Test that system configuration is cached and contains expected data.""" + config1 = TelemetryHelper.get_driver_system_configuration() + config2 = TelemetryHelper.get_driver_system_configuration() + + # Should be cached (same instance) + assert config1 is config2 + + def test_auth_mechanism_detection(self): + """Test authentication mechanism detection for different providers.""" + test_cases = [ + (AccessTokenAuthProvider("token"), AuthMech.PAT), + (MagicMock(spec=DatabricksOAuthProvider), AuthMech.OAUTH), + (MagicMock(spec=ExternalAuthProvider), AuthMech.OTHER), + (MagicMock(), AuthMech.OTHER), # Unknown provider + (None, None), + ] + + for provider, expected in test_cases: + assert TelemetryHelper.get_auth_mechanism(provider) == expected + + def test_auth_flow_detection(self): + """Test authentication flow detection for OAuth providers.""" + # OAuth with existing tokens + oauth_with_tokens = MagicMock(spec=DatabricksOAuthProvider) + oauth_with_tokens._access_token = "test-access-token" + oauth_with_tokens._refresh_token = "test-refresh-token" + assert ( + TelemetryHelper.get_auth_flow(oauth_with_tokens) + == AuthFlow.TOKEN_PASSTHROUGH + ) + + # Test OAuth with browser-based auth + oauth_with_browser = MagicMock(spec=DatabricksOAuthProvider) + oauth_with_browser._access_token = None + oauth_with_browser._refresh_token = None + oauth_with_browser.oauth_manager = MagicMock() + assert ( + TelemetryHelper.get_auth_flow(oauth_with_browser) + == AuthFlow.BROWSER_BASED_AUTHENTICATION + ) + + # Test non-OAuth provider + pat_auth = AccessTokenAuthProvider("test-token") + assert TelemetryHelper.get_auth_flow(pat_auth) is None + + # Test None auth provider + assert TelemetryHelper.get_auth_flow(None) is None + + def _make_real_federation(self, inner): + """Build a real TokenFederationProvider so attribute renames break tests.""" + return TokenFederationProvider( + hostname="example.cloud.databricks.com", + external_provider=inner, + http_client=MagicMock(), + ) + + def test_token_federation_unwraps_pat(self): + fed = self._make_real_federation(AccessTokenAuthProvider("test-token")) + assert TelemetryHelper.get_auth_mechanism(fed) == AuthMech.PAT + assert TelemetryHelper.get_auth_flow(fed) is None + + def test_token_federation_unwraps_m2m(self): + fed = self._make_real_federation(MagicMock(spec=ExternalAuthProvider)) + assert TelemetryHelper.get_auth_mechanism(fed) == AuthMech.OTHER + assert TelemetryHelper.get_auth_flow(fed) == AuthFlow.CLIENT_CREDENTIALS + + def test_token_federation_unwraps_oauth_browser(self): + oauth = MagicMock(spec=DatabricksOAuthProvider) + oauth._access_token = None + oauth._refresh_token = None + fed = self._make_real_federation(oauth) + assert TelemetryHelper.get_auth_mechanism(fed) == AuthMech.OAUTH + assert TelemetryHelper.get_auth_flow(fed) == AuthFlow.BROWSER_BASED_AUTHENTICATION + + def test_token_federation_unwraps_oauth_passthrough(self): + oauth = MagicMock(spec=DatabricksOAuthProvider) + oauth._access_token = "a" + oauth._refresh_token = "r" + fed = self._make_real_federation(oauth) + assert TelemetryHelper.get_auth_mechanism(fed) == AuthMech.OAUTH + assert TelemetryHelper.get_auth_flow(fed) == AuthFlow.TOKEN_PASSTHROUGH + + def test_token_federation_payload_serialization(self): + """End-to-end: federated PAT must serialize as PAT in the connection-params payload.""" + fed = self._make_real_federation(AccessTokenAuthProvider("test-token")) + params = DriverConnectionParameters( + http_path="/sql/1.0/warehouses/abc", + mode=DatabricksClientType.THRIFT, + host_info=HostDetails(host_url="https://example.cloud.databricks.com", port=443), + auth_mech=TelemetryHelper.get_auth_mechanism(fed), + auth_flow=TelemetryHelper.get_auth_flow(fed), + ) + payload = json.loads(params.to_json()) + assert payload["auth_mech"] == "PAT" + assert "auth_flow" not in payload # None-valued fields are stripped + + def test_token_federation_with_no_inner_provider(self): + """Federation with a None inner provider should not crash; both helpers return None.""" + fed = TokenFederationProvider( + hostname="example.cloud.databricks.com", + external_provider=None, + http_client=MagicMock(), + ) + 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_defaults_false(self): + kernel_kwargs = self._kernel_telemetry_kwargs_for_test({}) + + assert kernel_kwargs["telemetry_enabled"] is False + + +class TestTelemetryFactory: + """Tests for TelemetryClientFactory lifecycle and management.""" + + @pytest.fixture(autouse=True) + def telemetry_system_reset(self): + """Reset telemetry system state before each test.""" + TelemetryClientFactory._clients.clear() + if TelemetryClientFactory._executor: + TelemetryClientFactory._executor.shutdown(wait=True) + TelemetryClientFactory._executor = None + TelemetryClientFactory._initialized = False + yield + TelemetryClientFactory._clients.clear() + if TelemetryClientFactory._executor: + TelemetryClientFactory._executor.shutdown(wait=True) + TelemetryClientFactory._executor = None + TelemetryClientFactory._initialized = False + + def test_client_lifecycle_flow(self): + """Test complete client lifecycle: initialize -> use -> close.""" + session_id_hex = "test-session" + auth_provider = AccessTokenAuthProvider("token") + client_context = MagicMock() + + # Initialize enabled client + with patch( + "databricks.sql.common.unified_http_client.UnifiedHttpClient._setup_pool_managers" + ): + TelemetryClientFactory.initialize_telemetry_client( + telemetry_enabled=True, + session_id_hex=session_id_hex, + auth_provider=auth_provider, + host_url="test-host.com", + batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE, + client_context=client_context, + ) + + client = TelemetryClientFactory.get_telemetry_client("test-host.com") + assert isinstance(client, TelemetryClient) + assert client._session_id_hex == session_id_hex + + # Close client + with patch.object(client, "close") as mock_close: + TelemetryClientFactory.close(host_url="test-host.com") + mock_close.assert_called_once() + + # Should get NoopTelemetryClient after close + + def test_disabled_telemetry_creates_noop_client(self): + """Test that disabled telemetry creates NoopTelemetryClient.""" + session_id_hex = "test-session" + client_context = MagicMock() + + TelemetryClientFactory.initialize_telemetry_client( + telemetry_enabled=False, + session_id_hex=session_id_hex, + auth_provider=None, + host_url="test-host.com", + batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE, + client_context=client_context, + ) + + client = TelemetryClientFactory.get_telemetry_client("test-host.com") + assert isinstance(client, NoopTelemetryClient) + + def test_factory_error_handling(self): + """Test that factory errors fall back to NoopTelemetryClient.""" + session_id = "test-session" + client_context = MagicMock() + + # Simulate initialization error + with patch( + "databricks.sql.telemetry.telemetry_client.TelemetryClient", + side_effect=Exception("Init error"), + ): + TelemetryClientFactory.initialize_telemetry_client( + telemetry_enabled=True, + session_id_hex=session_id, + auth_provider=AccessTokenAuthProvider("token"), + host_url="test-host.com", + batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE, + client_context=client_context, + ) + + # Should fall back to NoopTelemetryClient + client = TelemetryClientFactory.get_telemetry_client("test-host.com") + assert isinstance(client, NoopTelemetryClient) + + def test_factory_shutdown_flow(self): + """Test factory shutdown when last client is removed.""" + session1 = "session-1" + session2 = "session-2" + client_context = MagicMock() + + # Initialize multiple clients + with patch( + "databricks.sql.common.unified_http_client.UnifiedHttpClient._setup_pool_managers" + ): + for session in [session1, session2]: + TelemetryClientFactory.initialize_telemetry_client( + telemetry_enabled=True, + session_id_hex=session, + auth_provider=AccessTokenAuthProvider("token"), + host_url="test-host.com", + batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE, + client_context=client_context, + ) + + # Factory should be initialized + assert TelemetryClientFactory._initialized is True + assert TelemetryClientFactory._executor is not None + + # Close first client - factory should stay initialized + TelemetryClientFactory.close(host_url="test-host.com") + assert TelemetryClientFactory._initialized is True + + # Close second client - factory should shut down + TelemetryClientFactory.close(host_url="test-host.com") + assert TelemetryClientFactory._initialized is False + assert TelemetryClientFactory._executor is None + + @patch( + "databricks.sql.telemetry.telemetry_client.TelemetryClient.export_failure_log" + ) + @patch("databricks.sql.client.Session") + def test_connection_failure_sends_correct_telemetry_payload( + self, mock_session, mock_export_failure_log + ): + """ + Verify that a connection failure constructs and sends the correct + telemetry payload via _send_telemetry. + """ + + error_message = "Could not connect to host" + # Set up the mock to create a session instance first, then make open() fail + mock_session_instance = MagicMock() + mock_session_instance.is_open = False # Ensure cleanup is safe + 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") + except Exception as e: + assert str(e) == error_message + + mock_export_failure_log.assert_called_once() + call_arguments = mock_export_failure_log.call_args + assert call_arguments[0][0] == "Exception" + assert call_arguments[0][1] == error_message + + +@patch("databricks.sql.client.Session") +class TestTelemetryFeatureFlag: + """Tests the interaction between the telemetry feature flag and connection parameters.""" + + def teardown_method(self): + """Clean up telemetry factory state after each test to prevent test pollution.""" + from databricks.sql.common.feature_flag import FeatureFlagsContextFactory + + TelemetryClientFactory._clients.clear() + FeatureFlagsContextFactory._context_map.clear() + + def _mock_ff_response(self, mock_http_request, enabled: bool): + """Helper method to mock feature flag response for unified HTTP client.""" + mock_response = MagicMock() + mock_response.status = 200 + mock_response.status_code = 200 # Compatibility attribute + payload = { + "flags": [ + { + "name": "databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForPythonDriver", + "value": str(enabled).lower(), + } + ], + "ttl_seconds": 3600, + } + mock_response.json.return_value = payload + mock_response.data = json.dumps(payload).encode() + mock_http_request.return_value = mock_response + + @patch("databricks.sql.common.unified_http_client.UnifiedHttpClient.request") + def test_telemetry_enabled_when_flag_is_true(self, mock_http_request, MockSession): + """Telemetry should be ON when enable_telemetry=True and server flag is 'true'.""" + 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 = ( + False # Connection starts closed for test cleanup + ) + + # Set up mock HTTP client on the session + mock_http_client = MagicMock() + mock_http_client.request = mock_http_request + mock_session_instance.http_client = mock_http_client + + conn = sql.client.Connection( + server_hostname="test", + http_path="test", + access_token="test", + enable_telemetry=True, + ) + + assert conn.telemetry_enabled is True + mock_http_request.assert_called_once() + client = TelemetryClientFactory.get_telemetry_client("test-host") + assert isinstance(client, TelemetryClient) + + @patch("databricks.sql.common.unified_http_client.UnifiedHttpClient.request") + def test_telemetry_disabled_when_flag_is_false( + self, mock_http_request, MockSession + ): + """Telemetry should be OFF when enable_telemetry=True but server 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 = ( + False # Connection starts closed for test cleanup + ) + + # Set up mock HTTP client on the session + mock_http_client = MagicMock() + mock_http_client.request = mock_http_request + mock_session_instance.http_client = mock_http_client + + conn = sql.client.Connection( + server_hostname="test", + http_path="test", + access_token="test", + enable_telemetry=True, + ) + + assert conn.telemetry_enabled is False + mock_http_request.assert_called_once() + client = TelemetryClientFactory.get_telemetry_client("test-host") + assert isinstance(client, NoopTelemetryClient) + + @patch("databricks.sql.common.unified_http_client.UnifiedHttpClient.request") + def test_telemetry_disabled_when_flag_request_fails( + self, mock_http_request, MockSession + ): + """Telemetry should default to OFF if the feature flag network 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 = ( + False # Connection starts closed for test cleanup + ) + + # Set up mock HTTP client on the session + mock_http_client = MagicMock() + mock_http_client.request = mock_http_request + mock_session_instance.http_client = mock_http_client + + conn = sql.client.Connection( + server_hostname="test", + http_path="test", + access_token="test", + enable_telemetry=True, + ) + + assert conn.telemetry_enabled is False + mock_http_request.assert_called_once() + client = TelemetryClientFactory.get_telemetry_client("test-host") + assert isinstance(client, NoopTelemetryClient) + + +class TestTelemetryEventModels: + """Tests for telemetry event model data structures and JSON serialization.""" + + def test_host_details_serialization(self): + """Test HostDetails model serialization.""" + host = HostDetails(host_url="test-host.com", port=443) + + # Test JSON string generation + json_str = host.to_json() + assert isinstance(json_str, str) + parsed = json.loads(json_str) + assert parsed["host_url"] == "test-host.com" + assert parsed["port"] == 443 + + def test_driver_connection_parameters_all_fields(self): + """Test DriverConnectionParameters with all fields populated.""" + host_info = HostDetails(host_url="workspace.databricks.com", port=443) + proxy_info = HostDetails(host_url="proxy.company.com", port=8080) + cf_proxy_info = HostDetails(host_url="cf-proxy.company.com", port=8080) + + params = DriverConnectionParameters( + http_path="/sql/1.0/warehouses/abc123", + mode=DatabricksClientType.SEA, + host_info=host_info, + auth_mech=AuthMech.OAUTH, + auth_flow=AuthFlow.BROWSER_BASED_AUTHENTICATION, + socket_timeout=30000, + azure_workspace_resource_id="/subscriptions/test/resourceGroups/test", + azure_tenant_id="tenant-123", + use_proxy=True, + use_system_proxy=True, + proxy_host_info=proxy_info, + use_cf_proxy=False, + cf_proxy_host_info=cf_proxy_info, + non_proxy_hosts=["localhost", "127.0.0.1"], + allow_self_signed_support=False, + use_system_trust_store=True, + enable_arrow=True, + enable_direct_results=True, + enable_sea_hybrid_results=True, + http_connection_pool_size=100, + rows_fetched_per_block=100000, + async_poll_interval_millis=2000, + support_many_parameters=True, + enable_complex_datatype_support=True, + allowed_volume_ingestion_paths="/Volumes/catalog/schema/volume", + query_tags="team:engineering,project:telemetry", + ) + + # Serialize to JSON and parse back + json_str = params.to_json() + json_dict = json.loads(json_str) + + # Verify all new fields are in JSON + assert json_dict["http_path"] == "/sql/1.0/warehouses/abc123" + assert json_dict["mode"] == "SEA" + assert json_dict["host_info"]["host_url"] == "workspace.databricks.com" + assert json_dict["auth_mech"] == "OAUTH" + assert json_dict["auth_flow"] == "BROWSER_BASED_AUTHENTICATION" + assert json_dict["socket_timeout"] == 30000 + assert json_dict["azure_workspace_resource_id"] == "/subscriptions/test/resourceGroups/test" + assert json_dict["azure_tenant_id"] == "tenant-123" + assert json_dict["use_proxy"] is True + assert json_dict["use_system_proxy"] is True + assert json_dict["proxy_host_info"]["host_url"] == "proxy.company.com" + assert json_dict["use_cf_proxy"] is False + assert json_dict["cf_proxy_host_info"]["host_url"] == "cf-proxy.company.com" + assert json_dict["non_proxy_hosts"] == ["localhost", "127.0.0.1"] + assert json_dict["allow_self_signed_support"] is False + assert json_dict["use_system_trust_store"] is True + assert json_dict["enable_arrow"] is True + assert json_dict["enable_direct_results"] is True + assert json_dict["enable_sea_hybrid_results"] is True + assert json_dict["http_connection_pool_size"] == 100 + assert json_dict["rows_fetched_per_block"] == 100000 + assert json_dict["async_poll_interval_millis"] == 2000 + assert json_dict["support_many_parameters"] is True + assert json_dict["enable_complex_datatype_support"] is True + assert json_dict["allowed_volume_ingestion_paths"] == "/Volumes/catalog/schema/volume" + assert json_dict["query_tags"] == "team:engineering,project:telemetry" + + def test_driver_connection_parameters_minimal_fields(self): + """Test DriverConnectionParameters with only required fields.""" + host_info = HostDetails(host_url="workspace.databricks.com", port=443) + + params = DriverConnectionParameters( + http_path="/sql/1.0/warehouses/abc123", + mode=DatabricksClientType.THRIFT, + host_info=host_info, + ) + + # Note: to_json() filters out None values, so we need to check asdict for complete structure + json_str = params.to_json() + json_dict = json.loads(json_str) + + # Required fields should be present + assert json_dict["http_path"] == "/sql/1.0/warehouses/abc123" + assert json_dict["mode"] == "THRIFT" + assert json_dict["host_info"]["host_url"] == "workspace.databricks.com" + + # Optional fields with None are filtered out by to_json() + # This is expected behavior - None values are excluded from JSON output + + def test_driver_system_configuration_serialization(self): + """Test DriverSystemConfiguration model serialization.""" + sys_config = DriverSystemConfiguration( + driver_name="Databricks SQL Connector for Python", + driver_version="3.0.0", + runtime_name="CPython", + runtime_version="3.11.0", + runtime_vendor="Python Software Foundation", + os_name="Darwin", + os_version="23.0.0", + os_arch="arm64", + char_set_encoding="utf-8", + locale_name="en_US", + client_app_name="MyApp", + ) + + json_str = sys_config.to_json() + json_dict = json.loads(json_str) + + assert json_dict["driver_name"] == "Databricks SQL Connector for Python" + assert json_dict["driver_version"] == "3.0.0" + assert json_dict["runtime_name"] == "CPython" + assert json_dict["runtime_version"] == "3.11.0" + assert json_dict["runtime_vendor"] == "Python Software Foundation" + assert json_dict["os_name"] == "Darwin" + assert json_dict["os_version"] == "23.0.0" + assert json_dict["os_arch"] == "arm64" + assert json_dict["locale_name"] == "en_US" + assert json_dict["char_set_encoding"] == "utf-8" + assert json_dict["client_app_name"] == "MyApp" + + def test_telemetry_event_complete_serialization(self): + """Test complete TelemetryEvent serialization with all nested objects.""" + host_info = HostDetails(host_url="workspace.databricks.com", port=443) + proxy_info = HostDetails(host_url="proxy.company.com", port=8080) + + connection_params = DriverConnectionParameters( + http_path="/sql/1.0/warehouses/abc123", + mode=DatabricksClientType.SEA, + host_info=host_info, + auth_mech=AuthMech.OAUTH, + use_proxy=True, + proxy_host_info=proxy_info, + enable_arrow=True, + rows_fetched_per_block=100000, + ) + + sys_config = DriverSystemConfiguration( + driver_name="Databricks SQL Connector for Python", + driver_version="3.0.0", + runtime_name="CPython", + runtime_version="3.11.0", + runtime_vendor="Python Software Foundation", + os_name="Darwin", + os_version="23.0.0", + os_arch="arm64", + char_set_encoding="utf-8", + ) + + error_info = DriverErrorInfo( + error_name="ConnectionError", + stack_trace="Traceback...", + ) + + event = TelemetryEvent( + session_id="test-session-123", + sql_statement_id="test-stmt-456", + operation_latency_ms=1500, + auth_type="OAUTH", + system_configuration=sys_config, + driver_connection_params=connection_params, + error_info=error_info, + ) + + # Test JSON serialization + json_str = event.to_json() + assert isinstance(json_str, str) + + # Parse and verify structure + parsed = json.loads(json_str) + assert parsed["session_id"] == "test-session-123" + assert parsed["sql_statement_id"] == "test-stmt-456" + assert parsed["operation_latency_ms"] == 1500 + assert parsed["auth_type"] == "OAUTH" + + # Verify nested objects + assert parsed["system_configuration"]["driver_name"] == "Databricks SQL Connector for Python" + assert parsed["driver_connection_params"]["http_path"] == "/sql/1.0/warehouses/abc123" + assert parsed["driver_connection_params"]["use_proxy"] is True + assert parsed["driver_connection_params"]["proxy_host_info"]["host_url"] == "proxy.company.com" + assert parsed["error_info"]["error_name"] == "ConnectionError" + + def test_json_serialization_excludes_none_values(self): + """Test that JSON serialization properly excludes None values.""" + host_info = HostDetails(host_url="workspace.databricks.com", port=443) + + params = DriverConnectionParameters( + http_path="/sql/1.0/warehouses/abc123", + mode=DatabricksClientType.SEA, + host_info=host_info, + # All optional fields left as None + ) + + json_str = params.to_json() + parsed = json.loads(json_str) + + # Required fields present + assert parsed["http_path"] == "/sql/1.0/warehouses/abc123" + + # None values should be EXCLUDED from JSON (not included as null) + # This is the behavior of JsonSerializableMixin + assert "auth_mech" not in parsed + assert "azure_tenant_id" not in parsed + assert "proxy_host_info" not in parsed + + +@patch("databricks.sql.client.Session") +@patch("databricks.sql.common.unified_http_client.UnifiedHttpClient._setup_pool_managers") +class TestConnectionParameterTelemetry: + """Tests for connection parameter population in telemetry.""" + + def test_connection_with_proxy_populates_telemetry(self, mock_setup_pools, mock_session): + """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 + mock_session_instance.port = 443 + mock_session_instance.host = "workspace.databricks.com" + mock_session.return_value = mock_session_instance + + with patch("databricks.sql.telemetry.telemetry_client.TelemetryClient.export_initial_telemetry_log") as mock_export: + conn = sql.connect( + server_hostname="workspace.databricks.com", + http_path="/sql/1.0/warehouses/test", + access_token="test-token", + enable_telemetry=True, + force_enable_telemetry=True, + ) + + # Verify export was called + mock_export.assert_called_once() + call_args = mock_export.call_args + + # Extract driver_connection_params + driver_params = call_args.kwargs.get("driver_connection_params") + assert driver_params is not None + assert isinstance(driver_params, DriverConnectionParameters) + + # Verify fields are populated + assert driver_params.http_path == "/sql/1.0/warehouses/test" + assert driver_params.mode == DatabricksClientType.SEA + assert driver_params.host_info.host_url == "workspace.databricks.com" + assert driver_params.host_info.port == 443 + + def test_connection_with_azure_params_populates_telemetry(self, mock_setup_pools, mock_session): + """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 + mock_session_instance.port = 443 + mock_session_instance.host = "workspace.azuredatabricks.net" + mock_session.return_value = mock_session_instance + + with patch("databricks.sql.telemetry.telemetry_client.TelemetryClient.export_initial_telemetry_log") as mock_export: + conn = sql.connect( + server_hostname="workspace.azuredatabricks.net", + http_path="/sql/1.0/warehouses/test", + access_token="test-token", + azure_workspace_resource_id="/subscriptions/test/resourceGroups/test", + azure_tenant_id="tenant-123", + enable_telemetry=True, + force_enable_telemetry=True, + ) + + mock_export.assert_called_once() + driver_params = mock_export.call_args.kwargs.get("driver_connection_params") + + # Verify Azure fields + assert driver_params.azure_workspace_resource_id == "/subscriptions/test/resourceGroups/test" + assert driver_params.azure_tenant_id == "tenant-123" + + def test_connection_populates_arrow_and_performance_params(self, mock_setup_pools, mock_session): + """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 + mock_session_instance.port = 443 + mock_session_instance.host = "workspace.databricks.com" + mock_session.return_value = mock_session_instance + + with patch("databricks.sql.telemetry.telemetry_client.TelemetryClient.export_initial_telemetry_log") as mock_export: + # Import pyarrow availability check + try: + import pyarrow + arrow_available = True + except ImportError: + arrow_available = False + + conn = sql.connect( + server_hostname="workspace.databricks.com", + http_path="/sql/1.0/warehouses/test", + access_token="test-token", + pool_maxsize=200, + enable_telemetry=True, + force_enable_telemetry=True, + ) + + mock_export.assert_called_once() + driver_params = mock_export.call_args.kwargs.get("driver_connection_params") + + # Verify performance fields + assert driver_params.enable_arrow == arrow_available + assert driver_params.enable_direct_results is True + assert driver_params.http_connection_pool_size == 200 + assert driver_params.rows_fetched_per_block == 100000 # DEFAULT_ARRAY_SIZE + assert driver_params.async_poll_interval_millis == 2000 + assert driver_params.support_many_parameters is True + + def test_federated_pat_populates_telemetry_as_pat(self, mock_setup_pools, mock_session): + """End-to-end: a TokenFederationProvider wrapping a PAT should report mech=PAT in the captured telemetry payload.""" + federated_pat = TokenFederationProvider( + hostname="workspace.databricks.com", + external_provider=AccessTokenAuthProvider("token"), + http_client=MagicMock(), + ) + 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 + mock_session_instance.port = 443 + mock_session_instance.host = "workspace.databricks.com" + mock_session.return_value = mock_session_instance + + with patch( + "databricks.sql.telemetry.telemetry_client.TelemetryClient.export_initial_telemetry_log" + ) as mock_export: + sql.connect( + server_hostname="workspace.databricks.com", + http_path="/sql/1.0/warehouses/test", + access_token="test-token", + enable_telemetry=True, + force_enable_telemetry=True, + ) + + mock_export.assert_called_once() + driver_params = mock_export.call_args.kwargs.get("driver_connection_params") + assert driver_params.auth_mech == AuthMech.PAT + assert driver_params.auth_flow is None + + def test_cf_proxy_fields_default_to_false_none(self, mock_setup_pools, mock_session): + """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 + mock_session_instance.port = 443 + mock_session_instance.host = "workspace.databricks.com" + mock_session.return_value = mock_session_instance + + with patch("databricks.sql.telemetry.telemetry_client.TelemetryClient.export_initial_telemetry_log") as mock_export: + conn = sql.connect( + server_hostname="workspace.databricks.com", + http_path="/sql/1.0/warehouses/test", + access_token="test-token", + enable_telemetry=True, + force_enable_telemetry=True, + ) + + mock_export.assert_called_once() + driver_params = mock_export.call_args.kwargs.get("driver_connection_params") + + # CF proxy not yet supported - should be False/None + assert driver_params.use_cf_proxy is False + assert driver_params.cf_proxy_host_info is None + + +class TestFeatureFlagsContextFactory: + """Tests for FeatureFlagsContextFactory host-level caching.""" + + @pytest.fixture(autouse=True) + def reset_factory(self): + """Reset factory state before/after each test.""" + FeatureFlagsContextFactory._context_map.clear() + if FeatureFlagsContextFactory._executor: + FeatureFlagsContextFactory._executor.shutdown(wait=False) + FeatureFlagsContextFactory._executor = None + yield + FeatureFlagsContextFactory._context_map.clear() + if FeatureFlagsContextFactory._executor: + FeatureFlagsContextFactory._executor.shutdown(wait=False) + FeatureFlagsContextFactory._executor = None + + @pytest.mark.parametrize( + "hosts,expected_contexts", + [ + (["host1.com", "host1.com"], 1), # Same host shares context + (["host1.com", "host2.com"], 2), # Different hosts get separate contexts + (["host1.com", "host1.com", "host2.com"], 2), # Mixed scenario + ], + ) + def test_host_level_caching(self, hosts, expected_contexts): + """Test that contexts are cached by host correctly.""" + contexts = [] + for host in hosts: + conn = MagicMock() + conn.session.host = host + conn.session.http_client = MagicMock() + contexts.append(FeatureFlagsContextFactory.get_instance(conn)) + + assert len(FeatureFlagsContextFactory._context_map) == expected_contexts + if expected_contexts == 1: + assert all(ctx is contexts[0] for ctx in contexts) + + def test_remove_instance_and_executor_cleanup(self): + """Test removal uses host key and cleans up executor when empty.""" + conn1 = MagicMock() + conn1.session.host = "host1.com" + conn1.session.http_client = MagicMock() + + conn2 = MagicMock() + conn2.session.host = "host2.com" + conn2.session.http_client = MagicMock() + + FeatureFlagsContextFactory.get_instance(conn1) + FeatureFlagsContextFactory.get_instance(conn2) + assert FeatureFlagsContextFactory._executor is not None + + FeatureFlagsContextFactory.remove_instance(conn1) + assert len(FeatureFlagsContextFactory._context_map) == 1 + assert FeatureFlagsContextFactory._executor is not None + + FeatureFlagsContextFactory.remove_instance(conn2) + assert len(FeatureFlagsContextFactory._context_map) == 0 + assert FeatureFlagsContextFactory._executor is None From e9a09eb3e42500d097f7bc33b03d0487308562bd Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 28 Aug 2026 00:19:18 +0000 Subject: [PATCH 09/17] ai: apply changes for #925 (1 review thread) Addresses: - #3876849398 at src/databricks/sql/session.py:276 Signed-off-by: peco-engineer-bot[bot] --- src/databricks/sql/session.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index ad23f1fcc..de2476029 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -273,8 +273,14 @@ def _create_backend( "telemetry_batch_size": kwargs.get( "telemetry_batch_size", TelemetryClientFactory.DEFAULT_BATCH_SIZE ), + # Match the connector-wide True default (auth/common.py:55, + # where ClientContext defaults telemetry_circuit_breaker_enabled + # to True on the Thrift/SEA path). Forwarding an explicit default + # here keeps parity with telemetry_batch_size above, rather than + # passing None and letting the kernel silently pick its own + # internal default. "telemetry_circuit_breaker_enabled": kwargs.get( - "_telemetry_circuit_breaker_enabled" + "_telemetry_circuit_breaker_enabled", True ), } return KernelDatabricksClient( From 76562500250c8ae488642fcdf38fbedabc64e669 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 28 Aug 2026 00:29:37 +0000 Subject: [PATCH 10/17] ai: apply changes for #925 (1 review thread) Addresses: - #3876884856 at src/databricks/sql/telemetry/telemetry_client.py:125 Signed-off-by: peco-engineer-bot[bot] --- src/databricks/sql/client.py | 8 ++++++-- tests/unit/test_telemetry.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 1f3a8f69d..805702b19 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -402,8 +402,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/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py index abcdd6767..f6888a3b6 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -488,6 +488,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: From 1b93410959596976f7a906777a65fd5410165112 Mon Sep 17 00:00:00 2001 From: Jay Xiao Date: Fri, 28 Aug 2026 01:17:51 +0000 Subject: [PATCH 11/17] Let kernel default telemetry enabled state Signed-off-by: Jay Xiao --- src/databricks/sql/backend/kernel/client.py | 6 ++---- src/databricks/sql/session.py | 10 +++------- tests/unit/test_session.py | 4 ++-- tests/unit/test_telemetry.py | 11 +++++++++-- 4 files changed, 16 insertions(+), 15 deletions(-) diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index f93ceb9de..75c259acf 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -191,11 +191,9 @@ def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]: # The Python telemetry model does not currently track process # name; omit it and let the kernel fill what it can derive. "process_name": None, - # Defaults to False by design: the kernel path deliberately diverges - # from the connector-wide True default (see session.py and client.py). - # Telemetry is off unless explicitly enabled on this backend. - "telemetry_enabled": bool(options.get("enable_telemetry", False)), } + if options.get("enable_telemetry") is not None: + out["telemetry_enabled"] = bool(options["enable_telemetry"]) if options.get("telemetry_batch_size") is not None: out["telemetry_batch_size"] = options["telemetry_batch_size"] if options.get("telemetry_circuit_breaker_enabled") is not None: diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index de2476029..a4310058f 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -258,13 +258,9 @@ def _create_backend( # identity at Session construction time so kernel-owned # telemetry can populate its system configuration. kernel_telemetry_options = { - # Intentionally defaults to False, diverging from the - # connector-wide True default on the Thrift/SEA path - # (client.py). The kernel path opts out of telemetry unless - # explicitly enabled; this is asserted by - # test_telemetry_enabled_defaults_false_for_kernel_client. - # Do not "fix" this back to True to match the other backends. - "enable_telemetry": kwargs.get("enable_telemetry", False), + # Preserve the caller's explicit telemetry choice. When unset, + # leave it as None so the kernel applies its own default. + "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 diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 6b3f28c12..e5339d9ca 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -619,7 +619,7 @@ def test_telemetry_kwargs_threaded_into_kernel_client(self): finally: conn.close() - def test_telemetry_enabled_defaults_false_for_kernel_client(self): + def test_telemetry_enabled_defaults_none_for_kernel_client(self): import sys import types @@ -651,7 +651,7 @@ def test_telemetry_enabled_defaults_false_for_kernel_client(self): try: _, kwargs = mock_kernel_client.call_args opts = kwargs["telemetry_options"] - assert opts["enable_telemetry"] is False + assert opts["enable_telemetry"] is None finally: conn.close() diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py index f6888a3b6..d2d69b9f9 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -332,10 +332,17 @@ def test_is_telemetry_enabled_returns_false_for_kernel( is expected_kernel_telemetry_enabled ) - def test_kernel_telemetry_enabled_defaults_false(self): + def test_kernel_telemetry_enabled_omitted_when_unset(self): kernel_kwargs = self._kernel_telemetry_kwargs_for_test({}) - assert kernel_kwargs["telemetry_enabled"] is False + 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: From 25879a14c8e54f4b58b3c3194cda34aa402cbd87 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 28 Aug 2026 01:40:30 +0000 Subject: [PATCH 12/17] ai: apply changes for #925 (1 review thread) Addresses: - #3877217286 at src/databricks/sql/session.py:263 Signed-off-by: peco-engineer-bot[bot] --- src/databricks/sql/session.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index a4310058f..c4139b444 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -259,7 +259,19 @@ def _create_backend( # 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 applies its own default. + # 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) From 02002a188b5e32e5d1aef478080fdeed97fafa32 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 28 Aug 2026 02:09:32 +0000 Subject: [PATCH 13/17] ai: apply changes for #925 (1 review thread) Addresses: - #3877271142 at src/databricks/sql/session.py:275 Signed-off-by: peco-engineer-bot[bot] --- src/databricks/sql/client.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 805702b19..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. From 46ef53174226dddbe74124e5c705718e43c78cc6 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 28 Aug 2026 02:17:10 +0000 Subject: [PATCH 14/17] ai: apply changes for #925 (1 review thread) Addresses: - #3877391550 at src/databricks/sql/session.py:290 Signed-off-by: peco-engineer-bot[bot] --- src/databricks/sql/session.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index c4139b444..a4166ef13 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -281,14 +281,23 @@ def _create_backend( "telemetry_batch_size": kwargs.get( "telemetry_batch_size", TelemetryClientFactory.DEFAULT_BATCH_SIZE ), - # Match the connector-wide True default (auth/common.py:55, - # where ClientContext defaults telemetry_circuit_breaker_enabled - # to True on the Thrift/SEA path). Forwarding an explicit default - # here keeps parity with telemetry_batch_size above, rather than - # passing None and letting the kernel silently pick its own - # internal default. + # 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", True + "_telemetry_circuit_breaker_enabled" ), } return KernelDatabricksClient( From a406b9ef9c832c316890827a4a9c80344edfb854 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 28 Aug 2026 02:32:44 +0000 Subject: [PATCH 15/17] ai: apply changes for #925 (1 review thread) Addresses: - #3877431848 at src/databricks/sql/backend/kernel/client.py:180 Signed-off-by: peco-engineer-bot[bot] --- src/databricks/sql/backend/kernel/client.py | 46 +++++++++-- tests/unit/_scratch_kernel_kwargs_probe.py | 0 tests/unit/test_kernel_client.py | 91 +++++++++++++++++++++ 3 files changed, 131 insertions(+), 6 deletions(-) create mode 100644 tests/unit/_scratch_kernel_kwargs_probe.py diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 75c259acf..3fb6fad6b 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 @@ -173,10 +174,39 @@ 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 open (returns ``True``) only when the signature + can't be introspected, so a future non-introspectable binding still gets + the kwargs. + """ + try: + params = inspect.signature(_kernel.Session).parameters + except (TypeError, ValueError): + return True + 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``.""" + """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() - out: Dict[str, Any] = { + candidates: Dict[str, Any] = { "driver_name": system.driver_name, "driver_version": system.driver_version, "runtime_name": system.runtime_name, @@ -193,14 +223,18 @@ def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]: "process_name": None, } if options.get("enable_telemetry") is not None: - out["telemetry_enabled"] = bool(options["enable_telemetry"]) + candidates["telemetry_enabled"] = bool(options["enable_telemetry"]) if options.get("telemetry_batch_size") is not None: - out["telemetry_batch_size"] = options["telemetry_batch_size"] + candidates["telemetry_batch_size"] = options["telemetry_batch_size"] if options.get("telemetry_circuit_breaker_enabled") is not None: - out["telemetry_circuit_breaker_enabled"] = options[ + candidates["telemetry_circuit_breaker_enabled"] = options[ "telemetry_circuit_breaker_enabled" ] - return out + return { + name: value + for name, value in candidates.items() + if _kernel_session_accepts_kwarg(name) + } # ─── Client ───────────────────────────────────────────────────────────────── diff --git a/tests/unit/_scratch_kernel_kwargs_probe.py b/tests/unit/_scratch_kernel_kwargs_probe.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index e707fd25c..8683d7bfa 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -428,6 +428,97 @@ def fake_session(**kw): 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_execute_command_forwards_parameters_to_bind_param(): """``execute_command(parameters=[...])`` routes each parameter through ``bind_tspark_params`` onto the kernel statement before From 7a3ac7ccfc4d6ca20786cb4b5e76aafb96b81da8 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Fri, 28 Aug 2026 02:41:54 +0000 Subject: [PATCH 16/17] ai: apply changes for #925 (1 review thread) Addresses: - #3877505593 at src/databricks/sql/backend/kernel/client.py:194 Signed-off-by: peco-engineer-bot[bot] --- src/databricks/sql/backend/kernel/client.py | 14 ++++-- tests/unit/test_kernel_client.py | 51 +++++++++++++++++++++ 2 files changed, 61 insertions(+), 4 deletions(-) diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 3fb6fad6b..caa7b93f5 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -184,14 +184,20 @@ def _kernel_session_accepts_kwarg(name: str) -> bool: (``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 open (returns ``True``) only when the signature - can't be introspected, so a future non-introspectable binding still gets - the kwargs. + 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 True + return False if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()): return True return name in params diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index 8683d7bfa..5fbf81ae7 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -519,6 +519,57 @@ def fake_session_v0_2_0( 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 From e101e491b07736dbf6cb16603b06dea8499e5da6 Mon Sep 17 00:00:00 2001 From: Jay Xiao Date: Fri, 28 Aug 2026 02:54:40 +0000 Subject: [PATCH 17/17] Remove scratch kernel kwargs probe Signed-off-by: Jay Xiao --- tests/unit/_scratch_kernel_kwargs_probe.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 tests/unit/_scratch_kernel_kwargs_probe.py diff --git a/tests/unit/_scratch_kernel_kwargs_probe.py b/tests/unit/_scratch_kernel_kwargs_probe.py deleted file mode 100644 index e69de29bb..000000000