diff --git a/CHANGELOG.md b/CHANGELOG.md index 2bae4c54e..c976f807e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ - Fix the kernel backend ignoring `use_cloud_fetch=False`, so complete results can be fetched without cloud-storage downloads. - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) - Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most. +- Speed up the first `connect()` in a process by not importing `requests`, `urllib3`, the OAuth client libraries and the telemetry HTTP client until something uses them. Kernel connections never import them. # 4.9.0 (2026-11-04) - Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes. diff --git a/src/databricks/sql/auth/auth.py b/src/databricks/sql/auth/auth.py index a4d4d6f2e..ba88149c4 100755 --- a/src/databricks/sql/auth/auth.py +++ b/src/databricks/sql/auth/auth.py @@ -1,4 +1,5 @@ -from typing import Optional, List +import importlib +from typing import Any, Optional, List, TYPE_CHECKING from databricks.sql.auth.authenticators import ( AuthProvider, @@ -8,7 +9,9 @@ AzureServicePrincipalCredentialProvider, ) from databricks.sql.auth.common import AuthType, ClientContext -from databricks.sql.auth.token_federation import TokenFederationProvider + +if TYPE_CHECKING: + from databricks.sql.auth.token_federation import TokenFederationProvider def get_auth_provider(cfg: ClientContext, http_client): @@ -67,7 +70,9 @@ def get_auth_provider(cfg: ClientContext, http_client): # Always wrap with token federation (falls back gracefully if not needed) if base_provider: - return TokenFederationProvider( + import databricks.sql.auth.auth as _module + + return _module.TokenFederationProvider( hostname=cfg.hostname, external_provider=base_provider, http_client=http_client, @@ -134,3 +139,17 @@ def get_python_sql_connector_auth_provider(hostname: str, http_client, **kwargs) identity_federation_client_id=kwargs.get("identity_federation_client_id"), ) return get_auth_provider(cfg, http_client) + + +_LAZY_NAMES = { + "TokenFederationProvider": "databricks.sql.auth.token_federation", +} + + +def __getattr__(name: str) -> Any: + # Imported on first use rather than at load, so connect() skips them; + # still importable, and patchable here, as module attributes. + source = _LAZY_NAMES.get(name) + if source is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + return getattr(importlib.import_module(source), name) diff --git a/src/databricks/sql/auth/authenticators.py b/src/databricks/sql/auth/authenticators.py index c30a287ee..ed3730afe 100644 --- a/src/databricks/sql/auth/authenticators.py +++ b/src/databricks/sql/auth/authenticators.py @@ -1,13 +1,16 @@ +import importlib import abc import logging -from typing import Callable, Dict, List +from typing import Any, Callable, Dict, List, TYPE_CHECKING from databricks.sql.common.http import HttpHeader -from databricks.sql.auth.oauth import ( - OAuthManager, - RefreshableTokenSource, - ClientCredentialsTokenSource, -) from databricks.sql.auth.endpoint import get_oauth_endpoints + +if TYPE_CHECKING: + from databricks.sql.auth.oauth import ( + ClientCredentialsTokenSource, + OAuthManager, + RefreshableTokenSource, + ) from databricks.sql.auth.common import ( AuthType, get_effective_azure_login_app_id, @@ -74,7 +77,9 @@ def __init__( # Convert to the corresponding scopes in the corresponding IdP cloud_scopes = idp_endpoint.get_scopes_mapping(scopes) - self.oauth_manager = OAuthManager( + import databricks.sql.auth.authenticators as _module + + self.oauth_manager = _module.OAuthManager( port_range=redirect_port_range, client_id=client_id, idp_endpoint=idp_endpoint, @@ -204,8 +209,10 @@ def __init__( def auth_type(self) -> str: return AuthType.AZURE_SP_M2M.value - def get_token_source(self, resource: str) -> RefreshableTokenSource: - return ClientCredentialsTokenSource( + def get_token_source(self, resource: str) -> "RefreshableTokenSource": + import databricks.sql.auth.authenticators as _module + + return _module.ClientCredentialsTokenSource( token_url=f"{self.AZURE_AAD_ENDPOINT}/{self.azure_tenant_id}/{self.AZURE_TOKEN_ENDPOINT}", client_id=self.azure_client_id, client_secret=self.azure_client_secret, @@ -236,3 +243,19 @@ def header_factory() -> Dict[str, str]: return headers return header_factory + + +_LAZY_NAMES = { + "OAuthManager": "databricks.sql.auth.oauth", + "RefreshableTokenSource": "databricks.sql.auth.oauth", + "ClientCredentialsTokenSource": "databricks.sql.auth.oauth", +} + + +def __getattr__(name: str) -> Any: + # Imported on first use rather than at load, so connect() skips them; + # still importable, and patchable here, as module attributes. + source = _LAZY_NAMES.get(name) + if source is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + return getattr(importlib.import_module(source), name) diff --git a/src/databricks/sql/auth/common.py b/src/databricks/sql/auth/common.py index cda30a528..7363086ea 100644 --- a/src/databricks/sql/auth/common.py +++ b/src/databricks/sql/auth/common.py @@ -1,13 +1,23 @@ from enum import Enum import logging +import sys from typing import Optional, List from urllib.parse import urlparse -from databricks.sql.auth.retry import DatabricksRetryPolicy from databricks.sql.common.http import HttpMethod logger = logging.getLogger(__name__) +def is_token_federation_provider(auth_provider) -> bool: + """``isinstance(auth_provider, TokenFederationProvider)`` without importing + token federation (and jwt) just to learn that it isn't in use: an instance + can only exist once its module has been imported.""" + module = sys.modules.get("databricks.sql.auth.token_federation") + return module is not None and isinstance( + auth_provider, module.TokenFederationProvider + ) + + class AuthType(Enum): DATABRICKS_OAUTH = "databricks-oauth" AZURE_OAUTH = "azure-oauth" diff --git a/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index 3477caf6b..867972064 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/src/databricks/sql/backend/kernel/auth_bridge.py @@ -80,7 +80,7 @@ PYSQL_OAUTH_SCOPES, ) from databricks.sql.auth.authenticators import AccessTokenAuthProvider, AuthProvider -from databricks.sql.auth.token_federation import TokenFederationProvider +from databricks.sql.auth.common import is_token_federation_provider from databricks.sql.exc import NotSupportedError, ProgrammingError logger = logging.getLogger(__name__) @@ -115,8 +115,8 @@ def _is_pat(auth_provider: Optional[AuthProvider]) -> bool: """ if isinstance(auth_provider, AccessTokenAuthProvider): return True - if isinstance(auth_provider, TokenFederationProvider) and isinstance( - auth_provider.external_provider, AccessTokenAuthProvider + if is_token_federation_provider(auth_provider) and isinstance( + getattr(auth_provider, "external_provider", None), AccessTokenAuthProvider ): return True return False diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 613b62386..4530db4ef 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -72,10 +72,11 @@ from databricks.sql.backend.types import CommandId, BackendType, CommandState, SessionId from databricks.sql.auth.common import AuthType, ClientContext -from databricks.sql.common.unified_http_client import UnifiedHttpClient from databricks.sql.common.http import HttpMethod if TYPE_CHECKING: + from databricks.sql.common.unified_http_client import UnifiedHttpClient + # Type-annotation-only imports (deferred by ``from __future__ import # annotations``). ``get_protocol_version`` and ``_prepare_native_parameters`` # are typed with these Thrift-generated types, but the Thrift backend and @@ -141,6 +142,11 @@ def __getattr__(name: str) -> Any: from databricks.sql.thrift_api.TCLIService import ttypes return getattr(ttypes, name) + if name == "UnifiedHttpClient": + # Same pattern, so kernel connections don't import urllib3. + from databricks.sql.common.unified_http_client import UnifiedHttpClient + + return UnifiedHttpClient raise AttributeError(f"module {__name__!r} has no attribute {name!r}") @@ -413,11 +419,14 @@ def read(self) -> Optional[OAuthToken]: ) client_context = build_client_context(server_hostname, __version__, **kwargs) - # The kernel owns its HTTP transport. + # The kernel owns its HTTP transport. Resolved through the module so + # the class (and urllib3) load only when a connection needs it. + import databricks.sql.client as _client_module + self.http_client: Optional[UnifiedHttpClient] = ( None if kwargs.get("use_kernel", False) - else UnifiedHttpClient(client_context) + else _client_module.UnifiedHttpClient(client_context) ) try: diff --git a/src/databricks/sql/cloudfetch/downloader.py b/src/databricks/sql/cloudfetch/downloader.py index 295f147cd..d21b54673 100644 --- a/src/databricks/sql/cloudfetch/downloader.py +++ b/src/databricks/sql/cloudfetch/downloader.py @@ -1,8 +1,9 @@ from __future__ import annotations +import importlib import logging from dataclasses import dataclass -from typing import Optional, TYPE_CHECKING +from typing import Any, Optional, TYPE_CHECKING import lz4.frame import time @@ -11,7 +12,6 @@ from databricks.sql.types import SSLOptions from databricks.sql.telemetry.latency_logger import log_latency from databricks.sql.telemetry.models.event import StatementType -from databricks.sql.common.unified_http_client import UnifiedHttpClient if TYPE_CHECKING: # Imported for type annotations only. ``from __future__ import annotations`` @@ -24,6 +24,20 @@ logger = logging.getLogger(__name__) +_LAZY_NAMES = { + "UnifiedHttpClient": "databricks.sql.common.unified_http_client", +} + + +def __getattr__(name: str) -> Any: + # Imported on first use rather than at load, so connect() skips them; + # still importable, and patchable here, as module attributes. + source = _LAZY_NAMES.get(name) + if source is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + return getattr(importlib.import_module(source), name) + + @dataclass class DownloadedFile: """ diff --git a/src/databricks/sql/common/http.py b/src/databricks/sql/common/http.py index cf76a5fba..84fa8751d 100644 --- a/src/databricks/sql/common/http.py +++ b/src/databricks/sql/common/http.py @@ -1,14 +1,6 @@ -import requests -from requests.adapters import HTTPAdapter -from urllib3.util.retry import Retry from enum import Enum -import threading from dataclasses import dataclass -from contextlib import contextmanager -from typing import Generator, Optional import logging -from requests.adapters import HTTPAdapter -from databricks.sql.auth.retry import DatabricksRetryPolicy, CommandType logger = logging.getLogger(__name__) diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index d20fe5df0..3616a529c 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -12,7 +12,6 @@ from databricks.sql import USER_AGENT_NAME from databricks.sql.backend.databricks_client import DatabricksClient 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.common.feature_flag import ( FeatureFlagsContext, @@ -23,6 +22,7 @@ if TYPE_CHECKING: from databricks.sql.backend.thrift_backend import ThriftDatabricksClient from databricks.sql.backend.sea.backend import SeaDatabricksClient + from databricks.sql.common.unified_http_client import UnifiedHttpClient logger = logging.getLogger(__name__) @@ -44,6 +44,11 @@ def __getattr__(name: str) -> Any: from databricks.sql.backend.sea.backend import SeaDatabricksClient return SeaDatabricksClient + if name == "UnifiedHttpClient": + # Same pattern, so connect() doesn't import urllib3. + from databricks.sql.common.unified_http_client import UnifiedHttpClient + + return UnifiedHttpClient raise AttributeError(f"module {__name__!r} has no attribute {name!r}") @@ -52,7 +57,7 @@ def __init__( self, server_hostname: str, http_path: str, - http_client: Optional[UnifiedHttpClient], + http_client: Optional["UnifiedHttpClient"], http_headers: Optional[List[Tuple[str, str]]] = None, session_configuration: Optional[Dict[str, Any]] = None, catalog: Optional[str] = None, diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py index fed2001cb..0ab9f00c8 100644 --- a/src/databricks/sql/telemetry/telemetry_client.py +++ b/src/databricks/sql/telemetry/telemetry_client.py @@ -1,3 +1,4 @@ +import importlib import threading import time import logging @@ -34,28 +35,46 @@ DatabricksOAuthProvider, ExternalAuthProvider, ) -from databricks.sql.auth.token_federation import TokenFederationProvider +from databricks.sql.auth.common import is_token_federation_provider import sys import platform import uuid import locale from databricks.sql.telemetry.utils import BaseTelemetryClient -from databricks.sql.common.unified_http_client import UnifiedHttpClient from databricks.sql.common.http import HttpMethod from databricks.sql.exc import RequestError -from databricks.sql.telemetry.telemetry_push_client import ( - ITelemetryPushClient, - TelemetryPushClient, - CircuitBreakerTelemetryPushClient, -) from databricks.sql.common.url_utils import normalize_host_with_protocol if TYPE_CHECKING: from databricks.sql.client import Connection + from databricks.sql.common.unified_http_client import UnifiedHttpClient + from databricks.sql.telemetry.telemetry_push_client import ( + CircuitBreakerTelemetryPushClient, + ITelemetryPushClient, + TelemetryPushClient, + ) logger = logging.getLogger(__name__) +_LAZY_NAMES = { + "UnifiedHttpClient": "databricks.sql.common.unified_http_client", + "TelemetryPushClient": "databricks.sql.telemetry.telemetry_push_client", + "CircuitBreakerTelemetryPushClient": "databricks.sql.telemetry.telemetry_push_client", + "ITelemetryPushClient": "databricks.sql.telemetry.telemetry_push_client", + "TokenFederationProvider": "databricks.sql.auth.token_federation", +} + + +def __getattr__(name: str) -> Any: + # Imported on first use rather than at load, so connect() skips them; + # still importable, and patchable here, as module attributes. + source = _LAZY_NAMES.get(name) + if source is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + return getattr(importlib.import_module(source), name) + + class TelemetryHelper: """Helper class for getting telemetry related information.""" @@ -90,7 +109,7 @@ def get_auth_mechanism(auth_provider): if not auth_provider: return None - if isinstance(auth_provider, TokenFederationProvider): + if is_token_federation_provider(auth_provider): return TelemetryHelper.get_auth_mechanism(auth_provider.external_provider) if isinstance(auth_provider, AccessTokenAuthProvider): return AuthMech.PAT @@ -107,7 +126,7 @@ def get_auth_flow(auth_provider): if not auth_provider: return None - if isinstance(auth_provider, TokenFederationProvider): + if is_token_federation_provider(auth_provider): return TelemetryHelper.get_auth_flow(auth_provider.external_provider) if isinstance(auth_provider, DatabricksOAuthProvider): if auth_provider._access_token and auth_provider._refresh_token: @@ -213,22 +232,26 @@ def __init__( self._host_url = host_url self._executor = executor + # Resolved here, not at module load, so connect() paths that never + # export telemetry don't import urllib3 and pybreaker. + import databricks.sql.telemetry.telemetry_client as _module + # Create own HTTP client from client context - self._http_client = UnifiedHttpClient(client_context) + self._http_client = _module.UnifiedHttpClient(client_context) # Create telemetry push client based on circuit breaker enabled flag if client_context.telemetry_circuit_breaker_enabled: # Create circuit breaker telemetry push client # (circuit breakers created on-demand) - self._telemetry_push_client: ITelemetryPushClient = ( - CircuitBreakerTelemetryPushClient( - TelemetryPushClient(self._http_client), + self._telemetry_push_client: "ITelemetryPushClient" = ( + _module.CircuitBreakerTelemetryPushClient( + _module.TelemetryPushClient(self._http_client), host_url, ) ) else: # Circuit breaker disabled - use direct telemetry push client - self._telemetry_push_client = TelemetryPushClient(self._http_client) + self._telemetry_push_client = _module.TelemetryPushClient(self._http_client) def _export_event(self, event): """Add an event to the batch queue and flush if batch is full""" diff --git a/tests/unit/test_lazy_http_import.py b/tests/unit/test_lazy_http_import.py new file mode 100644 index 000000000..3e64e9c3b --- /dev/null +++ b/tests/unit/test_lazy_http_import.py @@ -0,0 +1,259 @@ +"""Importing the connector must not import the Python HTTP and OAuth client stack. + +requests, urllib3, oauthlib, jwt and pybreaker are only needed by the auth flows, +the Python HTTP client and telemetry export that use them; a kernel connection uses +none of them. Each check runs in a fresh interpreter because ``sys.modules`` is +process-global. +""" + +import subprocess +import sys +import textwrap + +import pytest + +# Sentinel exit codes, distinct from the 1 of an uncaught exception. +_CLEAN, _LOADED, _MISSING_OPTIONAL_DEP = 10, 11, 12 +# Only a missing optional extra may skip; any other import error fails. +_OPTIONAL_DEPS = ("databricks_sql_kernel", "pyarrow") + +MODULES = [ + "databricks.sql", + "databricks.sql.client", + "databricks.sql.session", + "databricks.sql.auth.auth", + "databricks.sql.auth.authenticators", + "databricks.sql.common.http", + "databricks.sql.cloudfetch.downloader", + "databricks.sql.telemetry.telemetry_client", + "databricks.sql.backend.kernel.client", +] +UNWANTED = ["requests", "urllib3", "oauthlib", "jwt", "pybreaker"] + + +@pytest.mark.parametrize("module_name", MODULES) +def test_module_does_not_import_http_client_libraries(module_name): + script = f""" + import sys + try: + import {module_name} # noqa: F401 + except ImportError as e: + # The kernel backend re-raises a missing wheel as a friendlier + # ImportError, so look for the missing extra along the chain. + cause = e + while cause is not None and not ( + isinstance(cause, ModuleNotFoundError) + and cause.name in {_OPTIONAL_DEPS!r} + ): + cause = cause.__cause__ or cause.__context__ + if cause is None: + raise + print(cause.name) + sys.exit({_MISSING_OPTIONAL_DEP}) + loaded = [m for m in {UNWANTED!r} if m in sys.modules] + print(loaded) + sys.exit({_LOADED} if loaded else {_CLEAN}) + """ + proc = subprocess.run( + [sys.executable, "-c", textwrap.dedent(script)], + capture_output=True, + text=True, + ) + if proc.returncode == _MISSING_OPTIONAL_DEP: + pytest.skip(f"{module_name} needs {proc.stdout.strip()}, not installed") + assert proc.returncode != _LOADED, f"importing {module_name} loaded {proc.stdout}" + assert ( + proc.returncode == _CLEAN + ), f"unexpected exit {proc.returncode}:\n{proc.stderr}" + + +def test_oauth_provider_still_builds_its_manager(): + from unittest.mock import MagicMock, patch + + from databricks.sql.auth.authenticators import DatabricksOAuthProvider + + with patch("databricks.sql.auth.oauth.OAuthManager") as manager, patch.object( + DatabricksOAuthProvider, "_initial_get_token" + ): + DatabricksOAuthProvider( + hostname="example.cloud.databricks.com", + oauth_persistence=None, + redirect_port_range=[8020], + client_id="client", + scopes=["sql"], + http_client=MagicMock(), + ) + manager.assert_called_once() + + +# Names these modules used to import at load, now resolved on first access. +MOVED_NAMES = [ + ("databricks.sql.auth.authenticators", "OAuthManager", "databricks.sql.auth.oauth"), + ( + "databricks.sql.auth.authenticators", + "RefreshableTokenSource", + "databricks.sql.auth.oauth", + ), + ( + "databricks.sql.auth.authenticators", + "ClientCredentialsTokenSource", + "databricks.sql.auth.oauth", + ), + ( + "databricks.sql.auth.auth", + "TokenFederationProvider", + "databricks.sql.auth.token_federation", + ), + ( + "databricks.sql.telemetry.telemetry_client", + "UnifiedHttpClient", + "databricks.sql.common.unified_http_client", + ), + ( + "databricks.sql.telemetry.telemetry_client", + "TelemetryPushClient", + "databricks.sql.telemetry.telemetry_push_client", + ), + ( + "databricks.sql.telemetry.telemetry_client", + "CircuitBreakerTelemetryPushClient", + "databricks.sql.telemetry.telemetry_push_client", + ), + ( + "databricks.sql.telemetry.telemetry_client", + "TokenFederationProvider", + "databricks.sql.auth.token_federation", + ), + ( + "databricks.sql.session", + "UnifiedHttpClient", + "databricks.sql.common.unified_http_client", + ), + ( + "databricks.sql.cloudfetch.downloader", + "UnifiedHttpClient", + "databricks.sql.common.unified_http_client", + ), +] + + +@pytest.mark.parametrize("module_name, name, source", MOVED_NAMES) +def test_moved_names_stay_importable(module_name, name, source): + import importlib + + module = importlib.import_module(module_name) + assert getattr(module, name) is getattr(importlib.import_module(source), name) + + +def test_patching_a_moved_name_reaches_where_it_is_used(): + from unittest.mock import MagicMock, patch + + from databricks.sql.auth.authenticators import DatabricksOAuthProvider + from databricks.sql.telemetry.telemetry_client import TelemetryClient + + with patch( + "databricks.sql.auth.authenticators.OAuthManager" + ) as manager, patch.object(DatabricksOAuthProvider, "_initial_get_token"): + DatabricksOAuthProvider( + hostname="example.cloud.databricks.com", + oauth_persistence=None, + redirect_port_range=[8020], + client_id="client", + scopes=["sql"], + http_client=MagicMock(), + ) + manager.assert_called_once() + + context = MagicMock(telemetry_circuit_breaker_enabled=False) + with patch( + "databricks.sql.telemetry.telemetry_client.UnifiedHttpClient" + ) as http_client: + TelemetryClient( + telemetry_enabled=True, + session_id_hex="01ef", + auth_provider=None, + host_url="example.cloud.databricks.com", + executor=MagicMock(), + batch_size=10, + client_context=context, + ) + http_client.assert_called_once_with(context) + + +def test_patched_token_source_reaches_the_azure_sp_provider(): + from unittest.mock import MagicMock, patch + + from databricks.sql.auth.authenticators import ( + AzureServicePrincipalCredentialProvider, + ) + + provider = AzureServicePrincipalCredentialProvider( + hostname="adb-1.azuredatabricks.net", + azure_client_id="client", + azure_client_secret="secret", + http_client=MagicMock(), + azure_tenant_id="tenant", + ) + with patch( + "databricks.sql.auth.authenticators.ClientCredentialsTokenSource" + ) as source: + assert provider.get_token_source(resource="r") is source.return_value + assert source.call_args.kwargs["extra_params"] == {"resource": "r"} + + +def test_patched_push_clients_reach_the_circuit_breaker_branch(): + from unittest.mock import MagicMock, patch + + from databricks.sql.telemetry.telemetry_client import TelemetryClient + + module = "databricks.sql.telemetry.telemetry_client" + context = MagicMock(telemetry_circuit_breaker_enabled=True) + with patch(f"{module}.UnifiedHttpClient") as http_client, patch( + f"{module}.TelemetryPushClient" + ) as push, patch(f"{module}.CircuitBreakerTelemetryPushClient") as breaker: + client = TelemetryClient( + telemetry_enabled=True, + session_id_hex="01ef", + auth_provider=None, + host_url="example.cloud.databricks.com", + executor=MagicMock(), + batch_size=10, + client_context=context, + ) + push.assert_called_once_with(http_client.return_value) + breaker.assert_called_once_with(push.return_value, "example.cloud.databricks.com") + assert client._telemetry_push_client is breaker.return_value + + +def test_token_federation_check_does_not_import_it(): + script = """ + import sys + from databricks.sql.auth.common import is_token_federation_provider + from databricks.sql.auth.authenticators import AccessTokenAuthProvider + assert not is_token_federation_provider(AccessTokenAuthProvider("t")) + assert "databricks.sql.auth.token_federation" not in sys.modules + """ + proc = subprocess.run( + [sys.executable, "-c", textwrap.dedent(script)], + capture_output=True, + text=True, + ) + assert proc.returncode == 0, proc.stderr + + +def test_token_federation_check_matches_real_providers(): + from unittest.mock import MagicMock + + from databricks.sql.auth.authenticators import AccessTokenAuthProvider + from databricks.sql.auth.common import is_token_federation_provider + from databricks.sql.auth.token_federation import TokenFederationProvider + from databricks.sql.telemetry.models.enums import AuthMech + from databricks.sql.telemetry.telemetry_client import TelemetryHelper + + wrapped = TokenFederationProvider( + hostname="example.cloud.databricks.com", + external_provider=AccessTokenAuthProvider("t"), + http_client=MagicMock(), + ) + assert is_token_federation_provider(wrapped) + assert TelemetryHelper.get_auth_mechanism(wrapped) == AuthMech.PAT