import abc

import base64

import functools

import io

import json

import logging

import os

import pathlib

import platform

import subprocess

import sys

import threading

import time

from datetime import datetime

from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import google.auth # type: ignore

import requests

from google.auth import impersonated_credentials # type: ignore

from google.auth.transport.requests import Request # type: ignore

from google.oauth2 import service_account # type: ignore

from databricks.sdk.oauth import get_azure_entra_id_workspace_endpoints

from . import azure, oauth, oidc, oidc_token_supplier

from .client_types import ClientType

CredentialsProvider = Callable[[], Dict[str, str]]

logger = logging.getLogger("databricks.sdk")

class OAuthCredentialsProvider:

"""OAuthCredentialsProvider is a type of CredentialsProvider which exposes OAuth tokens."""

def __init__(

self,

credentials_provider: CredentialsProvider,

token_provider: Callable[[], oauth.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) -> oauth.Token:

return self._token_provider()

class CredentialsStrategy(abc.ABC):

"""CredentialsProvider is the protocol (call-side interface)

for authenticating requests to Databricks REST APIs"""

@abc.abstractmethod

def auth_type(self) -> str: ...

@abc.abstractmethod

def __call__(self, cfg: "Config") -> CredentialsProvider: ...

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") -> oauth.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"], CredentialsProvider],

) -> CredentialsStrategy:

@functools.wraps(func)

def wrapper(cfg: "Config") -> Optional[CredentialsProvider]:

for attr in require:

if not getattr(cfg, attr):

return None

return func(cfg)

wrapper.auth_type = lambda: name

return wrapper

return inner

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.

Args:

name: The name of the authentication strategy

require: List of config attributes that must be present

