Disable Python telemetry for kernel connections

Author: jay-xiao446

src/databricks/sql/backend/kernel/client.py

@@ -47,6 +47,7 @@
     NotSupportedError,
     ProgrammingError,
 )
+from databricks.sql.telemetry.telemetry_client import TelemetryHelper
 from databricks.sql.thrift_api.TCLIService import ttypes
 
 if TYPE_CHECKING:
@@ -165,6 +166,31 @@ def _is_staging_statement(operation: str) -> bool:
     return verb in _STAGING_VERBS
 
 
+def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]:
+    """Build phase-7 telemetry/system kwargs for ``databricks_sql_kernel.Session``."""
+    system = TelemetryHelper.get_driver_system_configuration()
+    out: Dict[str, Any] = {
+        "driver_name": system.driver_name,
+        "driver_version": system.driver_version,
+        "runtime_name": system.runtime_name,
+        "runtime_version": system.runtime_version,
+        "runtime_vendor": system.runtime_vendor,
+        "os_name": system.os_name,
+        "os_version": system.os_version,
+        "os_arch": system.os_arch,
+        "client_app_name": system.client_app_name,
+        "locale_name": system.locale_name,
+        "char_set_encoding": system.char_set_encoding,
+        # The Python telemetry model does not currently track process
+        # name; omit it and let the kernel fill what it can derive.
+        "process_name": None,
+        "telemetry_enabled": bool(options.get("enable_telemetry", True)),
+    }
+    if options.get("telemetry_batch_size") is not None:
+        out["telemetry_batch_size"] = options["telemetry_batch_size"]
+    return out
+
+
 # ─── Client ─────────────────────────────────────────────────────────────────
 
 
