rebase

Author: mgyucht

databricks/sdk/_base_client.py

@@ -18,6 +18,25 @@
 logger = logging.getLogger('databricks.sdk')
 
 
+def fix_host_if_needed(host: Optional[str]) -> Optional[str]:
+    if not host:
+        return host
+
+    # Add a default scheme if it's missing
+    if '://' not in host:
+        host = 'https://' + host
+
+    o = urllib.parse.urlparse(host)
+    # remove trailing slash
+    path = o.path.rstrip('/')
+    # remove port if 443
+    netloc = o.netloc
+    if o.port == 443:
+        netloc = netloc.split(':')[0]
+
+    return urllib.parse.urlunparse((o.scheme, netloc, path, o.params, o.query, o.fragment))
+
+
 class _BaseClient:
 
     def __init__(self,

databricks/sdk/config.py

@@ -14,7 +14,10 @@
 from .credentials_provider import CredentialsStrategy, DefaultCredentials
 from .environments import (ALL_ENVS, AzureEnvironment, Cloud,
                            DatabricksEnvironment, get_environment_for_hostname)
-from .oauth import OidcEndpoints, Token
+from .oauth import (OidcEndpoints, Token, get_account_endpoints,
+                    get_azure_entra_id_workspace_endpoints,
+                    get_workspace_endpoints)
+from ._base_client import fix_host_if_needed
 
 logger = logging.getLogger('databricks.sdk')
 
@@ -118,7 +121,9 @@ def __init__(self,
             self._set_inner_config(kwargs)
             self._load_from_env()
             self._known_file_config_loader()
-            self._fix_host_if_needed()
+            updated_host = fix_host_if_needed(self.host)
+            if updated_host:
+                self.host = updated_host
             self._validate()
             self.init_auth()
             self._init_product(product, product_version)
@@ -250,28 +255,14 @@ def with_user_agent_extra(self, key: str, value: str) -> 'Config':
 
     @property
     def oidc_endpoints(self) -> Optional[OidcEndpoints]:
-        self._fix_host_if_needed()
+        self.host = fix_host_if_needed(self.host)
         if not self.host:
             return None
         if self.is_azure and self.azure_client_id:
-            # Retrieve authorize endpoint to retrieve token endpoint after
-            res = requests.get(f'{self.host}/oidc/oauth2/v2.0/authorize', allow_redirects=False)
-            real_auth_url = res.headers.get('location')
-            if not real_auth_url:
-                return None
-            return OidcEndpoints(authorization_endpoint=real_auth_url,
-                                 token_endpoint=real_auth_url.replace('/authorize', '/token'))
+            return get_azure_entra_id_workspace_endpoints(self.host)
         if self.is_account_client and self.account_id:
-            prefix = f'{self.host}/oidc/accounts/{self.account_id}'
-            return OidcEndpoints(authorization_endpoint=f'{prefix}/v1/authorize',
-                                 token_endpoint=f'{prefix}/v1/token')
-        oidc = f'{self.host}/oidc/.well-known/oauth-authorization-server'
-        res = requests.get(oidc)
-        if res.status_code != 200:
-            return None
-        auth_metadata = res.json()
-        return OidcEndpoints(authorization_endpoint=auth_metadata.get('authorization_endpoint'),
-                             token_endpoint=auth_metadata.get('token_endpoint'))
+            return get_account_endpoints(self.host, self.account_id)
+        return get_workspace_endpoints(self.host)
 
     def debug_string(self) -> str:
         """ Returns log-friendly representation of configured attributes """
@@ -345,24 +336,6 @@ def attributes(cls) -> Iterable[ConfigAttribute]:
         cls._attributes = attrs
         return cls._attributes
 
-    def _fix_host_if_needed(self):
-        if not self.host:
-            return
-
-        # Add a default scheme if it's missing
-        if '://' not in self.host:
-            self.host = 'https://' + self.host
-
-        o = urllib.parse.urlparse(self.host)
-        # remove trailing slash
-        path = o.path.rstrip('/')
-        # remove port if 443
-        netloc = o.netloc
-        if o.port == 443:
-            netloc = netloc.split(':')[0]
-
-        self.host = urllib.parse.urlunparse((o.scheme, netloc, path, o.params, o.query, o.fragment))
-
     def load_azure_tenant_id(self):
         """[Internal] Load the Azure tenant ID from the Azure Databricks login page.
 

databricks/sdk/credentials_provider.py

@@ -197,19 +197,24 @@ def external_browser(cfg: 'Config') -> Optional[CredentialsProvider]:
         client_id = '6128a518-99a9-425b-8333-4cc94f04cacd'
     else:
         raise ValueError(f'local browser SSO is not supported')
-    oauth_client = OAuthClient(host=cfg.host,
-                               client_id=client_id,
-                               redirect_url='http://localhost:8020',
-                               client_secret=cfg.client_secret)
 
     # Load cached credentials from disk if they exist.
     # Note that these are local to the Python SDK and not reused by other SDKs.
-    token_cache = TokenCache(oauth_client)
+    oidc_endpoints = cfg.oidc_endpoints
+    token_cache = TokenCache(host=cfg.host,
+                             oidc_endpoints=oidc_endpoints,
+                             client_id=client_id,
+                             client_secret=cfg.client_secret,
+                             redirect_url='http://localhost:8020')
     credentials = token_cache.load()
     if credentials:
         # Force a refresh in case the loaded credentials are expired.
         credentials.token()
     else:
+        oauth_client = OAuthClient(oidc_endpoints=oidc_endpoints,
+                                   client_id=client_id,
+                                   redirect_url='http://localhost:8020',
+                                   client_secret=cfg.client_secret)
         consent = oauth_client.initiate_consent()
         if not consent:
             return None