From 07da392cd6e4da1e0924f7d5010d7cf3847463fa Mon Sep 17 00:00:00 2001
From: Vu Anh Phung <vu.phung@databricks.com>
Date: Fri, 13 Nov 2026 19:56:08 +0000
Subject: [PATCH 1/7] perf: don't import requests and oauthlib in connect()

common/http.py imported requests, urllib3 and the retry policy without
using them, auth/common.py imported the retry policy without using it,
and auth/authenticators.py imported the OAuth client (oauthlib, jwt) at
module load for two code paths that need it. Every connect() paid for
these in its first import. Drop the unused imports and import the OAuth
classes where they are used.

Co-authored-by: Isaac <no-reply@databricks.com>
Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
---
 CHANGELOG.md                              |  1 +
 src/databricks/sql/auth/authenticators.py | 16 +++---
 src/databricks/sql/auth/common.py         |  1 -
 src/databricks/sql/common/http.py         |  8 ---
 tests/unit/test_lazy_http_import.py       | 69 +++++++++++++++++++++++
 5 files changed, 79 insertions(+), 16 deletions(-)
 create mode 100644 tests/unit/test_lazy_http_import.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2bae4c54e..70c001184 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` and the OAuth client libraries until an auth flow that uses them runs.
 
 # 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/authenticators.py b/src/databricks/sql/auth/authenticators.py
index c30a287ee..179381b29 100644
--- a/src/databricks/sql/auth/authenticators.py
+++ b/src/databricks/sql/auth/authenticators.py
@@ -1,13 +1,11 @@
 import abc
 import logging
-from typing import Callable, Dict, List
+from typing import 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 RefreshableTokenSource
 from databricks.sql.auth.common import (
     AuthType,
     get_effective_azure_login_app_id,
@@ -74,6 +72,8 @@ def __init__(
             # Convert to the corresponding scopes in the corresponding IdP
             cloud_scopes = idp_endpoint.get_scopes_mapping(scopes)
 
+            from databricks.sql.auth.oauth import OAuthManager
+
             self.oauth_manager = OAuthManager(
                 port_range=redirect_port_range,
                 client_id=client_id,
@@ -204,7 +204,9 @@ def __init__(
     def auth_type(self) -> str:
         return AuthType.AZURE_SP_M2M.value
 
-    def get_token_source(self, resource: str) -> RefreshableTokenSource:
+    def get_token_source(self, resource: str) -> "RefreshableTokenSource":
+        from databricks.sql.auth.oauth import ClientCredentialsTokenSource
+
         return ClientCredentialsTokenSource(
             token_url=f"{self.AZURE_AAD_ENDPOINT}/{self.azure_tenant_id}/{self.AZURE_TOKEN_ENDPOINT}",
             client_id=self.azure_client_id,
diff --git a/src/databricks/sql/auth/common.py b/src/databricks/sql/auth/common.py
index cda30a528..245e513a6 100644
--- a/src/databricks/sql/auth/common.py
+++ b/src/databricks/sql/auth/common.py
@@ -2,7 +2,6 @@
 import logging
 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__)
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/tests/unit/test_lazy_http_import.py b/tests/unit/test_lazy_http_import.py
new file mode 100644
index 000000000..0c8e61e21
--- /dev/null
+++ b/tests/unit/test_lazy_http_import.py
@@ -0,0 +1,69 @@
+"""Importing the connector must not import requests or the OAuth client libraries.
+
+They are only needed by the auth flows and HTTP paths that use 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, _IMPORT_FAILED = 10, 11, 12
+
+MODULES = [
+    "databricks.sql",
+    "databricks.sql.client",
+    "databricks.sql.session",
+    "databricks.sql.auth.authenticators",
+    "databricks.sql.common.http",
+]
+UNWANTED = ["requests", "oauthlib"]
+
+
+@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 BaseException:
+            import traceback
+            traceback.print_exc()
+            sys.exit({_IMPORT_FAILED})
+        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 == _IMPORT_FAILED:
+        pytest.skip(f"{module_name} not importable here:\n{proc.stderr}")
+    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()

From 131ebf45fba0d995e838c74b94d4d97f49657ec9 Mon Sep 17 00:00:00 2001
From: Vu Anh Phung <vu.phung@databricks.com>
Date: Fri, 13 Nov 2026 20:12:55 +0000
Subject: [PATCH 2/7] perf: don't import urllib3, jwt or pybreaker in connect()
 either

telemetry_client imported the telemetry HTTP client (urllib3) and push
clients (pybreaker) at module load, and telemetry_client and the kernel
auth bridge imported token federation (jwt) for isinstance checks;
client.py, session.py and the cloudfetch downloader imported
UnifiedHttpClient at module load. Import the HTTP and push clients where
they are constructed, check for token federation through sys.modules
(an instance can't exist before its module is imported), and resolve
client.UnifiedHttpClient lazily so its patch seam keeps working.

Co-authored-by: Isaac <no-reply@databricks.com>
Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
---
 CHANGELOG.md                                  |  2 +-
 src/databricks/sql/auth/auth.py               |  3 +-
 src/databricks/sql/auth/common.py             | 11 +++++
 .../sql/backend/kernel/auth_bridge.py         |  6 +--
 src/databricks/sql/client.py                  | 15 ++++--
 src/databricks/sql/cloudfetch/downloader.py   |  1 -
 src/databricks/sql/session.py                 |  4 +-
 .../sql/telemetry/telemetry_client.py         | 23 +++++----
 tests/unit/test_lazy_http_import.py           | 48 +++++++++++++++++--
 9 files changed, 88 insertions(+), 25 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 70c001184..c976f807e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,7 +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` and the OAuth client libraries until an auth flow that uses them runs.
