databricks-sdk 0.27.1__py3-none-any.whl → 0.29.0__py3-none-any.whl

databricks/sdk/credentials_provider.py

@@ -22,12 +22,26 @@ from .azure import add_sp_management_token, add_workspace_id_header
 from .oauth import (ClientCredentials, OAuthClient, Refreshable, Token,
                     TokenCache, TokenSource)
 
-HeaderFactory = Callable[[], Dict[str, str]]
+CredentialsProvider = Callable[[], Dict[str, str]]
 
 logger = logging.getLogger('databricks.sdk')
 
 
-class CredentialsProvider(abc.ABC):
+class OAuthCredentialsProvider:
+    """ OAuthCredentialsProvider is a type of CredentialsProvider which exposes OAuth tokens. """
+
+    def __init__(self, credentials_provider: CredentialsProvider, token_provider: Callable[[], Token]):
+        self._credentials_provider = credentials_provider
+        self._token_provider = token_provider
+
+    def __call__(self) -> Dict[str, str]:
+        return self._credentials_provider()
+
+    def oauth_token(self) -> Token:
+        return self._token_provider()
+
+
+class CredentialsStrategy(abc.ABC):
     """ CredentialsProvider is the protocol (call-side interface)
      for authenticating requests to Databricks REST APIs"""
 
@@ -36,20 +50,39 @@ class CredentialsProvider(abc.ABC):
         ...
 
     @abc.abstractmethod
-    def __call__(self, cfg: 'Config') -> HeaderFactory:
+    def __call__(self, cfg: 'Config') -> CredentialsProvider:
         ...
 
 
-def credentials_provider(name: str, require: List[str]):
+class OauthCredentialsStrategy(CredentialsStrategy):
+    """ OauthCredentialsProvider is a CredentialsProvider which
+    supports Oauth tokens"""
+
+    def __init__(self, auth_type: str, headers_provider: Callable[['Config'], OAuthCredentialsProvider]):
+        self._headers_provider = headers_provider
+        self._auth_type = auth_type
+
+    def auth_type(self) -> str:
+        return self._auth_type
+
+    def __call__(self, cfg: 'Config') -> OAuthCredentialsProvider:
+        return self._headers_provider(cfg)
+
+    def oauth_token(self, cfg: 'Config') -> Token:
+        return self._headers_provider(cfg).oauth_token()
+
+
+def credentials_strategy(name: str, require: List[str]):
     """ Given the function that receives a Config and returns RequestVisitor,
     create CredentialsProvider with a given name and required configuration
     attribute names to be present for this function to be called. """
 
-    def inner(func: Callable[['Config'], HeaderFactory]) -> CredentialsProvider:
+    def inner(func: Callable[['Config'], CredentialsProvider]) -> CredentialsStrategy:
 
         @functools.wraps(func)
-        def wrapper(cfg: 'Config') -> Optional[HeaderFactory]:
+        def wrapper(cfg: 'Config') -> Optional[CredentialsProvider]:
             for attr in require:
+                getattr(cfg, attr)
                 if not getattr(cfg, attr):
                     return None
             return func(cfg)
@@ -60,8 +93,27 @@ def credentials_provider(name: str, require: List[str]):
     return inner
 
 
-@credentials_provider('basic', ['host', 'username', 'password'])
-def basic_auth(cfg: 'Config') -> HeaderFactory:
+def oauth_credentials_strategy(name: str, require: List[str]):
+    """ Given the function that receives a Config and returns an OauthHeaderFactory,
+    create an OauthCredentialsProvider with a given name and required configuration
+    attribute names to be present for this function to be called. """
+
+    def inner(func: Callable[['Config'], OAuthCredentialsProvider]) -> OauthCredentialsStrategy:
+
+        @functools.wraps(func)
+        def wrapper(cfg: 'Config') -> Optional[OAuthCredentialsProvider]:
+            for attr in require:
+                if not getattr(cfg, attr):
+                    return None
+            return func(cfg)
+
+        return OauthCredentialsStrategy(name, wrapper)
+
+    return inner
+
+
+@credentials_strategy('basic', ['host', 'username', 'password'])
+def basic_auth(cfg: 'Config') -> CredentialsProvider:
     """ Given username and password, add base64-encoded Basic credentials """
     encoded = base64.b64encode(f'{cfg.username}:{cfg.password}'.encode()).decode()
     static_credentials = {'Authorization': f'Basic {encoded}'}