"""

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}"}

def inner() -> Dict[str, str]:

return static_credentials

return inner

@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}"}

def inner() -> Dict[str, str]:

return static_credentials

return inner

@credentials_strategy("runtime", [])

def runtime_native_auth(cfg: "Config") -> Optional[CredentialsProvider]:

if "DATABRICKS_RUNTIME_VERSION" not in os.environ:

return None

# 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_native_unified,

init_runtime_repl_auth)

# Try the unified provider first (returns host, account_id, workspace_id, inner).

if init_runtime_native_unified is not None:

host, account_id, workspace_id, inner = init_runtime_native_unified()

if host is not None:

cfg.host = host

cfg.account_id = account_id

cfg.workspace_id = workspace_id

logger.debug("[init_runtime_native_unified] runtime native auth configured")

return inner

logger.debug("[init_runtime_native_unified] no host detected")

# Fall back to legacy providers (return host, inner).

for init in [

init_runtime_native_auth,

init_runtime_repl_auth,

init_runtime_legacy_auth,

]:

if init is None:

continue

host, inner = init()

if host is None:

logger.debug(f"[{init.__name__}] no host detected")

continue

cfg.host = host

logger.debug(f"[{init.__name__}] runtime native auth configured")

return inner

return None

@oauth_credentials_strategy("runtime-oauth", ["scopes"])

def runtime_oauth(cfg: "Config") -> Optional[CredentialsProvider]:

if "DATABRICKS_RUNTIME_VERSION" not in os.environ:

return None

def get_notebook_pat_token() -> Optional[str]:

native_auth = runtime_native_auth(cfg)

if native_auth is None:

return None

notebook_pat_token = None

notebook_pat_authorization = native_auth().get("Authorization", "").strip()

if notebook_pat_authorization.lower().startswith("bearer "):

notebook_pat_token = notebook_pat_authorization[len("bearer ") :].strip()

return notebook_pat_token

notebook_pat_token = get_notebook_pat_token()

if notebook_pat_token is None:

return None

token_source = oauth.PATOAuthTokenExchange(

get_original_token=get_notebook_pat_token,

host=cfg.host,

scopes=cfg.get_scopes_as_string(),

authorization_details=cfg.authorization_details,

)

def inner() -> Dict[str, str]:

token = token_source.token()

return {"Authorization": f"{token.token_type} {token.access_token}"}

def token() -> oauth.Token:

return token_source.token()

return OAuthCredentialsProvider(inner, token)

@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.databricks_oidc_endpoints

if oidc is None:

return None

token_source = oauth.ClientCredentials(

client_id=cfg.client_id,

client_secret=cfg.client_secret,

token_url=oidc.token_endpoint,

scopes=cfg.get_scopes_as_string(),

use_header=True,

disable_async=cfg.disable_async_token_refresh,

authorization_details=cfg.authorization_details,

)

def inner() -> Dict[str, str]:

token = token_source.token()

return {"Authorization": f"{token.token_type} {token.access_token}"}

def token() -> oauth.Token:

return token_source.token()

return OAuthCredentialsProvider(inner, token)

@credentials_strategy("external-browser", ["host", "auth_type"])

def external_browser(cfg: "Config") -> Optional[CredentialsProvider]:

if cfg.auth_type != "external-browser":

return None

client_id, client_secret = None, None

oidc_endpoints = None

if cfg.client_id:

client_id = cfg.client_id

client_secret = cfg.client_secret

oidc_endpoints = cfg.databricks_oidc_endpoints

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)

if not client_id:

client_id = "databricks-cli"

oidc_endpoints = cfg.databricks_oidc_endpoints

if not oidc_endpoints:

return None

scopes = cfg.get_scopes()

if not cfg.disable_oauth_refresh_token:

if "offline_access" not in scopes:

scopes = scopes + ["offline_access"]

# Load cached credentials from disk if they exist. Note that these are

# local to the Python SDK and not reused by other SDKs.

redirect_url = "http://localhost:8020"

token_cache = oauth.TokenCache(

host=cfg.host,

oidc_endpoints=oidc_endpoints,

client_id=client_id,

client_secret=client_secret,

redirect_url=redirect_url,

scopes=scopes,

profile=cfg.profile,

)

credentials = token_cache.load()

if credentials:

try:

# Pro-actively refresh the loaded credentials. This is done

# to detect if the token is expired and needs to be refreshed

# by going through the OAuth login flow.

credentials.token()

return credentials(cfg)

# TODO: We should ideally use more specific exceptions.

except Exception as e:

logger.warning(f"Failed to refresh cached token: {e}. Initiating new OAuth login flow")

oauth_client = oauth.OAuthClient(

oidc_endpoints=oidc_endpoints,

client_id=client_id,

redirect_url=redirect_url,

client_secret=client_secret,

scopes=scopes,

)

consent = oauth_client.initiate_consent()

if not consent:

return None

credentials = consent.launch_external_browser()

token_cache.save(credentials)

return credentials(cfg)

def _ensure_host_present(cfg: "Config", token_source_for: Callable[[str], oauth.TokenSource]):

"""Resolves Azure Databricks workspace URL from ARM Resource ID"""

if cfg.host:

return

if not cfg.azure_workspace_resource_id:

return

arm = cfg.arm_environment.resource_manager_endpoint

token = token_source_for(arm).token()

resp = requests.get(

f"{arm}{cfg.azure_workspace_resource_id}?api-version=2018-04-01",

headers={"Authorization": f"Bearer {token.access_token}"},

)

if not resp.ok:

raise ValueError(f"Cannot resolve Azure Databricks workspace: {resp.content}")

cfg.host = f"https://{resp.json()['properties']['workspaceUrl']}"

@oauth_credentials_strategy(

"azure-client-secret",

["azure_client_id", "azure_client_secret"],

)

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.

"""

def token_source_for(resource: str) -> oauth.TokenSource:

aad_endpoint = cfg.arm_environment.active_directory_endpoint

return oauth.ClientCredentials(

client_id=cfg.azure_client_id,

client_secret=cfg.azure_client_secret,

token_url=f"{aad_endpoint}{cfg.azure_tenant_id}/oauth2/token",

endpoint_params={"resource": resource},

use_params=True,

disable_async=cfg.disable_async_token_refresh,

scopes=cfg.get_scopes_as_string(),

authorization_details=cfg.authorization_details,

)

_ensure_host_present(cfg, token_source_for)

cfg.load_azure_tenant_id()

logger.info("Configured AAD token for Service Principal (%s)", cfg.azure_client_id)

inner = token_source_for(cfg.effective_azure_login_app_id)

cloud = token_source_for(cfg.arm_environment.service_management_endpoint)

def refreshed_headers() -> Dict[str, str]:

headers = {

"Authorization": f"Bearer {inner.token().access_token}",

}

