fix: add timeout to requests.post()/get() calls in oauth.py

requests.post() and requests.get() calls in oauth.py's retrieve_token(),
get_azure_entra_id_workspace_endpoints(), and PATOAuthTokenExchange.refresh()
do not pass a timeout= parameter. When the OAuth endpoint is unreachable or
slow, these calls block indefinitely. The SDK's per-request timeout
(session.request(timeout=60)) does not protect against this because the
token refresh runs inside session.auth, before the timeout takes effect.

Add an http_timeout_seconds field to ClientCredentials and
PATOAuthTokenExchange dataclasses (default 60, matching _BaseClient), and
a timeout parameter to retrieve_token() and
get_azure_entra_id_workspace_endpoints(). All call sites in
credentials_provider.py and config.py now pass cfg.http_timeout_seconds
so the timeout is user-configurable via Config.

Fixes #1338

Signed-off-by: Chris Grierson <[email protected]>

Author: Chris Grierson

databricks/sdk/config.py

@@ -15,13 +15,15 @@
 from ._base_client import _fix_host_if_needed
 from .client_types import ClientType, HostType
 from .clock import Clock, RealClock
-from .credentials_provider import (CredentialsStrategy, DefaultCredentials,
-                                   OAuthCredentialsProvider)
-from .environments import (ALL_ENVS, AzureEnvironment, Cloud,
-                           DatabricksEnvironment, get_environment_for_hostname)
-from .oauth import (OidcEndpoints, Token,
-                    get_azure_entra_id_workspace_endpoints,
-                    get_endpoints_from_url, get_host_metadata)
+from .credentials_provider import CredentialsStrategy, DefaultCredentials, OAuthCredentialsProvider
+from .environments import ALL_ENVS, AzureEnvironment, Cloud, DatabricksEnvironment, get_environment_for_hostname
+from .oauth import (
+    OidcEndpoints,
+    Token,
+    get_azure_entra_id_workspace_endpoints,
+    get_endpoints_from_url,
+    get_host_metadata,
+)
 
 logger = logging.getLogger("databricks.sdk")
 
@@ -546,7 +548,7 @@ def oidc_endpoints(self) -> Optional[OidcEndpoints]:
         if not self.host:
             return None
         if self.is_azure and self.azure_client_id:
-            return get_azure_entra_id_workspace_endpoints(self.host)
+            return get_azure_entra_id_workspace_endpoints(self.host, timeout=self.http_timeout_seconds or 60)
         return self.databricks_oidc_endpoints
 
     def debug_string(self) -> str:

databricks/sdk/credentials_provider.py

@@ -158,9 +158,7 @@ def runtime_native_auth(cfg: "Config") -> Optional[CredentialsProvider]:
     # This import MUST be after the "DATABRICKS_RUNTIME_VERSION" check
     # above, so that we are not throwing import errors when not in
     # runtime and no config variables are set.
-    from databricks.sdk.runtime import (init_runtime_legacy_auth,
-                                        init_runtime_native_auth,
-                                        init_runtime_repl_auth)
+    from databricks.sdk.runtime import init_runtime_legacy_auth, init_runtime_native_auth, init_runtime_repl_auth
 
     for init in [
         init_runtime_native_auth,
@@ -203,6 +201,7 @@ def get_notebook_pat_token() -> Optional[str]:
         host=cfg.host,
         scopes=cfg.get_scopes_as_string(),
         authorization_details=cfg.authorization_details,
+        http_timeout_seconds=cfg.http_timeout_seconds or 60,
     )
 
     def inner() -> Dict[str, str]:
@@ -232,6 +231,7 @@ def oauth_service_principal(cfg: "Config") -> Optional[CredentialsProvider]:
         use_header=True,
         disable_async=cfg.disable_async_token_refresh,
         authorization_details=cfg.authorization_details,
+        http_timeout_seconds=cfg.http_timeout_seconds or 60,
     )
 
     def inner() -> Dict[str, str]:
@@ -258,7 +258,7 @@ def external_browser(cfg: "Config") -> Optional[CredentialsProvider]:
     elif cfg.azure_client_id:
         client_id = cfg.azure_client_id
         client_secret = cfg.azure_client_secret
-        oidc_endpoints = get_azure_entra_id_workspace_endpoints(cfg.host)
+        oidc_endpoints = get_azure_entra_id_workspace_endpoints(cfg.host, timeout=cfg.http_timeout_seconds or 60)
     if not client_id:
         client_id = "databricks-cli"
         oidc_endpoints = cfg.databricks_oidc_endpoints
@@ -348,6 +348,7 @@ def token_source_for(resource: str) -> oauth.TokenSource:
             disable_async=cfg.disable_async_token_refresh,
             scopes=cfg.get_scopes_as_string(),
             authorization_details=cfg.authorization_details,
