1240 lines
46 KiB
Python
1240 lines
46 KiB
Python
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_repl_auth)
|
|
|
|
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,
|
|
)
|
|
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 and cfg.client_type == ClientType.ACCOUNT:
|
|
audience = cfg.account_id
|
|
if audience is None and cfg.client_type != ClientType.ACCOUNT:
|
|
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}"}
|
|
if cfg.client_type == ClientType.ACCOUNT:
|
|
gcp_credentials.refresh(request)
|
|
headers["X-Databricks-GCP-SA-Access-Token"] = gcp_credentials.token
|
|
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}"}
|
|
if cfg.client_type == ClientType.ACCOUNT:
|
|
gcp_impersonated_credentials.refresh(request)
|
|
headers["X-Databricks-GCP-SA-Access-Token"] = gcp_impersonated_credentials.token
|
|
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)
|
|
|
|
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,
|
|
)
|
|
|
|
@staticmethod
|
|
def _build_host_args(cfg: "Config") -> List[str]:
|
|
"""Build CLI arguments using --host (legacy path)."""
|
|
args = ["auth", "token", "--host", cfg.host]
|
|
if cfg.experimental_is_unified_host:
|
|
# For unified hosts, pass account_id, workspace_id, and experimental flag
|
|
args += ["--experimental-is-unified-host"]
|
|
if cfg.account_id:
|
|
args += ["--account-id", cfg.account_id]
|
|
if cfg.workspace_id:
|
|
args += ["--workspace-id", str(cfg.workspace_id)]
|
|
elif cfg.client_type == ClientType.ACCOUNT:
|
|
args += ["--account-id", cfg.account_id]
|
|
return args
|
|
|
|
@staticmethod
|
|
def _find_executable(name) -> str:
|
|
err = FileNotFoundError("Most likely the Databricks CLI is not installed")
|
|
for dir in os.getenv("PATH", default="").split(os.path.pathsep):
|
|
path = pathlib.Path(dir).joinpath(name).resolve()
|
|
if not path.is_file():
|
|
continue
|
|
|
|
# The new Databricks CLI is a single binary with size > 1MB.
|
|
# We use the size as a signal to determine which Databricks CLI is installed.
|
|
stat = path.stat()
|
|
if stat.st_size < (1024 * 1024):
|
|
err = FileNotFoundError("Databricks CLI version <0.100.0 detected")
|
|
continue
|
|
|
|
return str(path)
|
|
|
|
raise err
|
|
|
|
|
|
@oauth_credentials_strategy("databricks-cli", ["host"])
|
|
def databricks_cli(cfg: "Config") -> Optional[CredentialsProvider]:
|
|
try:
|
|
token_source = DatabricksCliTokenSource(cfg)
|
|
except FileNotFoundError as e:
|
|
logger.debug(e)
|
|
return None
|
|
|
|
try:
|
|
token_source.token()
|
|
except IOError as e:
|
|
if "databricks OAuth is not" in str(e):
|
|
logger.debug(f"OAuth not configured or not available: {e}")
|
|
return None
|
|
raise e
|
|
|
|
logger.info("Using Databricks CLI authentication")
|
|
|
|
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)
|
|
|
|
|
|
class MetadataServiceTokenSource(oauth.Refreshable):
|
|
"""Obtain the token granted by Databricks Metadata Service"""
|
|
|
|
METADATA_SERVICE_VERSION = "1"
|
|
METADATA_SERVICE_VERSION_HEADER = "X-Databricks-Metadata-Version"
|
|
METADATA_SERVICE_HOST_HEADER = "X-Databricks-Host"
|
|
_metadata_service_timeout = 10 # seconds
|
|
|
|
def __init__(self, cfg: "Config"):
|
|
super().__init__()
|
|
self.url = cfg.metadata_service_url
|
|
self.host = cfg.host
|
|
|
|
def refresh(self) -> oauth.Token:
|
|
resp = requests.get(
|
|
self.url,
|
|
timeout=self._metadata_service_timeout,
|
|
headers={
|
|
self.METADATA_SERVICE_VERSION_HEADER: self.METADATA_SERVICE_VERSION,
|
|
self.METADATA_SERVICE_HOST_HEADER: self.host,
|
|
},
|
|
proxies={
|
|
# Explicitly exclude localhost from being proxied. This is necessary
|
|
# for Metadata URLs which typically point to localhost.
|
|
"no_proxy": "localhost,127.0.0.1"
|
|
},
|
|
)
|
|
json_resp: dict[str, Union[str, float]] = resp.json()
|
|
access_token = json_resp.get("access_token", None)
|
|
if access_token is None:
|
|
raise ValueError("Metadata Service returned empty token")
|
|
token_type = json_resp.get("token_type", None)
|
|
if token_type is None:
|
|
raise ValueError("Metadata Service returned empty token type")
|
|
if json_resp["expires_on"] in ["", None]:
|
|
raise ValueError("Metadata Service returned invalid expiry")
|
|
try:
|
|
expiry = datetime.fromtimestamp(json_resp["expires_on"])
|
|
except:
|
|
raise ValueError("Metadata Service returned invalid expiry")
|
|
|
|
return oauth.Token(access_token=access_token, token_type=token_type, expiry=expiry)
|
|
|
|
|
|
@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)
|
|
token_source.token()
|
|
logger.info("Using Databricks Metadata Service authentication")
|
|
|
|
def inner() -> Dict[str, str]:
|
|
token = token_source.token()
|
|
return {"Authorization": f"{token.token_type} {token.access_token}"}
|
|
|
|
return inner
|
|
|
|
|
|
# This Code is derived from Mlflow DatabricksModelServingConfigProvider
|
|
# https://github.com/mlflow/mlflow/blob/1219e3ef1aac7d337a618a352cd859b336cf5c81/mlflow/legacy_databricks_cli/configure/provider.py#L332
|
|
class ModelServingAuthProvider:
|
|
USER_CREDENTIALS = "user_credentials"
|
|
|
|
_MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH = "/var/credentials-secret/model-dependencies-oauth-token"
|
|
|
|
def __init__(self, credential_type: Optional[str]):
|
|
self.expiry_time = -1
|
|
self.current_token = None
|
|
self.refresh_duration = 300 # 300 Seconds
|
|
self.credential_type = credential_type
|
|
|
|
def should_fetch_model_serving_environment_oauth() -> bool:
|
|
"""
|
|
Check whether this is the model serving environment
|
|
Additionally check if the oauth token file path exists
|
|
"""
|
|
|
|
is_in_model_serving_env = (
|
|
os.environ.get("IS_IN_DB_MODEL_SERVING_ENV")
|
|
or os.environ.get("IS_IN_DATABRICKS_MODEL_SERVING_ENV")
|
|
or "false"
|
|
)
|
|
return is_in_model_serving_env == "true" and os.path.isfile(
|
|
ModelServingAuthProvider._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH
|
|
)
|
|
|
|
def _get_model_dependency_oauth_token(self, should_retry=True) -> str:
|
|
# Use Cached value if it is valid
|
|
if self.current_token is not None and self.expiry_time > time.time():
|
|
return self.current_token
|
|
|
|
try:
|
|
with open(ModelServingAuthProvider._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH) as f:
|
|
oauth_dict = json.load(f)
|
|
self.current_token = oauth_dict["OAUTH_TOKEN"][0]["oauthTokenValue"]
|
|
self.expiry_time = time.time() + self.refresh_duration
|
|
except Exception as e:
|
|
# sleep and retry in case of any race conditions with OAuth refreshing
|
|
if should_retry:
|
|
logger.warning(
|
|
"Unable to read oauth token on first attmept in Model Serving Environment",
|
|
exc_info=e,
|
|
)
|
|
time.sleep(0.5)
|
|
return self._get_model_dependency_oauth_token(should_retry=False)
|
|
else:
|
|
raise RuntimeError(
|
|
"Unable to read OAuth credentials from the file mounted in Databricks Model Serving"
|
|
) from e
|
|
return self.current_token
|
|
|
|
def _get_invokers_token(self):
|
|
main_thread = threading.main_thread()
|
|
thread_data = main_thread.__dict__
|
|
invokers_token = None
|
|
if "invokers_token" in thread_data:
|
|
invokers_token = thread_data["invokers_token"]
|
|
|
|
if invokers_token is None:
|
|
raise RuntimeError("Unable to read Invokers Token in Databricks Model Serving")
|
|
|
|
return invokers_token
|
|
|
|
def get_databricks_host_token(self) -> Optional[Tuple[str, str]]:
|
|
if not ModelServingAuthProvider.should_fetch_model_serving_environment_oauth():
|
|
return None
|
|
|
|
# read from DB_MODEL_SERVING_HOST_ENV_VAR if available otherwise MODEL_SERVING_HOST_ENV_VAR
|
|
host = os.environ.get("DATABRICKS_MODEL_SERVING_HOST_URL") or os.environ.get("DB_MODEL_SERVING_HOST_URL")
|
|
|
|
if self.credential_type == ModelServingAuthProvider.USER_CREDENTIALS:
|
|
return (host, self._get_invokers_token())
|
|
else:
|
|
return (host, self._get_model_dependency_oauth_token())
|
|
|
|
|
|
def model_serving_auth_visitor(cfg: "Config", credential_type: Optional[str] = None) -> Optional[CredentialsProvider]:
|
|
try:
|
|
model_serving_auth_provider = ModelServingAuthProvider(credential_type)
|
|
host, token = model_serving_auth_provider.get_databricks_host_token()
|
|
if token is None:
|
|
raise ValueError(
|
|
"Got malformed auth (empty token) when fetching auth implicitly available in Model Serving Environment. Please contact Databricks support"
|
|
)
|
|
if cfg.host is None:
|
|
cfg.host = host
|
|
except Exception as e:
|
|
logger.warning(
|
|
"Unable to get auth from Databricks Model Serving Environment",
|
|
exc_info=e,
|
|
)
|
|
return None
|
|
logger.info("Using Databricks Model Serving Authentication")
|
|
|
|
def inner() -> Dict[str, str]:
|
|
# Call here again to get the refreshed token
|
|
_, token = model_serving_auth_provider.get_databricks_host_token()
|
|
return {"Authorization": f"Bearer {token}"}
|
|
|
|
return inner
|
|
|
|
|
|
@credentials_strategy("model-serving", [])
|
|
def model_serving_auth(cfg: "Config") -> Optional[CredentialsProvider]:
|
|
if not ModelServingAuthProvider.should_fetch_model_serving_environment_oauth():
|
|
logger.debug("model-serving: Not in Databricks Model Serving, skipping")
|
|
return None
|
|
|
|
return model_serving_auth_visitor(cfg)
|
|
|
|
|
|
class DefaultCredentials:
|
|
"""Select the first applicable credential provider from the chain"""
|
|
|
|
def __init__(self) -> None:
|
|
self._auth_type = "default"
|
|
self._auth_providers = [
|
|
pat_auth,
|
|
basic_auth,
|
|
metadata_service,
|
|
oauth_service_principal,
|
|
env_oidc,
|
|
file_oidc,
|
|
github_oidc,
|
|
azure_service_principal,
|
|
github_oidc_azure,
|
|
azure_cli,
|
|
azure_devops_oidc,
|
|
external_browser,
|
|
databricks_cli,
|
|
runtime_oauth,
|
|
runtime_native_auth,
|
|
google_credentials,
|
|
google_id,
|
|
model_serving_auth,
|
|
]
|
|
|
|
def auth_type(self) -> str:
|
|
return self._auth_type
|
|
|
|
def oauth_token(self, cfg: "Config") -> oauth.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
|
|
logger.debug(f"Ignoring {auth_type} auth, because {cfg.auth_type} is preferred")
|
|
continue
|
|
logger.debug(f"Attempting to configure auth: {auth_type}")
|
|
try:
|
|
# The header factory might be None if the provider cannot be
|
|
# configured for the current environment. For example, if the
|
|
# provider requires some missing environment variables.
|
|
header_factory = provider(cfg)
|
|
if not header_factory:
|
|
continue
|
|
|
|
self._auth_type = auth_type
|
|
return header_factory
|
|
except Exception as e:
|
|
raise ValueError(f"{auth_type}: {e}") from e
|
|
auth_flow_url = "https://docs.databricks.com/en/dev-tools/auth.html#databricks-client-unified-authentication"
|
|
raise ValueError(
|
|
f"cannot configure default credentials, please check {auth_flow_url} to configure credentials for your preferred authentication method."
|
|
)
|
|
|
|
|
|
class ModelServingUserCredentials(CredentialsStrategy):
|
|
"""
|
|
This credential strategy is designed for authenticating the Databricks SDK in the model serving environment using user-specific rights.
|
|
In the model serving environment, the strategy retrieves a downscoped user token from the thread-local variable.
|
|
In any other environments, the class defaults to the DefaultCredentialStrategy.
|
|
To use this credential strategy, instantiate the WorkspaceClient with the ModelServingUserCredentials strategy as follows:
|
|
|
|
invokers_client = WorkspaceClient(credential_strategy = ModelServingUserCredentials())
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.credential_type = ModelServingAuthProvider.USER_CREDENTIALS
|
|
self.default_credentials = DefaultCredentials()
|
|
|
|
def auth_type(self):
|
|
if ModelServingAuthProvider.should_fetch_model_serving_environment_oauth():
|
|
return "model_serving_" + self.credential_type
|
|
else:
|
|
return self.default_credentials.auth_type()
|
|
|
|
def __call__(self, cfg: "Config") -> CredentialsProvider:
|
|
if ModelServingAuthProvider.should_fetch_model_serving_environment_oauth():
|
|
header_factory = model_serving_auth_visitor(cfg, self.credential_type)
|
|
if not header_factory:
|
|
raise ValueError(
|
|
f"Unable to authenticate using {self.credential_type} in Databricks Model Serving Environment"
|
|
)
|
|
return header_factory
|
|
else:
|
|
return self.default_credentials(cfg)
|