@@ -72,8 +124,8 @@ def basic_auth(cfg: 'Config') -> HeaderFactory:
     return inner
 
 
-@credentials_provider('pat', ['host', 'token'])
-def pat_auth(cfg: 'Config') -> HeaderFactory:
+@credentials_strategy('pat', ['host', 'token'])
+def pat_auth(cfg: 'Config') -> CredentialsProvider:
     """ Adds Databricks Personal Access Token to every request """
     static_credentials = {'Authorization': f'Bearer {cfg.token}'}
 
@@ -83,8 +135,8 @@ def pat_auth(cfg: 'Config') -> HeaderFactory:
     return inner
 
 
-@credentials_provider('runtime', [])
-def runtime_native_auth(cfg: 'Config') -> Optional[HeaderFactory]:
+@credentials_strategy('runtime', [])
+def runtime_native_auth(cfg: 'Config') -> Optional[CredentialsProvider]:
     if 'DATABRICKS_RUNTIME_VERSION' not in os.environ:
         return None
 
@@ -107,8 +159,8 @@ def runtime_native_auth(cfg: 'Config') -> Optional[HeaderFactory]:
     return None
 
 
-@credentials_provider('oauth-m2m', ['host', 'client_id', 'client_secret'])
-def oauth_service_principal(cfg: 'Config') -> Optional[HeaderFactory]:
+@oauth_credentials_strategy('oauth-m2m', ['host', 'client_id', 'client_secret'])
+def oauth_service_principal(cfg: 'Config') -> Optional[CredentialsProvider]:
     """ Adds refreshed Databricks machine-to-machine OAuth Bearer token to every request,
     if /oidc/.well-known/oauth-authorization-server is available on the given host. """
     oidc = cfg.oidc_endpoints
@@ -124,11 +176,14 @@ def oauth_service_principal(cfg: 'Config') -> Optional[HeaderFactory]:
         token = token_source.token()
         return {'Authorization': f'{token.token_type} {token.access_token}'}
 
-    return inner
+    def token() -> Token:
+        return token_source.token()
+
+    return OAuthCredentialsProvider(inner, token)
 
 
-@credentials_provider('external-browser', ['host', 'auth_type'])
-def external_browser(cfg: 'Config') -> Optional[HeaderFactory]:
+@credentials_strategy('external-browser', ['host', 'auth_type'])
+def external_browser(cfg: 'Config') -> Optional[CredentialsProvider]:
     if cfg.auth_type != 'external-browser':
         return None
     if cfg.client_id:
@@ -178,9 +233,9 @@ def _ensure_host_present(cfg: 'Config', token_source_for: Callable[[str], TokenS
     cfg.host = f"https://{resp.json()['properties']['workspaceUrl']}"
 
 
-@credentials_provider('azure-client-secret',
-                      ['is_azure', 'azure_client_id', 'azure_client_secret', 'azure_tenant_id'])
-def azure_service_principal(cfg: 'Config') -> HeaderFactory:
+@oauth_credentials_strategy('azure-client-secret',
+                            ['is_azure', 'azure_client_id', 'azure_client_secret', 'azure_tenant_id'])
+def azure_service_principal(cfg: 'Config') -> CredentialsProvider:
     """ Adds refreshed Azure Active Directory (AAD) Service Principal OAuth tokens
     to every request, while automatically resolving different Azure environment endpoints. """
 
@@ -203,11 +258,14 @@ def azure_service_principal(cfg: 'Config') -> HeaderFactory:
         add_sp_management_token(cloud, headers)
         return headers
 
-    return refreshed_headers
+    def token() -> Token:
+        return inner.token()
+
+    return OAuthCredentialsProvider(refreshed_headers, token)
 
 
-@credentials_provider('github-oidc-azure', ['host', 'azure_client_id'])
-def github_oidc_azure(cfg: 'Config') -> Optional[HeaderFactory]:
+@oauth_credentials_strategy('github-oidc-azure', ['host', 'azure_client_id'])
+def github_oidc_azure(cfg: 'Config') -> Optional[CredentialsProvider]:
     if 'ACTIONS_ID_TOKEN_REQUEST_TOKEN' not in os.environ:
         # not in GitHub actions
         return None
@@ -250,14 +308,17 @@ def github_oidc_azure(cfg: 'Config') -> Optional[HeaderFactory]:
         token = inner.token()
         return {'Authorization': f'{token.token_type} {token.access_token}'}
 
-    return refreshed_headers
+    def token() -> Token:
+        return inner.token()
+
+    return OAuthCredentialsProvider(refreshed_headers, token)
 
 
 GcpScopes = ["https://www.googleapis.com/auth/cloud-platform", "https://www.googleapis.com/auth/compute"]
 
 
-@credentials_provider('google-credentials', ['host', 'google_credentials'])
-def google_credentials(cfg: 'Config') -> Optional[HeaderFactory]:
+@oauth_credentials_strategy('google-credentials', ['host', 'google_credentials'])
+def google_credentials(cfg: 'Config') -> Optional[CredentialsProvider]:
     if not cfg.is_gcp:
         return None
     # Reads credentials as JSON. Credentials can be either a path to JSON file, or actual JSON string.
@@ -277,6 +338,10 @@ def google_credentials(cfg: 'Config') -> Optional[HeaderFactory]:
     gcp_credentials = service_account.Credentials.from_service_account_info(info=account_info,
                                                                             scopes=GcpScopes)
 
+    def token() -> Token:
+        credentials.refresh(request)
+        return credentials.token
+
     def refreshed_headers() -> Dict[str, str]:
         credentials.refresh(request)
         headers = {'Authorization': f'Bearer {credentials.token}'}
@@ -285,11 +350,11 @@ def google_credentials(cfg: 'Config') -> Optional[HeaderFactory]:
             headers["X-Databricks-GCP-SA-Access-Token"] = gcp_credentials.token
         return headers
 
-    return refreshed_headers
+    return OAuthCredentialsProvider(refreshed_headers, token)
 
 
-@credentials_provider('google-id', ['host', 'google_service_account'])
-def google_id(cfg: 'Config') -> Optional[HeaderFactory]:
+@oauth_credentials_strategy('google-id', ['host', 'google_service_account'])
+def google_id(cfg: 'Config') -> Optional[CredentialsProvider]:
     if not cfg.is_gcp:
         return None
     credentials, _project_id = google.auth.default()
@@ -309,6 +374,10 @@ def google_id(cfg: 'Config') -> Optional[HeaderFactory]:
 
     request = Request()
 
+    def token() -> Token:
+        id_creds.refresh(request)
+        return id_creds.token
+
     def refreshed_headers() -> Dict[str, str]:
         id_creds.refresh(request)
         headers = {'Authorization': f'Bearer {id_creds.token}'}
@@ -317,7 +386,7 @@ def google_id(cfg: 'Config') -> Optional[HeaderFactory]:
             headers["X-Databricks-GCP-SA-Access-Token"] = gcp_impersonated_credentials.token
         return headers
 
-    return refreshed_headers
+    return OAuthCredentialsProvider(refreshed_headers, token)
 
 
 class CliTokenSource(Refreshable):
@@ -422,8 +491,8 @@ class AzureCliTokenSource(CliTokenSource):
         return components[2]
 
 
-@credentials_provider('azure-cli', ['is_azure'])
-def azure_cli(cfg: 'Config') -> Optional[HeaderFactory]:
+@credentials_strategy('azure-cli', ['is_azure'])
+def azure_cli(cfg: 'Config') -> Optional[CredentialsProvider]:
     """ Adds refreshed OAuth token granted by `az login` command to every request. """
     token_source = None
     mgmt_token_source = None
@@ -516,8 +585,8 @@ class DatabricksCliTokenSource(CliTokenSource):
         raise err
 
 
-@credentials_provider('databricks-cli', ['host'])
-def databricks_cli(cfg: 'Config') -> Optional[HeaderFactory]:
+@oauth_credentials_strategy('databricks-cli', ['host'])
+def databricks_cli(cfg: 'Config') -> Optional[CredentialsProvider]:
     try:
         token_source = DatabricksCliTokenSource(cfg)
     except FileNotFoundError as e:
@@ -538,7 +607,7 @@ def databricks_cli(cfg: 'Config') -> Optional[HeaderFactory]:
         token = token_source.token()
         return {'Authorization': f'{token.token_type} {token.access_token}'}
 
-    return inner
+    return OAuthCredentialsProvider(inner, token_source.token)
 
 
 class MetadataServiceTokenSource(Refreshable):
@@ -577,8 +646,8 @@ class MetadataServiceTokenSource(Refreshable):
         return Token(access_token=access_token, token_type=token_type, expiry=expiry)
 
 
-@credentials_provider('metadata-service', ['host', 'metadata_service_url'])
-def metadata_service(cfg: 'Config') -> Optional[HeaderFactory]:
+@credentials_strategy('metadata-service', ['host', 'metadata_service_url'])
+def metadata_service(cfg: 'Config') -> Optional[CredentialsProvider]:
     """ Adds refreshed token granted by Databricks Metadata Service to every request. """
 
     token_source = MetadataServiceTokenSource(cfg)
@@ -597,17 +666,25 @@ class DefaultCredentials:
 
     def __init__(self) -> None:
         self._auth_type = 'default'
-
-    def auth_type(self) -> str:
-        return self._auth_type
-
-    def __call__(self, cfg: 'Config') -> HeaderFactory:
-        auth_providers = [
+        self._auth_providers = [
             pat_auth, basic_auth, metadata_service, oauth_service_principal, azure_service_principal,
             github_oidc_azure, azure_cli, external_browser, databricks_cli, runtime_native_auth,
             google_credentials, google_id
         ]
-        for provider in auth_providers:
+
+    def auth_type(self) -> str:
+        return self._auth_type
+
+    def oauth_token(self, cfg: 'Config') -> Token:
+        for provider in self._auth_providers:
+            auth_type = provider.auth_type()
+            if auth_type != self._auth_type:
+                # ignore other auth types if they don't match the selected one
+                continue
+            return provider.oauth_token(cfg)
+
+    def __call__(self, cfg: 'Config') -> CredentialsProvider:
+        for provider in self._auth_providers:
             auth_type = provider.auth_type()
             if cfg.auth_type and auth_type != cfg.auth_type:
                 # ignore other auth types if one is explicitly enforced

databricks/sdk/dbutils.py

@@ -1,10 +1,11 @@
 import base64
 import json
 import logging
-import os.path
+import os
 import threading
 from collections import namedtuple
-from typing import Callable, Dict, List
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional
 
 from .core import ApiClient, Config, DatabricksError
 from .mixins import compute as compute_ext
@@ -240,6 +241,76 @@ class RemoteDbUtils:
                           name=util)
 
 
+@dataclass
+class OverrideResult:
+    result: Any
+
+
+def get_local_notebook_path():
+    value = os.getenv("DATABRICKS_SOURCE_FILE")
+    if value is None:
+        raise ValueError(
+            "Getting the current notebook path is only supported when running a notebook using the `Databricks Connect: Run as File` or `Databricks Connect: Debug as File` commands in the Databricks extension for VS Code. To bypass this error, set environment variable `DATABRICKS_SOURCE_FILE` to the desired notebook path."
+        )
+
+    return value
+
+
+class _OverrideProxyUtil:
+
+    @classmethod
+    def new(cls, path: str):
+        if len(cls.__get_matching_overrides(path)) > 0:
+            return _OverrideProxyUtil(path)
+        return None
+
+    def __init__(self, name: str):
+        self._name = name
+
+    # These are the paths that we want to override and not send to remote dbutils. NOTE, for each of these paths, no prefixes
+    # are sent to remote either. This could lead to unintentional breakage.
+    # Our current proxy implementation (which sends everything to remote dbutils) uses `{util}.{method}(*args, **kwargs)` ONLY.
+    # This means, it is completely safe to override paths starting with `{util}.{attribute}.<other_parts>`, since none of the prefixes
+    # are being proxied to remote dbutils currently.
+    proxy_override_paths = {
+        'notebook.entry_point.getDbutils().notebook().getContext().notebookPath().get()':
+        get_local_notebook_path,
+    }
+
+    @classmethod
+    def __get_matching_overrides(cls, path: str):
+        return [x for x in cls.proxy_override_paths.keys() if x.startswith(path)]
+
+    def __run_override(self, path: str) -> Optional[OverrideResult]:
+        overrides = self.__get_matching_overrides(path)
+        if len(overrides) == 1 and overrides[0] == path:
+            return OverrideResult(self.proxy_override_paths[overrides[0]]())
+
+        if len(overrides) > 0:
+            return OverrideResult(_OverrideProxyUtil(name=path))
+
+        return None
+
+    def __call__(self, *args, **kwds) -> Any:
+        if len(args) != 0 or len(kwds) != 0:
+            raise TypeError(
+                f"Arguments are not supported for overridden method {self._name}. Invoke as: {self._name}()")
+
+        callable_path = f"{self._name}()"
+        result = self.__run_override(callable_path)
+        if result:
+            return result.result
+
+        raise TypeError(f"{self._name} is not callable")
+
+    def __getattr__(self, method: str) -> Any:
+        result = self.__run_override(f"{self._name}.{method}")
+        if result:
+            return result.result
+
+        raise AttributeError(f"module {self._name} has no attribute {method}")
+
+
 class _ProxyUtil:
     """Enables temporary workaround to call remote in-REPL dbutils without having to re-implement them"""
 
@@ -250,7 +321,14 @@ class _ProxyUtil:
         self._context_factory = context_factory
         self._name = name
 
-    def __getattr__(self, method: str) -> '_ProxyCall':
+    def __call__(self):
+        raise NotImplementedError(f"dbutils.{self._name} is not callable")
+
+    def __getattr__(self, method: str) -> '_ProxyCall | _ProxyUtil | _OverrideProxyUtil':
+        override = _OverrideProxyUtil.new(f"{self._name}.{method}")
+        if override:
+            return override
+
         return _ProxyCall(command_execution=self._commands,
                           cluster_id=self._cluster_id,
                           context_factory=self._context_factory,

databricks/sdk/environments.py

@@ -2,7 +2,31 @@ from dataclasses import dataclass
 from enum import Enum
 from typing import Optional
 
-from .azure import ARM_DATABRICKS_RESOURCE_ID, ENVIRONMENTS, AzureEnvironment
+
+@dataclass
+class AzureEnvironment:
+    name: str
+    service_management_endpoint: str
+    resource_manager_endpoint: str
+    active_directory_endpoint: str
+
+
+ARM_DATABRICKS_RESOURCE_ID = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
+
+ENVIRONMENTS = dict(
+    PUBLIC=AzureEnvironment(name="PUBLIC",
+                            service_management_endpoint="https://management.core.windows.net/",
+                            resource_manager_endpoint="https://management.azure.com/",
+                            active_directory_endpoint="https://login.microsoftonline.com/"),
+    USGOVERNMENT=AzureEnvironment(name="USGOVERNMENT",
+                                  service_management_endpoint="https://management.core.usgovcloudapi.net/",
+                                  resource_manager_endpoint="https://management.usgovcloudapi.net/",
+                                  active_directory_endpoint="https://login.microsoftonline.us/"),
+    CHINA=AzureEnvironment(name="CHINA",
+                           service_management_endpoint="https://management.core.chinacloudapi.cn/",
+                           resource_manager_endpoint="https://management.chinacloudapi.cn/",
+                           active_directory_endpoint="https://login.chinacloudapi.cn/"),
+)
 
 
 class Cloud(Enum):
@@ -70,3 +94,12 @@ ALL_ENVS = [
     DatabricksEnvironment(Cloud.GCP, ".staging.gcp.databricks.com"),
     DatabricksEnvironment(Cloud.GCP, ".gcp.databricks.com")
 ]
+
+
+def get_environment_for_hostname(hostname: str) -> DatabricksEnvironment:
+    if not hostname:
+        return DEFAULT_ENVIRONMENT
+    for env in ALL_ENVS:
+        if hostname.endswith(env.dns_zone):
+            return env
+    return DEFAULT_ENVIRONMENT

databricks/sdk/errors/__init__.py

@@ -1,4 +1,5 @@
 from .base import DatabricksError, ErrorDetail
 from .mapper import error_mapper
 from .platform import *
+from .private_link import PrivateLinkValidationError
 from .sdk import *

databricks/sdk/errors/mapper.py

@@ -4,6 +4,8 @@ from databricks.sdk.errors import platform
 from databricks.sdk.errors.base import DatabricksError
 
 from .overrides import _ALL_OVERRIDES
+from .private_link import (_get_private_link_validation_error,
+                           _is_private_link_redirect)
 
 
 def error_mapper(response: requests.Response, raw: dict) -> DatabricksError:
@@ -21,6 +23,8 @@ def error_mapper(response: requests.Response, raw: dict) -> DatabricksError:
         # where there's a default exception class per HTTP status code, and we do
         # rely on Databricks platform exception mapper to do the right thing.
         return platform.STATUS_CODE_MAPPING[status_code](**raw)
+    if _is_private_link_redirect(response):
+        return _get_private_link_validation_error(response.url)
 
     # backwards-compatible error creation for cases like using older versions of
     # the SDK on way never releases of the platform.

databricks/sdk/errors/private_link.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass
+from urllib import parse
+
+import requests
+
+from ..environments import Cloud, get_environment_for_hostname
+from .platform import PermissionDenied
+
+
+@dataclass
+class _PrivateLinkInfo:
+    serviceName: str
+    endpointName: str
+    referencePage: str
+
+    def error_message(self):
+        return (
+            f'The requested workspace has {self.serviceName} enabled and is not accessible from the current network. '
+            f'Ensure that {self.serviceName} is properly configured and that your device has access to the '
+            f'{self.endpointName}. For more information, see {self.referencePage}.')
+
+
+_private_link_info_map = {
+    Cloud.AWS:
+    _PrivateLinkInfo(serviceName='AWS PrivateLink',
+                     endpointName='AWS VPC endpoint',
+                     referencePage='https://docs.databricks.com/en/security/network/classic/privatelink.html',
+                     ),
+    Cloud.AZURE:
+    _PrivateLinkInfo(
+        serviceName='Azure Private Link',
+        endpointName='Azure Private Link endpoint',
+        referencePage='https://learn.microsoft.com/en-us/azure/databricks/security/network/classic/private-link-standard#authentication-troubleshooting',
+    ),
+    Cloud.GCP:
+    _PrivateLinkInfo(
+        serviceName='Private Service Connect',
+        endpointName='GCP VPC endpoint',
+        referencePage='https://docs.gcp.databricks.com/en/security/network/classic/private-service-connect.html',
+    )
+}
+
+
+class PrivateLinkValidationError(PermissionDenied):
+    """Raised when a user tries to access a Private Link-enabled workspace, but the user's network does not have access
+    to the workspace."""
+
+
+def _is_private_link_redirect(resp: requests.Response) -> bool:
+    parsed = parse.urlparse(resp.url)
+    return parsed.path == '/login.html' and 'error=private-link-validation-error' in parsed.query
+
+
+def _get_private_link_validation_error(url: str) -> _PrivateLinkInfo:
+    parsed = parse.urlparse(url)
+    env = get_environment_for_hostname(parsed.hostname)
+    return PrivateLinkValidationError(message=_private_link_info_map[env.cloud].error_message(),
+                                      error_code='PRIVATE_LINK_VALIDATION_ERROR',
+                                      status_code=403,
+                                      )

databricks/sdk/oauth.py

@@ -21,6 +21,10 @@ import requests.auth
 # See https://stackoverflow.com/a/75466778/277035 for more info
 NO_ORIGIN_FOR_SPA_CLIENT_ERROR = 'AADSTS9002327'
 
+URL_ENCODED_CONTENT_TYPE = "application/x-www-form-urlencoded"
+JWT_BEARER_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"
+OIDC_TOKEN_PATH = "/oidc/v1/token"
+
 logger = logging.getLogger(__name__)
 
 
@@ -358,18 +362,15 @@ class OAuthClient:
                  client_secret: str = None):
         # TODO: is it a circular dependency?..
         from .core import Config
-        from .credentials_provider import credentials_provider
+        from .credentials_provider import credentials_strategy
 
-        @credentials_provider('noop', [])
+        @credentials_strategy('noop', [])
         def noop_credentials(_: any):
             return lambda: {}
 
-        config = Config(host=host, credentials_provider=noop_credentials)
+        config = Config(host=host, credentials_strategy=noop_credentials)
         if not scopes:
             scopes = ['all-apis']
-        if config.is_azure:
-            # Azure AD only supports full access to Azure Databricks.
-            scopes = [f'{config.effective_azure_login_app_id}/user_impersonation', 'offline_access']
         oidc = config.oidc_endpoints
         if not oidc:
             raise ValueError(f'{host} does not support OAuth')
@@ -381,6 +382,7 @@ class OAuthClient:
         self.token_url = oidc.token_endpoint
         self.is_aws = config.is_aws
         self.is_azure = config.is_azure
+        self.is_gcp = config.is_gcp
 
         self._auth_url = oidc.authorization_endpoint
         self._scopes = scopes