+- 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..201630dce 100755
--- a/src/databricks/sql/auth/auth.py
+++ b/src/databricks/sql/auth/auth.py
@@ -8,7 +8,6 @@
     AzureServicePrincipalCredentialProvider,
 )
 from databricks.sql.auth.common import AuthType, ClientContext
-from databricks.sql.auth.token_federation import TokenFederationProvider
 
 
 def get_auth_provider(cfg: ClientContext, http_client):
@@ -67,6 +66,8 @@ def get_auth_provider(cfg: ClientContext, http_client):
 
     # Always wrap with token federation (falls back gracefully if not needed)
     if base_provider:
+        from databricks.sql.auth.token_federation import TokenFederationProvider
+
         return TokenFederationProvider(
             hostname=cfg.hostname,
             external_provider=base_provider,
diff --git a/src/databricks/sql/auth/common.py b/src/databricks/sql/auth/common.py
index 245e513a6..7363086ea 100644
--- a/src/databricks/sql/auth/common.py
+++ b/src/databricks/sql/auth/common.py
@@ -1,5 +1,6 @@
 from enum import Enum
 import logging
+import sys
 from typing import Optional, List
 from urllib.parse import urlparse
 from databricks.sql.common.http import HttpMethod
@@ -7,6 +8,16 @@
 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..5596ba527 100644
--- a/src/databricks/sql/cloudfetch/downloader.py
+++ b/src/databricks/sql/cloudfetch/downloader.py
@@ -11,7 +11,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``
diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py
index d20fe5df0..9ecf7b28d 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__)
 
@@ -52,7 +52,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..c7af74402 100644
--- a/src/databricks/sql/telemetry/telemetry_client.py
+++ b/src/databricks/sql/telemetry/telemetry_client.py
@@ -34,24 +34,19 @@
     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.telemetry.telemetry_push_client import ITelemetryPushClient
 
 logger = logging.getLogger(__name__)
 
