diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 3d3a33b86..f93ceb9de 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,38 @@ 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, + # 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: + out["telemetry_batch_size"] = options["telemetry_batch_size"] + if options.get("telemetry_circuit_breaker_enabled") is not None: + out["telemetry_circuit_breaker_enabled"] = options[ + "telemetry_circuit_breaker_enabled" + ] + return out + + # ─── Client ───────────────────────────────────────────────────────────────── @@ -226,6 +259,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 +375,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 +419,7 @@ def open_session( **auth_kwargs, **tls_kwargs, **retry_kwargs, + **telemetry_kwargs, **http_headers_kwargs, ) except Exception as exc: diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 1f3a8f69d..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/src/databricks/sql/session.py b/src/databricks/sql/session.py index bced2db97..de2476029 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -13,6 +13,7 @@ from databricks.sql.backend.types import SessionId, BackendType from databricks.sql.common.unified_http_client import UnifiedHttpClient from databricks.sql.common.agent import detect as detect_agent +from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory if TYPE_CHECKING: from databricks.sql.backend.thrift_backend import ThriftDatabricksClient @@ -250,6 +251,38 @@ 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 = { + # 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), + # 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 + ), + # 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", True + ), + } return KernelDatabricksClient( server_hostname=server_hostname, http_path=http_path, @@ -263,6 +296,7 @@ def _create_backend( auth_options=kernel_auth_options, retry_options=kernel_retry_options, request_timeout_secs=kwargs.get("_socket_timeout"), + telemetry_options=kernel_telemetry_options, ) # These reference the lazily-resolved module attributes defined via diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py index a1ab4b5e0..2051fb2f8 100644 --- a/src/databricks/sql/telemetry/telemetry_client.py +++ b/src/databricks/sql/telemetry/telemetry_client.py @@ -122,6 +122,9 @@ def get_auth_flow(auth_provider): @staticmethod def is_telemetry_enabled(connection: "Connection") -> bool: + if getattr(connection.session, "use_kernel", False): + return False + # Fast path: force enabled - skip feature flag fetch entirely if connection.force_enable_telemetry: return True diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index 7e8249553..e707fd25c 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -368,6 +368,66 @@ def fake_session(**kw): assert captured["request_timeout_secs"] == timeout +def test_open_session_passes_phase_7_telemetry_kwargs_to_kernel(monkeypatch): + """Kernel telemetry phase 7 added binding/runtime identity and + telemetry config kwargs to ``databricks_sql_kernel.Session``.""" + captured = {} + + def fake_session(**kw): + captured.update(kw) + sess = MagicMock() + sess.session_id = "sess-id" + return sess + + monkeypatch.setattr(kernel_client._kernel, "Session", fake_session) + monkeypatch.setattr( + kernel_client.TelemetryHelper, + "get_driver_system_configuration", + lambda: types.SimpleNamespace( + driver_name="Databricks SQL Python Connector", + driver_version="1.2.3", + runtime_name="Python 3.12.0", + runtime_version="3.12.0", + runtime_vendor="CPython", + os_name="Linux", + os_version="6.1", + os_arch="x86_64", + client_app_name=None, + locale_name="en_US", + char_set_encoding="utf-8", + ), + ) + + c = kernel_client.KernelDatabricksClient( + server_hostname="example.cloud.databricks.com", + http_path="/sql/1.0/warehouses/abc", + auth_provider=AccessTokenAuthProvider("dapi-test"), + ssl_options=None, + telemetry_options={ + "enable_telemetry": True, + "telemetry_batch_size": 17, + "telemetry_circuit_breaker_enabled": False, + }, + ) + c.open_session(session_configuration=None, catalog=None, schema=None) + + assert captured["driver_name"] == "Databricks SQL Python Connector" + assert captured["driver_version"] == "1.2.3" + assert captured["runtime_name"] == "Python 3.12.0" + assert captured["runtime_version"] == "3.12.0" + assert captured["runtime_vendor"] == "CPython" + assert captured["os_name"] == "Linux" + assert captured["os_version"] == "6.1" + assert captured["os_arch"] == "x86_64" + assert captured["client_app_name"] is None + assert captured["locale_name"] == "en_US" + assert captured["char_set_encoding"] == "utf-8" + assert captured["process_name"] is None + assert captured["telemetry_enabled"] is True + assert captured["telemetry_batch_size"] == 17 + assert captured["telemetry_circuit_breaker_enabled"] is False + + def test_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..6b3f28c12 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -570,6 +570,92 @@ def test_oauth_token_cache_enabled_threaded_into_kernel_auth_options(self): conn.close() +class TestKernelTelemetryOptionsThreading: + """The kernel path must forward telemetry options from connect() + into ``KernelDatabricksClient`` so phase-7 PyO3 Session kwargs can + be populated before the kernel opens its session.""" + + PACKAGE = "databricks.sql" + + def test_telemetry_kwargs_threaded_into_kernel_client(self): + import sys + import types + + pytest.importorskip( + "pyarrow", + reason="kernel client module imports pyarrow at load", + ) + + fake = types.ModuleType("databricks_sql_kernel") + fake.KernelError = type("KernelError", (Exception,), {}) + fake.Session = MagicMock() + + with patch.dict(sys.modules, {"databricks_sql_kernel": fake}), patch( + "databricks.sql.backend.kernel.client.KernelDatabricksClient" + ) as mock_kernel_client, patch( + "%s.session.get_python_sql_connector_auth_provider" % self.PACKAGE + ): + instance = mock_kernel_client.return_value + instance.open_session.return_value = SessionId( + BackendType.SEA, "sess-id", None + ) + + conn = databricks.sql.connect( + server_hostname="foo", + http_path="/sql/1.0/warehouses/abc", + use_kernel=True, + access_token="dapi-xyz", + enable_telemetry=True, + force_enable_telemetry=False, + telemetry_batch_size=17, + _telemetry_circuit_breaker_enabled=False, + ) + try: + _, kwargs = mock_kernel_client.call_args + opts = kwargs["telemetry_options"] + assert opts["enable_telemetry"] is True + assert opts["telemetry_batch_size"] == 17 + assert opts["telemetry_circuit_breaker_enabled"] is False + finally: + conn.close() + + def test_telemetry_enabled_defaults_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 — 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..f6888a3b6 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -270,6 +270,73 @@ def test_token_federation_with_no_inner_provider(self): assert TelemetryHelper.get_auth_mechanism(fed) is None assert TelemetryHelper.get_auth_flow(fed) is None + @staticmethod + def _kernel_telemetry_kwargs_for_test(options): + import importlib + import sys + import types + + pytest.importorskip( + "pyarrow", + reason="kernel client module imports pyarrow at load", + ) + + fake = types.ModuleType("databricks_sql_kernel") + fake.KernelError = type("KernelError", (Exception,), {}) + fake.Session = MagicMock() + + sys.modules.pop("databricks.sql.backend.kernel.client", None) + import databricks.sql.backend.kernel as kernel_pkg + + if hasattr(kernel_pkg, "client"): + delattr(kernel_pkg, "client") + + try: + with patch.dict(sys.modules, {"databricks_sql_kernel": fake}): + kernel_client = importlib.import_module( + "databricks.sql.backend.kernel.client" + ) + return kernel_client._kernel_telemetry_kwargs(options) + finally: + sys.modules.pop("databricks.sql.backend.kernel.client", None) + if hasattr(kernel_pkg, "client"): + delattr(kernel_pkg, "client") + + @pytest.mark.parametrize( + ("enable_telemetry", "expected_kernel_telemetry_enabled"), + [ + (True, True), + (False, False), + ], + ) + def test_is_telemetry_enabled_returns_false_for_kernel( + self, + enable_telemetry, + expected_kernel_telemetry_enabled, + ): + connection = MagicMock() + connection.session.use_kernel = True + connection.force_enable_telemetry = True + connection.enable_telemetry = enable_telemetry + + assert TelemetryHelper.is_telemetry_enabled(connection) is False + + kernel_kwargs = self._kernel_telemetry_kwargs_for_test( + { + "enable_telemetry": enable_telemetry, + "force_enable_telemetry": True, + } + ) + assert ( + kernel_kwargs["telemetry_enabled"] + is expected_kernel_telemetry_enabled + ) + + def test_kernel_telemetry_enabled_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.""" @@ -421,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: @@ -457,6 +554,7 @@ def test_telemetry_enabled_when_flag_is_true(self, mock_http_request, MockSessio self._mock_ff_response(mock_http_request, enabled=True) mock_session_instance = MockSession.return_value mock_session_instance.guid_hex = "test-session-ff-true" + mock_session_instance.use_kernel = False mock_session_instance.host = "test-host" # Set host for telemetry client lookup mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = ( @@ -488,6 +586,7 @@ def test_telemetry_disabled_when_flag_is_false( self._mock_ff_response(mock_http_request, enabled=False) mock_session_instance = MockSession.return_value mock_session_instance.guid_hex = "test-session-ff-false" + mock_session_instance.use_kernel = False mock_session_instance.host = "test-host" # Set host for telemetry client lookup mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = ( @@ -519,6 +618,7 @@ def test_telemetry_disabled_when_flag_request_fails( mock_http_request.side_effect = Exception("Network is down") mock_session_instance = MockSession.return_value mock_session_instance.guid_hex = "test-session-ff-fail" + mock_session_instance.use_kernel = False mock_session_instance.host = "test-host" # Set host for telemetry client lookup mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = ( @@ -771,6 +871,7 @@ def test_connection_with_proxy_populates_telemetry(self, mock_setup_pools, mock_ """Test that proxy configuration is captured in telemetry.""" mock_session_instance = MagicMock() mock_session_instance.guid_hex = "test-session-proxy" + mock_session_instance.use_kernel = False mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = False mock_session_instance.use_sea = True @@ -806,6 +907,7 @@ def test_connection_with_azure_params_populates_telemetry(self, mock_setup_pools """Test that Azure-specific parameters are captured in telemetry.""" mock_session_instance = MagicMock() mock_session_instance.guid_hex = "test-session-azure" + mock_session_instance.use_kernel = False mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = False mock_session_instance.use_sea = False @@ -835,6 +937,7 @@ def test_connection_populates_arrow_and_performance_params(self, mock_setup_pool """Test that Arrow and performance parameters are captured in telemetry.""" mock_session_instance = MagicMock() mock_session_instance.guid_hex = "test-session-perf" + mock_session_instance.use_kernel = False mock_session_instance.auth_provider = AccessTokenAuthProvider("token") mock_session_instance.is_open = False mock_session_instance.use_sea = True @@ -879,6 +982,7 @@ def test_federated_pat_populates_telemetry_as_pat(self, mock_setup_pools, mock_s ) mock_session_instance = MagicMock() mock_session_instance.guid_hex = "test-session-fed-pat" + mock_session_instance.use_kernel = False mock_session_instance.auth_provider = federated_pat mock_session_instance.is_open = False mock_session_instance.use_sea = False @@ -906,6 +1010,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