@@ -217,6 +243,9 @@ def __init__(
         # to the kernel ``Session``'s ``retry_*`` kwargs in
         # ``open_session`` via ``_kernel_retry_kwargs``.
         self._retry_options = kwargs.get("retry_options") or {}
+        # Kernel telemetry phase 7 adds binding/runtime identity and
+        # telemetry config kwargs directly to ``databricks_sql_kernel.Session``.
+        self._telemetry_options = kwargs.get("telemetry_options") or {}
         self._catalog = catalog
         self._schema = schema
         # ``_use_arrow_native_complex_types`` is the connector-side
@@ -316,6 +345,7 @@ def open_session(
             # Translate the connector's ``_retry_*`` kwargs into the
             # kernel's ``retry_*`` kwargs. Empty when at defaults.
             retry_kwargs = _kernel_retry_kwargs(self._retry_options)
+            telemetry_kwargs = _kernel_telemetry_kwargs(self._telemetry_options)
             # Forward caller / connector HTTP headers. The kernel applies
             # them on every request; a caller ``User-Agent`` is appended
             # to the kernel's base UA. Only pass the kwarg when there's
@@ -358,6 +388,7 @@ def open_session(
                 **auth_kwargs,
                 **tls_kwargs,
                 **retry_kwargs,
+                **telemetry_kwargs,
                 **http_headers_kwargs,
             )
         except Exception as exc:

src/databricks/sql/session.py

@@ -204,6 +204,16 @@ def _create_backend(
                     "_retry_stop_after_attempts_duration"
                 ),
             }
+            # Forward the binding/runtime identity and telemetry knobs
+            # added by kernel telemetry phase 7. Python-side telemetry
+            # still owns feature-flag evaluation and event export for the
+            # Thrift/SEA paths; the kernel path needs the same driver
+            # identity at Session construction time so kernel-owned
+            # telemetry can populate its system configuration.
+            kernel_telemetry_options = {
+                "enable_telemetry": kwargs.get("enable_telemetry", True),
+                "telemetry_batch_size": kwargs.get("telemetry_batch_size"),
+            }
             return KernelDatabricksClient(
                 server_hostname=server_hostname,
                 http_path=http_path,
@@ -216,6 +226,7 @@ def _create_backend(
                 _use_arrow_native_complex_types=_use_arrow_native_complex_types,
                 auth_options=kernel_auth_options,
                 retry_options=kernel_retry_options,
+                telemetry_options=kernel_telemetry_options,
             )
 
         databricks_client_class: Type[DatabricksClient]

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

tests/unit/test_kernel_client.py

@@ -344,6 +344,64 @@ def fake_session(**kw):
     assert captured.get("complex_types_as_json") is expected_flag
 
 
+def test_open_session_passes_phase_7_telemetry_kwargs_to_kernel(monkeypatch):
+    """Kernel telemetry phase 7 added binding/runtime identity and
+    telemetry config kwargs to ``databricks_sql_kernel.Session``."""
+    captured = {}
+
+    def fake_session(**kw):
+        captured.update(kw)
+        sess = MagicMock()
+        sess.session_id = "sess-id"
+        return sess
+
+    monkeypatch.setattr(kernel_client._kernel, "Session", fake_session)
+    monkeypatch.setattr(
+        kernel_client.TelemetryHelper,
+        "get_driver_system_configuration",
+        lambda: types.SimpleNamespace(
+            driver_name="Databricks SQL Python Connector",
+            driver_version="1.2.3",
+            runtime_name="Python 3.12.0",
+            runtime_version="3.12.0",
+            runtime_vendor="CPython",
+            os_name="Linux",
+            os_version="6.1",
+            os_arch="x86_64",
+            client_app_name=None,
+            locale_name="en_US",
+            char_set_encoding="utf-8",
+        ),
+    )
+
+    c = kernel_client.KernelDatabricksClient(
+        server_hostname="example.cloud.databricks.com",
+        http_path="/sql/1.0/warehouses/abc",
+        auth_provider=AccessTokenAuthProvider("dapi-test"),
+        ssl_options=None,
+        telemetry_options={
+            "enable_telemetry": True,
+            "telemetry_batch_size": 17,
+        },
+    )
+    c.open_session(session_configuration=None, catalog=None, schema=None)
+
+    assert captured["driver_name"] == "Databricks SQL Python Connector"
+    assert captured["driver_version"] == "1.2.3"
+    assert captured["runtime_name"] == "Python 3.12.0"
+    assert captured["runtime_version"] == "3.12.0"
+    assert captured["runtime_vendor"] == "CPython"
+    assert captured["os_name"] == "Linux"
+    assert captured["os_version"] == "6.1"
+    assert captured["os_arch"] == "x86_64"
+    assert captured["client_app_name"] is None
+    assert captured["locale_name"] == "en_US"
+    assert captured["char_set_encoding"] == "utf-8"
+    assert captured["process_name"] is None
+    assert captured["telemetry_enabled"] is True
+    assert captured["telemetry_batch_size"] == 17
+
+
 def test_execute_command_forwards_parameters_to_bind_param():
     """``execute_command(parameters=[...])`` routes each parameter
     through ``bind_tspark_params`` onto the kernel statement before

tests/unit/test_session.py

@@ -1,4 +1,5 @@
 import pytest
+import sys
 from unittest.mock import patch, MagicMock, Mock, PropertyMock
 import gc
 
@@ -14,6 +15,14 @@
 import databricks.sql
 
 
+def _forget_kernel_client_module():
+    sys.modules.pop("databricks.sql.backend.kernel.client", None)
+    import databricks.sql.backend.kernel as kernel_pkg
+
+    if hasattr(kernel_pkg, "client"):
+        delattr(kernel_pkg, "client")
+
+
 class TestSession:
     """
     Unit tests for Session functionality
@@ -427,7 +436,6 @@ class TestKernelRetryOptionsThreading:
     PACKAGE = "databricks.sql"
 
     def test_retry_kwargs_threaded_into_kernel_client(self):
-        import sys
         import types
 
         pytest.importorskip(
@@ -442,6 +450,7 @@ def test_retry_kwargs_threaded_into_kernel_client(self):
         fake = types.ModuleType("databricks_sql_kernel")
         fake.KernelError = type("KernelError", (Exception,), {})
         fake.Session = MagicMock()
+        _forget_kernel_client_module()
 
         # Patch the kernel client class (imported lazily inside
         # _create_backend) and the provider builder; capture the kwargs
@@ -478,6 +487,54 @@ def test_retry_kwargs_threaded_into_kernel_client(self):
                 conn.close()
 
 
+class TestKernelTelemetryOptionsThreading:
+    """The kernel path must forward telemetry options from connect()
+    into ``KernelDatabricksClient`` so phase-7 PyO3 Session kwargs can
+    be populated before the kernel opens its session."""
+
+    PACKAGE = "databricks.sql"
+
+    def test_telemetry_kwargs_threaded_into_kernel_client(self):
+        import types
+
+        pytest.importorskip(
+            "pyarrow",
+            reason="kernel client module imports pyarrow at load",
+        )
+
+        fake = types.ModuleType("databricks_sql_kernel")
+        fake.KernelError = type("KernelError", (Exception,), {})
+        fake.Session = MagicMock()
+        _forget_kernel_client_module()
+
+        with patch.dict(sys.modules, {"databricks_sql_kernel": fake}), patch(
+            "databricks.sql.backend.kernel.client.KernelDatabricksClient"
+        ) as mock_kernel_client, patch(
+            "%s.session.get_python_sql_connector_auth_provider" % self.PACKAGE
+        ):
+            instance = mock_kernel_client.return_value
+            instance.open_session.return_value = SessionId(
+                BackendType.SEA, "sess-id", None
+            )
+
+            conn = databricks.sql.connect(
+                server_hostname="foo",
+                http_path="/sql/1.0/warehouses/abc",
+                use_kernel=True,
+                access_token="dapi-xyz",
+                enable_telemetry=True,
+                force_enable_telemetry=False,
+                telemetry_batch_size=17,
+            )
+            try:
+                _, kwargs = mock_kernel_client.call_args
+                opts = kwargs["telemetry_options"]
+                assert opts["enable_telemetry"] is True
+                assert opts["telemetry_batch_size"] == 17
+            finally:
+                conn.close()
+
+
 class TestKernelUserAgentForwarding:
     """user_agent_entry must reach the kernel on the use_kernel path —
     session.py folds it into the composed User-Agent and includes it in
@@ -488,7 +545,6 @@ class TestKernelUserAgentForwarding:
     PACKAGE = "databricks.sql"
 
     def test_user_agent_entry_reaches_kernel_client_http_headers(self):
-        import sys
         import types
 
         pytest.importorskip(
@@ -498,6 +554,7 @@ def test_user_agent_entry_reaches_kernel_client_http_headers(self):
         fake = types.ModuleType("databricks_sql_kernel")
         fake.KernelError = type("KernelError", (Exception,), {})
         fake.Session = MagicMock()
+        _forget_kernel_client_module()
 
         with patch.dict(sys.modules, {"databricks_sql_kernel": fake}), patch(
             "databricks.sql.backend.kernel.client.KernelDatabricksClient"

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."""