@@ -90,7 +85,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 +102,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,6 +208,14 @@ def __init__(
         self._host_url = host_url
         self._executor = executor
 
+        # Imported here, not at module load, so connect() paths that never
+        # export telemetry don't import urllib3 and pybreaker.
+        from databricks.sql.common.unified_http_client import UnifiedHttpClient
+        from databricks.sql.telemetry.telemetry_push_client import (
+            CircuitBreakerTelemetryPushClient,
+            TelemetryPushClient,
+        )
+
         # Create own HTTP client from client context
         self._http_client = UnifiedHttpClient(client_context)
 
@@ -220,7 +223,7 @@ def __init__(
         if client_context.telemetry_circuit_breaker_enabled:
             # Create circuit breaker telemetry push client
             # (circuit breakers created on-demand)
-            self._telemetry_push_client: ITelemetryPushClient = (
+            self._telemetry_push_client: "ITelemetryPushClient" = (
                 CircuitBreakerTelemetryPushClient(
                     TelemetryPushClient(self._http_client),
                     host_url,
diff --git a/tests/unit/test_lazy_http_import.py b/tests/unit/test_lazy_http_import.py
index 0c8e61e21..d8cda176b 100644
--- a/tests/unit/test_lazy_http_import.py
+++ b/tests/unit/test_lazy_http_import.py
@@ -1,7 +1,9 @@
-"""Importing the connector must not import requests or the OAuth client libraries.
+"""Importing the connector must not import the Python HTTP and OAuth client stack.
 
-They are only needed by the auth flows and HTTP paths that use them. Each check runs
-in a fresh interpreter because ``sys.modules`` is process-global.
+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
@@ -17,10 +19,14 @@
     "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", "oauthlib"]
+UNWANTED = ["requests", "urllib3", "oauthlib", "jwt", "pybreaker"]
 
 
 @pytest.mark.parametrize("module_name", MODULES)
@@ -67,3 +73,37 @@ def test_oauth_provider_still_builds_its_manager():
             http_client=MagicMock(),
         )
     manager.assert_called_once()
+
+
+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

From 34a7759a532a330a495e588627e7a30fa07ccb72 Mon Sep 17 00:00:00 2001
From: Vu Anh Phung <vu.phung@databricks.com>
Date: Fri, 13 Nov 2026 20:32:14 +0000
Subject: [PATCH 3/7] test: skip the HTTP import guard only when an optional
 extra is missing

Co-authored-by: Isaac <no-reply@databricks.com>
Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
---
 tests/unit/test_lazy_http_import.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/tests/unit/test_lazy_http_import.py b/tests/unit/test_lazy_http_import.py
index d8cda176b..047adf1d2 100644
--- a/tests/unit/test_lazy_http_import.py
+++ b/tests/unit/test_lazy_http_import.py
@@ -13,7 +13,9 @@
 import pytest
 
 # Sentinel exit codes, distinct from the 1 of an uncaught exception.
-_CLEAN, _LOADED, _IMPORT_FAILED = 10, 11, 12
+_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",
@@ -35,10 +37,11 @@ def test_module_does_not_import_http_client_libraries(module_name):
         import sys
         try:
             import {module_name}  # noqa: F401
-        except BaseException:
-            import traceback
-            traceback.print_exc()
-            sys.exit({_IMPORT_FAILED})
+        except ModuleNotFoundError as e:
+            if e.name not in {_OPTIONAL_DEPS!r}:
+                raise
+            print(e.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})
@@ -48,8 +51,8 @@ def test_module_does_not_import_http_client_libraries(module_name):
         capture_output=True,
         text=True,
     )
-    if proc.returncode == _IMPORT_FAILED:
-        pytest.skip(f"{module_name} not importable here:\n{proc.stderr}")
+    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

From 8320850fe142245d1a0edde50ea89167ca6ce6f9 Mon Sep 17 00:00:00 2001
From: Vu Anh Phung <vu.phung@databricks.com>
Date: Fri, 13 Nov 2026 20:53:25 +0000
Subject: [PATCH 4/7] test: skip the HTTP import guard when the kernel wheel is
 missing

The kernel backend re-raises a missing databricks_sql_kernel as an
ImportError, so look for the missing extra along the exception chain.

Co-authored-by: Isaac <no-reply@databricks.com>
Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
---
 tests/unit/test_lazy_http_import.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/tests/unit/test_lazy_http_import.py b/tests/unit/test_lazy_http_import.py
index 047adf1d2..a0f759521 100644
--- a/tests/unit/test_lazy_http_import.py
+++ b/tests/unit/test_lazy_http_import.py
@@ -37,10 +37,18 @@ def test_module_does_not_import_http_client_libraries(module_name):
         import sys
         try:
             import {module_name}  # noqa: F401
-        except ModuleNotFoundError as e:
-            if e.name not in {_OPTIONAL_DEPS!r}:
+        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(e.name)
+            print(cause.name)
             sys.exit({_MISSING_OPTIONAL_DEP})
         loaded = [m for m in {UNWANTED!r} if m in sys.modules]
         print(loaded)

From f4839f14967fd73edbf0a3282607d48c88c0cbc0 Mon Sep 17 00:00:00 2001
From: Vu Anh Phung <vu.phung@databricks.com>
Date: Sat, 14 Nov 2026 00:32:50 +0000
Subject: [PATCH 5/7] fix: keep the names modules no longer import at load
 importable

authenticators, auth, telemetry_client, session and the cloudfetch
downloader stopped importing OAuthManager, ClientCredentialsTokenSource,
RefreshableTokenSource, TokenFederationProvider, UnifiedHttpClient and the
telemetry push clients at module load. Downstream imports or patch()es of
those module attributes would then fail. Each module now resolves them on
first access (PEP 562), as client.py already does, and the places that use
them read the module attribute, so a patch there still takes effect.

Co-authored-by: Isaac <no-reply@databricks.com>
Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
---
 src/databricks/sql/auth/auth.py               | 19 +++-
 src/databricks/sql/auth/authenticators.py     | 30 ++++--
 src/databricks/sql/cloudfetch/downloader.py   | 12 ++-
 src/databricks/sql/session.py                 |  5 +
 .../sql/telemetry/telemetry_client.py         | 42 ++++++---
 tests/unit/test_lazy_http_import.py           | 94 +++++++++++++++++++
 6 files changed, 181 insertions(+), 21 deletions(-)

diff --git a/src/databricks/sql/auth/auth.py b/src/databricks/sql/auth/auth.py
index 201630dce..7b2e61606 100755
--- a/src/databricks/sql/auth/auth.py
+++ b/src/databricks/sql/auth/auth.py
@@ -1,4 +1,4 @@
-from typing import Optional, List
+from typing import Any, Optional, List, TYPE_CHECKING
 
 from databricks.sql.auth.authenticators import (
     AuthProvider,
@@ -9,6 +9,9 @@
 )
 from databricks.sql.auth.common import AuthType, ClientContext
 
+if TYPE_CHECKING:
+    from databricks.sql.auth.token_federation import TokenFederationProvider
+
 
 def get_auth_provider(cfg: ClientContext, http_client):
     # Determine the base auth provider
@@ -66,9 +69,9 @@ def get_auth_provider(cfg: ClientContext, http_client):
 
     # Always wrap with token federation (falls back gracefully if not needed)
     if base_provider:
-        from databricks.sql.auth.token_federation import TokenFederationProvider
+        import databricks.sql.auth.auth as _module
 
-        return TokenFederationProvider(
+        return _module.TokenFederationProvider(
             hostname=cfg.hostname,
             external_provider=base_provider,
             http_client=http_client,
@@ -135,3 +138,13 @@ 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)
+
+
+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.
+    if name in ("TokenFederationProvider",):
+        import databricks.sql.auth.token_federation as _source
+
+        return getattr(_source, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/src/databricks/sql/auth/authenticators.py b/src/databricks/sql/auth/authenticators.py
index 179381b29..fdc0b6dcf 100644
--- a/src/databricks/sql/auth/authenticators.py
+++ b/src/databricks/sql/auth/authenticators.py
@@ -1,11 +1,15 @@
 import abc
 import logging
-from typing import Callable, Dict, List, TYPE_CHECKING
+from typing import Any, Callable, Dict, List, TYPE_CHECKING
 from databricks.sql.common.http import HttpHeader
 from databricks.sql.auth.endpoint import get_oauth_endpoints
 
 if TYPE_CHECKING:
-    from databricks.sql.auth.oauth import RefreshableTokenSource
+    from databricks.sql.auth.oauth import (
+        ClientCredentialsTokenSource,
+        OAuthManager,
+        RefreshableTokenSource,
+    )
 from databricks.sql.auth.common import (
     AuthType,
     get_effective_azure_login_app_id,
@@ -72,9 +76,9 @@ def __init__(
             # Convert to the corresponding scopes in the corresponding IdP
             cloud_scopes = idp_endpoint.get_scopes_mapping(scopes)
 
-            from databricks.sql.auth.oauth import OAuthManager
+            import databricks.sql.auth.authenticators as _module
 
-            self.oauth_manager = OAuthManager(
+            self.oauth_manager = _module.OAuthManager(
                 port_range=redirect_port_range,
                 client_id=client_id,
                 idp_endpoint=idp_endpoint,
@@ -205,9 +209,9 @@ def auth_type(self) -> str:
         return AuthType.AZURE_SP_M2M.value
 
     def get_token_source(self, resource: str) -> "RefreshableTokenSource":
-        from databricks.sql.auth.oauth import ClientCredentialsTokenSource
+        import databricks.sql.auth.authenticators as _module
 
-        return ClientCredentialsTokenSource(
+        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,
@@ -238,3 +242,17 @@ def header_factory() -> Dict[str, str]:
             return headers
 
         return header_factory
+
+
+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.
+    if name in (
+        "OAuthManager",
+        "RefreshableTokenSource",
+        "ClientCredentialsTokenSource",
+    ):
+        import databricks.sql.auth.oauth as _source
+
+        return getattr(_source, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/src/databricks/sql/cloudfetch/downloader.py b/src/databricks/sql/cloudfetch/downloader.py
index 5596ba527..c91007bef 100644
--- a/src/databricks/sql/cloudfetch/downloader.py
+++ b/src/databricks/sql/cloudfetch/downloader.py
@@ -2,7 +2,7 @@
 
 import logging
 from dataclasses import dataclass
-from typing import Optional, TYPE_CHECKING
+from typing import Any, Optional, TYPE_CHECKING
 
 import lz4.frame
 import time
@@ -23,6 +23,16 @@
 logger = logging.getLogger(__name__)
 
 
+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.
+    if name in ("UnifiedHttpClient",):
+        import databricks.sql.common.unified_http_client as _source
+
+        return getattr(_source, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 @dataclass
 class DownloadedFile:
     """
diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py
index 9ecf7b28d..3616a529c 100644
--- a/src/databricks/sql/session.py
+++ b/src/databricks/sql/session.py
@@ -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}")
 
 
diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py
index c7af74402..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
@@ -46,11 +47,34 @@
 
 if TYPE_CHECKING:
     from databricks.sql.client import Connection
-    from databricks.sql.telemetry.telemetry_push_client import ITelemetryPushClient
+    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."""
 
@@ -208,30 +232,26 @@ def __init__(
         self._host_url = host_url
         self._executor = executor
 
-        # Imported here, not at module load, so connect() paths that never
+        # Resolved here, not at module load, so connect() paths that never
         # export telemetry don't import urllib3 and pybreaker.
-        from databricks.sql.common.unified_http_client import UnifiedHttpClient
-        from databricks.sql.telemetry.telemetry_push_client import (
-            CircuitBreakerTelemetryPushClient,
-            TelemetryPushClient,
-        )
+        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),
+                _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
index a0f759521..ec68b621c 100644
--- a/tests/unit/test_lazy_http_import.py
+++ b/tests/unit/test_lazy_http_import.py
@@ -86,6 +86,100 @@ def test_oauth_provider_still_builds_its_manager():
     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_token_federation_check_does_not_import_it():
     script = """
         import sys

From cbdae71cdae646b89c2102faf53b290dff8dd299 Mon Sep 17 00:00:00 2001
From: Vu Anh Phung <vu.phung@databricks.com>
Date: Sat, 14 Nov 2026 00:44:35 +0000
Subject: [PATCH 6/7] refactor: write the new lazy-name lookups the same way

auth, authenticators and the cloudfetch downloader now use the same
_LAZY_NAMES map as telemetry_client instead of if-chains.

Co-authored-by: Isaac <no-reply@databricks.com>
Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
---
 src/databricks/sql/auth/auth.py             | 15 ++++++++++-----
 src/databricks/sql/auth/authenticators.py   | 21 ++++++++++++---------
 src/databricks/sql/cloudfetch/downloader.py | 15 ++++++++++-----
 3 files changed, 32 insertions(+), 19 deletions(-)

diff --git a/src/databricks/sql/auth/auth.py b/src/databricks/sql/auth/auth.py
index 7b2e61606..ba88149c4 100755
--- a/src/databricks/sql/auth/auth.py
+++ b/src/databricks/sql/auth/auth.py
@@ -1,3 +1,4 @@
+import importlib
 from typing import Any, Optional, List, TYPE_CHECKING
 
 from databricks.sql.auth.authenticators import (
@@ -140,11 +141,15 @@ def get_python_sql_connector_auth_provider(hostname: str, http_client, **kwargs)
     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.
-    if name in ("TokenFederationProvider",):
-        import databricks.sql.auth.token_federation as _source
-
-        return getattr(_source, name)
-    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    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 fdc0b6dcf..ed3730afe 100644
--- a/src/databricks/sql/auth/authenticators.py
+++ b/src/databricks/sql/auth/authenticators.py
@@ -1,3 +1,4 @@
+import importlib
 import abc
 import logging
 from typing import Any, Callable, Dict, List, TYPE_CHECKING
@@ -244,15 +245,17 @@ def header_factory() -> Dict[str, str]:
         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.
-    if name in (
-        "OAuthManager",
-        "RefreshableTokenSource",
-        "ClientCredentialsTokenSource",
-    ):
-        import databricks.sql.auth.oauth as _source
-
-        return getattr(_source, name)
-    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    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/cloudfetch/downloader.py b/src/databricks/sql/cloudfetch/downloader.py
index c91007bef..d21b54673 100644
--- a/src/databricks/sql/cloudfetch/downloader.py
+++ b/src/databricks/sql/cloudfetch/downloader.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import importlib
 import logging
 from dataclasses import dataclass
 from typing import Any, Optional, TYPE_CHECKING
@@ -23,14 +24,18 @@
 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.
-    if name in ("UnifiedHttpClient",):
-        import databricks.sql.common.unified_http_client as _source
-
-        return getattr(_source, name)
-    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    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

From 1b2feb6b5ac3fb7a9c18b290c9ce5a60b87fced8 Mon Sep 17 00:00:00 2001
From: Vu Anh Phung <vu.phung@databricks.com>
Date: Sat, 14 Nov 2026 00:58:01 +0000
Subject: [PATCH 7/7] test: cover the Azure SP and circuit-breaker uses of the
 lazy names

The new tests reached OAuthManager and the plain telemetry push client, but
not get_token_source's ClientCredentialsTokenSource or the circuit-breaker
branch of TelemetryClient, so a mistyped lazy-name key there would only
fail at runtime.

Co-authored-by: Isaac <no-reply@databricks.com>
Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
---
 tests/unit/test_lazy_http_import.py | 45 +++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/tests/unit/test_lazy_http_import.py b/tests/unit/test_lazy_http_import.py
index ec68b621c..3e64e9c3b 100644
--- a/tests/unit/test_lazy_http_import.py
+++ b/tests/unit/test_lazy_http_import.py
@@ -180,6 +180,51 @@ def test_patching_a_moved_name_reaches_where_it_is_used():
     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