+            http_timeout_seconds=cfg.http_timeout_seconds or 60,
         )
 
     _ensure_host_present(cfg, token_source_for)
@@ -470,6 +471,7 @@ def token_source_for(audience: str) -> oauth.TokenSource:
             use_params=True,
             disable_async=cfg.disable_async_token_refresh,
             authorization_details=cfg.authorization_details,
+            http_timeout_seconds=cfg.http_timeout_seconds or 60,
         )
 
     def refreshed_headers() -> Dict[str, str]:
@@ -532,7 +534,9 @@ def github_oidc_azure(cfg: "Config") -> Optional[CredentialsProvider]:
     aad_endpoint = cfg.arm_environment.active_directory_endpoint
     if not cfg.azure_tenant_id:
         # detect Azure AD Tenant ID if it's not specified directly
-        token_endpoint = get_azure_entra_id_workspace_endpoints(cfg.host).token_endpoint
+        token_endpoint = get_azure_entra_id_workspace_endpoints(
+            cfg.host, timeout=cfg.http_timeout_seconds or 60
+        ).token_endpoint
         cfg.azure_tenant_id = token_endpoint.replace(aad_endpoint, "").split("/")[0]
 
     inner = oauth.ClientCredentials(
@@ -548,6 +552,7 @@ def github_oidc_azure(cfg: "Config") -> Optional[CredentialsProvider]:
         disable_async=cfg.disable_async_token_refresh,
         scopes=cfg.get_scopes_as_string(),
         authorization_details=cfg.authorization_details,
+        http_timeout_seconds=cfg.http_timeout_seconds or 60,
     )
 
     def refreshed_headers() -> Dict[str, str]:

databricks/sdk/oauth.py

@@ -194,6 +194,7 @@ def retrieve_token(
     use_params=False,
     use_header=False,
     headers=None,
+    timeout=60,
 ) -> Token:
     logger.debug(f"Retrieving token for {client_id}")
     if use_params:
@@ -206,7 +207,7 @@ def retrieve_token(
         auth = requests.auth.HTTPBasicAuth(client_id, client_secret)
     else:
         auth = IgnoreNetrcAuth()
-    resp = requests.post(token_url, params, auth=auth, headers=headers)
+    resp = requests.post(token_url, params, auth=auth, headers=headers, timeout=timeout)
     if not resp.ok:
         if resp.headers["Content-Type"].startswith("application/json"):
             err = resp.json()
@@ -513,16 +514,18 @@ def get_unified_endpoints(host: str, account_id: str, client: _BaseClient = _Bas
 
 def get_azure_entra_id_workspace_endpoints(
     host: str,
+    timeout: int = 60,
 ) -> Optional[OidcEndpoints]:
     """
     Get the Azure Entra ID endpoints for a given workspace. Can only be used when authenticating to Azure Databricks
     using an application registered in Azure Entra ID.
     :param host: The Databricks workspace host.
+    :param timeout: HTTP request timeout in seconds.
     :return: The OIDC endpoints for the workspace's Azure Entra ID tenant.
     """
     # In Azure, this workspace endpoint redirects to the Entra ID authorization endpoint
     host = _fix_host_if_needed(host)
-    res = requests.get(f"{host}/oidc/oauth2/v2.0/authorize", allow_redirects=False)
+    res = requests.get(f"{host}/oidc/oauth2/v2.0/authorize", allow_redirects=False, timeout=timeout)
     real_auth_url = res.headers.get("location")
     if not real_auth_url:
         return None
@@ -828,6 +831,7 @@ class ClientCredentials(Refreshable):
     use_header: bool = False
     disable_async: bool = True
     authorization_details: str = None
+    http_timeout_seconds: int = 60
 
     def __post_init__(self):
         super().__init__(disable_async=self.disable_async)
@@ -848,6 +852,7 @@ def refresh(self) -> Token:
             params,
             use_params=self.use_params,
             use_header=self.use_header,
+            timeout=self.http_timeout_seconds,
         )
 
 
@@ -874,6 +879,7 @@ class PATOAuthTokenExchange(Refreshable):
     scopes: str
     authorization_details: str = None
     disable_async: bool = True
+    http_timeout_seconds: int = 60
 
     def __post_init__(self):
         super().__init__(disable_async=self.disable_async)
@@ -890,7 +896,7 @@ def refresh(self) -> Token:
         if self.authorization_details:
             params["authorization_details"] = self.authorization_details
 
-        resp = requests.post(token_exchange_url, params)
+        resp = requests.post(token_exchange_url, params, timeout=self.http_timeout_seconds)
         if not resp.ok:
             if resp.headers["Content-Type"].startswith("application/json"):
                 err = resp.json()