azure.add_workspace_id_header(cfg, headers)

azure.add_sp_management_token(cloud, headers)

return headers

def token() -> oauth.Token:

return inner.token()

return OAuthCredentialsProvider(refreshed_headers, token)

@credentials_strategy("env-oidc", ["host"])

def env_oidc(cfg) -> Optional[CredentialsProvider]:

# Search for an OIDC ID token in DATABRICKS_OIDC_TOKEN environment variable

# by default. This can be overridden by setting DATABRICKS_OIDC_TOKEN_ENV

# to the name of an environment variable that contains the OIDC ID token.

env_var = "DATABRICKS_OIDC_TOKEN"

if cfg.oidc_token_env:

env_var = cfg.oidc_token_env

return oidc_credentials_provider(cfg, oidc.EnvIdTokenSource(env_var))

@credentials_strategy("file-oidc", ["host", "oidc_token_filepath"])

def file_oidc(cfg) -> Optional[CredentialsProvider]:

return oidc_credentials_provider(cfg, oidc.FileIdTokenSource(cfg.oidc_token_filepath))

def oidc_credentials_provider(cfg, id_token_source: oidc.IdTokenSource) -> Optional[CredentialsProvider]:

"""Creates a CredentialsProvider to sign requests with an OAuth token obtained

by automatically performing the token exchange using the given IdTokenSource."""

try:

id_token_source.id_token() # validate the id_token_source

except Exception as e:

logger.debug(f"Failed to get OIDC token: {e}")

return None

token_source = oidc.DatabricksOidcTokenSource(

host=cfg.host,

token_endpoint=cfg.databricks_oidc_endpoints.token_endpoint,

client_id=cfg.client_id,

account_id=cfg.account_id,

id_token_source=id_token_source,

disable_async=cfg.disable_async_token_refresh,

scopes=cfg.get_scopes_as_string(),

)

def refreshed_headers() -> Dict[str, str]:

token = token_source.token()

return {"Authorization": f"{token.token_type} {token.access_token}"}

def token() -> oauth.Token:

return token_source.token()

return OAuthCredentialsProvider(refreshed_headers, token)

def _oidc_credentials_provider(

cfg: "Config", supplier_factory: Callable[[], Any], provider_name: str

) -> Optional[CredentialsProvider]:

"""

Generic OIDC credentials provider that works with any OIDC token supplier.

Args:

cfg: Databricks configuration

supplier_factory: Callable that returns an OIDC token supplier instance

provider_name: Human-readable name (e.g., "GitHub OIDC", "Azure DevOps OIDC")

Returns:

OAuthCredentialsProvider if successful, None if supplier unavailable or token retrieval fails

"""

# Try to create the supplier

try:

supplier = supplier_factory()

except Exception as e:

logger.debug(f"{provider_name}: {str(e)}")

return None

# Determine the audience for token exchange

audience = cfg.token_audience

if audience is None:

audience = cfg.databricks_oidc_endpoints.token_endpoint

# Try to get an OIDC token. If no supplier returns a token, we cannot use this authentication mode.

id_token = supplier.get_oidc_token(audience)

if not id_token:

logger.debug(f"{provider_name}: no token available, skipping authentication method")

return None

logger.info(f"Configured {provider_name} authentication")

def token_source_for(audience: str) -> oauth.TokenSource:

id_token = supplier.get_oidc_token(audience)

if not id_token:

# Should not happen, since we checked it above.

raise Exception(f"Cannot get {provider_name} token")

return oauth.ClientCredentials(

client_id=cfg.client_id,

client_secret="", # we have no (rotatable) secrets in OIDC flow

token_url=cfg.databricks_oidc_endpoints.token_endpoint,

endpoint_params={

"subject_token_type": "urn:ietf:params:oauth:token-type:jwt",

"subject_token": id_token,

"grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",

},

scopes=cfg.get_scopes_as_string(),

use_params=True,

disable_async=cfg.disable_async_token_refresh,

authorization_details=cfg.authorization_details,

)

def refreshed_headers() -> Dict[str, str]:

token = token_source_for(audience).token()

return {"Authorization": f"{token.token_type} {token.access_token}"}

def token() -> oauth.Token:

return token_source_for(audience).token()

return OAuthCredentialsProvider(refreshed_headers, token)

@oauth_credentials_strategy("github-oidc", ["host", "client_id"])

def github_oidc(cfg: "Config") -> Optional[CredentialsProvider]:

"""

GitHub OIDC authentication uses a Token Supplier to get a JWT Token and exchanges

it for a Databricks Token.

Supported in GitHub Actions with OIDC service connections.

"""

return _oidc_credentials_provider(

cfg=cfg,

supplier_factory=lambda: oidc_token_supplier.GitHubOIDCTokenSupplier(),

provider_name="GitHub OIDC",

)

@oauth_credentials_strategy("azure-devops-oidc", ["host", "client_id"])

def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]:

"""

Azure DevOps OIDC authentication uses a Token Supplier to get a JWT Token

and exchanges it for a Databricks Token.

Supported in Azure DevOps pipelines with OIDC service connections.

"""

return _oidc_credentials_provider(

cfg=cfg,

supplier_factory=lambda: oidc_token_supplier.AzureDevOpsOIDCTokenSupplier(),

provider_name="Azure DevOps OIDC",

)

# Azure Client ID is the minimal thing we need, as otherwise we get AADSTS700016: Application with

# identifier 'https://token.actions.githubusercontent.com' was not found in the directory '...'.

@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

token = oidc_token_supplier.GitHubOIDCTokenSupplier().get_oidc_token("api://AzureADTokenExchange")

if not token:

return None

logger.info(

"Configured AAD token for GitHub Actions OIDC (%s)",

cfg.azure_client_id,

)

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

cfg.azure_tenant_id = token_endpoint.replace(aad_endpoint, "").split("/")[0]

inner = oauth.ClientCredentials(

client_id=cfg.azure_client_id,

client_secret="", # we have no (rotatable) secrets in OIDC flow

token_url=f"{aad_endpoint}{cfg.azure_tenant_id}/oauth2/token",

endpoint_params={

"client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",

"resource": cfg.effective_azure_login_app_id,

"client_assertion": token,

},

use_params=True,

disable_async=cfg.disable_async_token_refresh,

scopes=cfg.get_scopes_as_string(),

authorization_details=cfg.authorization_details,

)

def refreshed_headers() -> Dict[str, str]:

token = inner.token()

return {"Authorization": f"{token.token_type} {token.access_token}"}

def token() -> oauth.Token:

return inner.token()

return OAuthCredentialsProvider(refreshed_headers, token)

GcpScopes = [

"https://www.googleapis.com/auth/cloud-platform",

"https://www.googleapis.com/auth/compute",

]

@oauth_credentials_strategy("google-credentials", ["host", "google_credentials"])

def google_credentials(cfg: "Config") -> Optional[CredentialsProvider]:

# Reads credentials as JSON. Credentials can be either a path to JSON file, or actual JSON string.

# Obtain the id token by providing the json file path and target audience.

if os.path.isfile(cfg.google_credentials):

with io.open(cfg.google_credentials, "r", encoding="utf-8") as json_file:

account_info = json.load(json_file)

else:

# If the file doesn't exist, assume that the config is the actual JSON content.

account_info = json.loads(cfg.google_credentials)

credentials = service_account.IDTokenCredentials.from_service_account_info(

info=account_info, target_audience=cfg.host

)

request = Request()

gcp_credentials = service_account.Credentials.from_service_account_info(info=account_info, scopes=GcpScopes)

def token() -> oauth.Token:

credentials.refresh(request)

return credentials.token

def refreshed_headers() -> Dict[str, str]:

credentials.refresh(request)

headers = {"Authorization": f"Bearer {credentials.token}"}

# GCP SA Access token is only required for specific account level operations.

# It is possible that a user does not have persomissions to mint the GCP SA access token,

# but this is not a blocking error at this point.

try:

gcp_credentials.refresh(request)

headers["X-Databricks-GCP-SA-Access-Token"] = gcp_credentials.token

except Exception as e:

logger.warning(f"Failed to refresh GCP credentials: {e}")

return headers

return OAuthCredentialsProvider(refreshed_headers, token)

@oauth_credentials_strategy("google-id", ["host", "google_service_account"])

def google_id(cfg: "Config") -> Optional[CredentialsProvider]:

credentials, _project_id = google.auth.default()

# Create the impersonated credential.

target_credentials = impersonated_credentials.Credentials(

source_credentials=credentials,

target_principal=cfg.google_service_account,

target_scopes=[],

)

# Set the impersonated credential, target audience and token options.

id_creds = impersonated_credentials.IDTokenCredentials(

target_credentials, target_audience=cfg.host, include_email=True

)

gcp_impersonated_credentials = impersonated_credentials.Credentials(

source_credentials=credentials,

target_principal=cfg.google_service_account,

target_scopes=GcpScopes,

)

request = Request()

def token() -> oauth.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}"}

# GCP SA Access token is only required for specific account level operations.

# It is possible that a user does not have persomissions to mint the GCP SA access token,

# but this is not a blocking error at this point.

try:

gcp_impersonated_credentials.refresh(request)

headers["X-Databricks-GCP-SA-Access-Token"] = gcp_impersonated_credentials.token

except Exception as e:

logger.warning(f"Failed to refresh GCP impersonated credentials: {e}")

return headers

return OAuthCredentialsProvider(refreshed_headers, token)

class CliTokenSource(oauth.Refreshable):

def __init__(

self,

cmd: List[str],

token_type_field: str,

access_token_field: str,

expiry_field: str,

disable_async: bool = True,

fallback_cmd: Optional[List[str]] = None,

):

super().__init__(disable_async=disable_async)

self._cmd = cmd

# fallback_cmd is tried when the primary command fails with "unknown flag: --profile",

# indicating the CLI is too old to support --profile. Can be removed once support

# for CLI versions predating --profile is dropped.

# See: https://github.com/databricks/databricks-sdk-go/pull/1497

self._fallback_cmd = fallback_cmd

self._token_type_field = token_type_field

self._access_token_field = access_token_field

self._expiry_field = expiry_field

@staticmethod

def _parse_expiry(expiry: str) -> datetime:

expiry = expiry.rstrip("Z").split(".")[0]

for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%S"):

try:

return datetime.strptime(expiry, fmt)

except ValueError as e:

last_e = e

if last_e:

raise last_e

def _exec_cli_command(self, cmd: List[str]) -> oauth.Token:

try:

out = _run_subprocess(cmd, capture_output=True, check=True)

it = json.loads(out.stdout.decode())

expires_on = self._parse_expiry(it[self._expiry_field])

return oauth.Token(

access_token=it[self._access_token_field],

token_type=it[self._token_type_field],

expiry=expires_on,

)

except ValueError as e:

raise ValueError(f"cannot unmarshal CLI result: {e}")

except subprocess.CalledProcessError as e:

stdout = e.stdout.decode().strip()

stderr = e.stderr.decode().strip()

message = "\n".join(filter(None, [stdout, stderr]))

raise IOError(f"cannot get access token: {message}") from e

def refresh(self) -> oauth.Token:

try:

return self._exec_cli_command(self._cmd)

except IOError as e:

if self._fallback_cmd is not None and "unknown flag: --profile" in str(e):

logger.warning(

"Databricks CLI does not support --profile flag. Falling back to --host. "

"Please upgrade your CLI to the latest version."

)

return self._exec_cli_command(self._fallback_cmd)

raise

def _run_subprocess(

popenargs,

input=None,

capture_output=True,

timeout=None,

check=False,

**kwargs,

) -> subprocess.CompletedProcess:

"""Runs subprocess with given arguments.

This handles OS-specific modifications that need to be made to the invocation of subprocess.run.

"""

kwargs["shell"] = sys.platform.startswith("win")

# windows requires shell=True to be able to execute 'az login' or other commands

# cannot use shell=True all the time, as it breaks macOS

logging.debug(f"Running command: {' '.join(popenargs)}")

return subprocess.run(

popenargs,

input=input,

capture_output=capture_output,

timeout=timeout,

check=check,

**kwargs,

)

class AzureCliTokenSource(CliTokenSource):

"""Obtain the token granted by `az login` CLI command"""

def __init__(

self,

resource: str,

subscription: Optional[str] = None,

tenant: Optional[str] = None,

):

cmd = [

"az",

"account",

"get-access-token",

"--resource",

resource,

"--output",

"json",

]

if subscription is not None:

cmd.append("--subscription")

cmd.append(subscription)

if tenant and not self.__is_cli_using_managed_identity():

cmd.extend(["--tenant", tenant])

super().__init__(

cmd=cmd,

token_type_field="tokenType",

access_token_field="accessToken",

expiry_field="expiresOn",

)

@staticmethod

def __is_cli_using_managed_identity() -> bool:

"""Checks whether the current CLI session is authenticated using managed identity."""

try:

cmd = ["az", "account", "show", "--output", "json"]

out = _run_subprocess(cmd, capture_output=True, check=True)

account = json.loads(out.stdout.decode())

user = account.get("user")

if user is None:

return False

return user.get("type") == "servicePrincipal" and user.get("name") in [

"systemAssignedIdentity",

"userAssignedIdentity",

]

except subprocess.CalledProcessError as e:

logger.debug("Failed to get account information from Azure CLI", exc_info=e)

return False

def is_human_user(self) -> bool:

"""The UPN claim is the username of the user, but not the Service Principal.

Azure CLI can be authenticated by both human users (`az login`) and service principals. In case of service

principals, it can be either OIDC from GitHub or login with a password:

~ $ az login --service-principal --user $clientID --password $clientSecret --tenant $tenantID

Human users get more claims:

- 'amr' - how the subject of the token was authenticated

- 'name', 'family_name', 'given_name' - human-readable values that identifies the subject of the token

- 'scp' with `user_impersonation` value, that shows the set of scopes exposed by your application for which

the client application has requested (and received) consent

- 'unique_name' - a human-readable value that identifies the subject of the token. This value is not

guaranteed to be unique within a tenant and should be used only for display purposes.

- 'upn' - The username of the user.

"""

return "upn" in self.token().jwt_claims()

@staticmethod

def for_resource(cfg: "Config", resource: str) -> "AzureCliTokenSource":

subscription = AzureCliTokenSource.get_subscription(cfg)

if subscription is not None:

token_source = AzureCliTokenSource(resource, subscription=subscription, tenant=cfg.azure_tenant_id)

try:

# This will fail if the user has access to the workspace, but not to the subscription

# itself.

# In such case, we fall back to not using the subscription.

token_source.token()

return token_source

except OSError:

logger.warning("Failed to get token for subscription. Using resource only token.")

token_source = AzureCliTokenSource(resource, subscription=None, tenant=cfg.azure_tenant_id)

token_source.token()

return token_source

@staticmethod

def get_subscription(cfg: "Config") -> Optional[str]:

resource = cfg.azure_workspace_resource_id

if resource is None or resource == "":

return None

components = resource.split("/")

if len(components) < 3:

logger.warning("Invalid azure workspace resource ID")

return None

return components[2]

@credentials_strategy("azure-cli", ["effective_azure_login_app_id"])

def azure_cli(cfg: "Config") -> Optional[CredentialsProvider]:

"""Adds refreshed OAuth token granted by `az login` command to every request."""

cfg.load_azure_tenant_id()

token_source = None

mgmt_token_source = None

try:

token_source = AzureCliTokenSource.for_resource(cfg, cfg.effective_azure_login_app_id)

except FileNotFoundError:

doc = "https://docs.microsoft.com/en-us/cli/azure/?view=azure-cli-latest"

logger.debug(f"Most likely Azure CLI is not installed. See {doc} for details")

return None

except OSError as e:

logger.debug("skipping Azure CLI auth", exc_info=e)

logger.debug("This may happen if you are attempting to login to a dev or staging workspace")

return None

if not token_source.is_human_user():

try:

management_endpoint = cfg.arm_environment.service_management_endpoint

mgmt_token_source = AzureCliTokenSource.for_resource(cfg, management_endpoint)

except Exception as e:

logger.debug(

"Not including service management token in headers",

exc_info=e,

)

mgmt_token_source = None

_ensure_host_present(cfg, lambda resource: AzureCliTokenSource.for_resource(cfg, resource))

logger.info("Using Azure CLI authentication with AAD tokens")

def inner() -> Dict[str, str]:

token = token_source.token()

headers = {"Authorization": f"{token.token_type} {token.access_token}"}

azure.add_workspace_id_header(cfg, headers)

if mgmt_token_source:

azure.add_sp_management_token(mgmt_token_source, headers)

return headers

return inner

class DatabricksCliTokenSource(CliTokenSource):

"""Obtain the token granted by `databricks auth login` CLI command"""

def __init__(self, cfg: "Config"):

cli_path = cfg.databricks_cli_path

# If the path is not specified look for "databricks" / "databricks.exe" in PATH.

if not cli_path:

try:

# Try to find "databricks" in PATH

cli_path = self.__class__._find_executable("databricks")

except FileNotFoundError as e:

# If "databricks" is not found, try to find "databricks.exe" in PATH (Windows)

if platform.system() == "Windows":

cli_path = self.__class__._find_executable("databricks.exe")

else:

raise e

# If the path is unqualified, look it up in PATH.

elif cli_path.count("/") == 0:

cli_path = self.__class__._find_executable(cli_path)

fallback_cmd = None

if cfg.profile:

# When profile is set, use --profile as the primary command.

# The profile contains the full config (host, account_id, etc.).

args = ["auth", "token", "--profile", cfg.profile]

# Build a --host fallback for older CLIs that don't support --profile.

if cfg.host:

fallback_cmd = [cli_path, *self.__class__._build_host_args(cfg)]

else:

args = self.__class__._build_host_args(cfg)

# get_scopes() defaults to ["all-apis"] when nothing is configured, which would

# cause false-positive mismatches against every token that wasn't issued with

# exactly ["all-apis"]. Only validate when scopes are explicitly set (either

# directly in code or loaded from a CLI profile).

self._requested_scopes = cfg.get_scopes() if cfg.scopes else None

self._host = cfg.host

super().__init__(

cmd=[cli_path, *args],

token_type_field="token_type",

access_token_field="access_token",

expiry_field="expiry",

disable_async=cfg.disable_async_token_refresh,

fallback_cmd=fallback_cmd,

)

def refresh(self) -> oauth.Token:

# The scope validation lives in refresh() because this is the only method that

# produces new tokens (see Refreshable._token assignments). By overriding here,

# every token is validated — both at initial auth and on every subsequent refresh

# when the cached token expires. This catches cases where a user re-authenticates

# mid-session with different scopes.

token = super().refresh()

if self._requested_scopes:

self._validate_token_scopes(token)

return token

# offline_access controls whether the IdP issues a refresh token. It does not

# grant any API permissions, so its presence or absence should not cause a

# scope mismatch error.

_SCOPES_IGNORED_FOR_COMPARISON = {"offline_access"}

def _validate_token_scopes(self, token: oauth.Token):

"""Validate that the token's scopes match the requested scopes from the config.

The `databricks auth token` command does not accept scopes yet. It returns whatever

token was cached from the last `databricks auth login`. If a user configures

specific scopes in the SDK config but their cached CLI token was issued with

different scopes, requests will silently use the wrong scopes. This check

surfaces that mismatch early with an actionable error telling the user how to

re-authenticate with the correct scopes.

"""

claims = token.jwt_claims()

if not claims:

logger.debug("Could not decode token as JWT to validate scopes")

return

token_scopes_raw = claims.get("scope")

if token_scopes_raw is None:

logger.debug("Token does not contain 'scope' claim, skipping scope validation")

return

if isinstance(token_scopes_raw, str):

token_scopes = set(token_scopes_raw.split())

elif isinstance(token_scopes_raw, list):

token_scopes = {str(s) for s in token_scopes_raw}

else:

logger.debug(f"Unexpected 'scope' claim type: {type(token_scopes_raw)}")

return

token_scopes -= self._SCOPES_IGNORED_FOR_COMPARISON

requested_scopes = set(self._requested_scopes) - self._SCOPES_IGNORED_FOR_COMPARISON

if token_scopes != requested_scopes:

raise ValueError(

f"Token issued by Databricks CLI has scopes {sorted(token_scopes)} which do not match "

f"the configured scopes {sorted(requested_scopes)}. Please re-authenticate "

f"with the desired scopes by running `databricks auth login` with the --scopes flag."

f"Scopes default to all-apis."

)

@staticmethod

def _build_host_args(cfg: "Config") -> List[str]:

"""Build CLI arguments using --host (legacy path)."""

args = ["auth", "token", "--host", cfg.host]

# This is here to support older versions of the Databricks CLI, so we need to keep the client type check.