From 3e3bafa2fd8c3796e7eac2044c6e3cdc3e042509 Mon Sep 17 00:00:00 2001 From: Mihai Mitrea Date: Wed, 26 Aug 2026 15:04:59 +0000 Subject: [PATCH] Add group role assumption support Signed-off-by: Mihai Mitrea --- README.md | 1 + databricks/sdk/__init__.py | 4 + databricks/sdk/config.py | 1 + databricks/sdk/credentials_provider.py | 91 ++++-- databricks/sdk/oauth.py | 11 + databricks/sdk/oidc.py | 16 +- docs/auth-types-reference.md | 14 +- docs/authentication.md | 21 ++ tests/integration/test_auth.py | 152 +++++++++- tests/test_client.py | 16 +- tests/test_config.py | 20 ++ tests/test_core.py | 15 + tests/test_credentials_provider.py | 370 ++++++++++++++++++++++++- tests/test_notebook_oauth.py | 41 +++ tests/test_oauth.py | 113 ++++++++ tests/test_oidc.py | 83 ++++++ tests/testdata/.databrickscfg | 4 + 17 files changed, 923 insertions(+), 50 deletions(-) diff --git a/README.md b/README.md index 2a9e63a58..ed5a421b7 100644 --- a/README.md +++ b/README.md @@ -186,6 +186,7 @@ During initialization, the SDK automatically resolves missing configuration fiel | `host` | _(String)_ The Databricks host URL for either the Databricks workspace endpoint or the Databricks accounts endpoint. | `DATABRICKS_HOST` | | `account_id` | _(String)_ The Databricks account ID for the Databricks accounts endpoint. Auto-discovered if not provided. | `DATABRICKS_ACCOUNT_ID` | | `workspace_id` | _(String)_ The Databricks workspace ID for the Databricks workspace endpoint. Auto-discovered if not provided. | `DATABRICKS_WORKSPACE_ID` | +| `group_id` | _(String)_ The ID of the group role to assume when using a supported Databricks OAuth authentication method. | `DATABRICKS_GROUP_ID` | | `cloud` | _(String)_ The cloud provider for the Databricks workspace (`AWS`, `AZURE`, or `GCP`). Auto-discovered if not provided. When set, `is_aws`, `is_azure`, and `is_gcp` use this value directly instead of inferring from hostname. | `DATABRICKS_CLOUD` | | `discovery_url` | _(String)_ The OpenID Connect discovery URL. Auto-discovered if not provided. When set, OIDC endpoints are fetched directly from this URL instead of using the default host-based well-known endpoint logic. | `DATABRICKS_DISCOVERY_URL` | | `token` | _(String)_ The Databricks personal access token (PAT) _(AWS, Azure, and GCP)_ or Azure Active Directory (Azure AD) token _(Azure)_. | `DATABRICKS_TOKEN` | diff --git a/databricks/sdk/__init__.py b/databricks/sdk/__init__.py index c97433b4f..ef270453e 100644 --- a/databricks/sdk/__init__.py +++ b/databricks/sdk/__init__.py @@ -294,6 +294,7 @@ def __init__( token: Optional[str] = None, profile: Optional[str] = None, config_file: Optional[str] = None, + group_id: Optional[str] = None, azure_workspace_resource_id: Optional[str] = None, azure_client_secret: Optional[str] = None, azure_client_id: Optional[str] = None, @@ -327,6 +328,7 @@ def __init__( token=token, profile=profile, config_file=config_file, + group_id=group_id, azure_workspace_resource_id=azure_workspace_resource_id, azure_client_secret=azure_client_secret, azure_client_id=azure_client_id, @@ -1217,6 +1219,7 @@ def __init__( token: Optional[str] = None, profile: Optional[str] = None, config_file: Optional[str] = None, + group_id: Optional[str] = None, azure_workspace_resource_id: Optional[str] = None, azure_client_secret: Optional[str] = None, azure_client_id: Optional[str] = None, @@ -1247,6 +1250,7 @@ def __init__( token=token, profile=profile, config_file=config_file, + group_id=group_id, azure_workspace_resource_id=azure_workspace_resource_id, azure_client_secret=azure_client_secret, azure_client_id=azure_client_id, diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index 219a115c1..0c2adcb91 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -95,6 +95,7 @@ class Config: # can route to the right workspace. Accepts a classic numeric workspace ID # or another workspace identifier format that the server understands. workspace_id: str = ConfigAttribute(env="DATABRICKS_WORKSPACE_ID") + group_id: str = ConfigAttribute(env="DATABRICKS_GROUP_ID") # Cloud provider. When set, is_aws/is_azure/is_gcp use this value directly # instead of inferring from hostname. Populated automatically from /.well-known/databricks-config. diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index dbe8bde1e..7382d77bc 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -85,7 +85,13 @@ def oauth_token(self, cfg: "Config") -> oauth.Token: return self._headers_provider(cfg).oauth_token() -def credentials_strategy(name: str, require: List[str]): +def _unsupported_group_role_assumption(auth_type: str) -> ValueError: + return ValueError( + f'auth type "{auth_type}" does not support group role assumption. Use Databricks OAuth authentication' + ) + + +def credentials_strategy(name: str, require: List[str], supports_group: Optional[bool] = None): """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.""" @@ -95,6 +101,12 @@ def inner( ) -> CredentialsStrategy: @functools.wraps(func) def wrapper(cfg: "Config") -> Optional[CredentialsProvider]: + if cfg.group_id and supports_group is False: + # Explicit auth has no fallback, so report why the requested strategy cannot be used. + # Default auth is a discovery chain, so decline and let it find a group-capable strategy. + if cfg.auth_type == name: + raise _unsupported_group_role_assumption(name) + return None for attr in require: if not getattr(cfg, attr): return None @@ -106,7 +118,7 @@ def wrapper(cfg: "Config") -> Optional[CredentialsProvider]: return inner -def oauth_credentials_strategy(name: str, require: List[str]): +def oauth_credentials_strategy(name: str, require: List[str], supports_group: Optional[bool] = None): """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. @@ -121,6 +133,12 @@ def inner( ) -> OauthCredentialsStrategy: @functools.wraps(func) def wrapper(cfg: "Config") -> Optional[OAuthCredentialsProvider]: + if cfg.group_id and supports_group is False: + # Explicit auth has no fallback, so report why the requested strategy cannot be used. + # Default auth is a discovery chain, so decline and let it find a group-capable strategy. + if cfg.auth_type == name: + raise _unsupported_group_role_assumption(name) + return None for attr in require: if not getattr(cfg, attr): return None @@ -131,7 +149,7 @@ def wrapper(cfg: "Config") -> Optional[OAuthCredentialsProvider]: return inner -@credentials_strategy("basic", ["host", "username", "password"]) +@credentials_strategy("basic", ["host", "username", "password"], supports_group=False) 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() @@ -143,7 +161,7 @@ def inner() -> Dict[str, str]: return inner -@credentials_strategy("pat", ["host", "token"]) +@credentials_strategy("pat", ["host", "token"], supports_group=False) def pat_auth(cfg: "Config") -> CredentialsProvider: """Adds Databricks Personal Access Token to every request""" static_credentials = {"Authorization": f"Bearer {cfg.token}"} @@ -154,8 +172,7 @@ def inner() -> Dict[str, str]: return inner -@credentials_strategy("runtime", []) -def runtime_native_auth(cfg: "Config") -> Optional[CredentialsProvider]: +def _runtime_native_auth(cfg: "Config") -> Optional[CredentialsProvider]: if "DATABRICKS_RUNTIME_VERSION" not in os.environ: return None @@ -198,13 +215,18 @@ def runtime_native_auth(cfg: "Config") -> Optional[CredentialsProvider]: return None -@oauth_credentials_strategy("runtime-oauth", ["scopes"]) +@credentials_strategy("runtime", [], supports_group=False) +def runtime_native_auth(cfg: "Config") -> Optional[CredentialsProvider]: + return _runtime_native_auth(cfg) + + +@oauth_credentials_strategy("runtime-oauth", ["scopes"], supports_group=True) 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) + native_auth = _runtime_native_auth(cfg) if native_auth is None: return None notebook_pat_token = None @@ -222,6 +244,7 @@ def get_notebook_pat_token() -> Optional[str]: host=cfg.host, scopes=cfg.get_scopes_as_string(), authorization_details=cfg.authorization_details, + group_id=cfg.group_id, ) def inner() -> Dict[str, str]: @@ -234,7 +257,7 @@ def token() -> oauth.Token: return OAuthCredentialsProvider(inner, token) -@oauth_credentials_strategy("oauth-m2m", ["host", "client_id", "client_secret"]) +@oauth_credentials_strategy("oauth-m2m", ["host", "client_id", "client_secret"], supports_group=True) 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. @@ -243,6 +266,7 @@ def oauth_service_principal(cfg: "Config") -> Optional[CredentialsProvider]: if oidc is None: return None + endpoint_params = {"assume_group": cfg.group_id} if cfg.group_id else None token_source = oauth.ClientCredentials( client_id=cfg.client_id, client_secret=cfg.client_secret, @@ -251,6 +275,7 @@ def oauth_service_principal(cfg: "Config") -> Optional[CredentialsProvider]: use_header=True, disable_async=cfg.disable_async_token_refresh, authorization_details=cfg.authorization_details, + endpoint_params=endpoint_params, ) def inner() -> Dict[str, str]: @@ -263,7 +288,7 @@ def token() -> oauth.Token: return OAuthCredentialsProvider(inner, token) -@credentials_strategy("external-browser", ["host", "auth_type"]) +@credentials_strategy("external-browser", ["host", "auth_type"], supports_group=True) def external_browser(cfg: "Config") -> Optional[CredentialsProvider]: if cfg.auth_type != "external-browser": return None @@ -275,6 +300,10 @@ def external_browser(cfg: "Config") -> Optional[CredentialsProvider]: client_secret = cfg.client_secret oidc_endpoints = cfg.databricks_oidc_endpoints elif cfg.azure_client_id: + # This branch authenticates against Azure Entra ID rather than Databricks OAuth, + # so it cannot mint a group-scoped token. + if cfg.group_id: + raise _unsupported_group_role_assumption("external-browser with Azure Entra ID") client_id = cfg.azure_client_id client_secret = cfg.azure_client_secret oidc_endpoints = get_azure_entra_id_workspace_endpoints(cfg.host) @@ -301,6 +330,7 @@ def external_browser(cfg: "Config") -> Optional[CredentialsProvider]: redirect_url=redirect_url, scopes=scopes, profile=cfg.profile, + group_id=cfg.group_id, ) credentials = token_cache.load() if credentials: @@ -320,6 +350,7 @@ def external_browser(cfg: "Config") -> Optional[CredentialsProvider]: redirect_url=redirect_url, client_secret=client_secret, scopes=scopes, + group_id=cfg.group_id, ) consent = oauth_client.initiate_consent() if not consent: @@ -350,6 +381,7 @@ def _ensure_host_present(cfg: "Config", token_source_for: Callable[[str], oauth. @oauth_credentials_strategy( "azure-client-secret", ["azure_client_id", "azure_client_secret"], + supports_group=False, ) def azure_service_principal(cfg: "Config") -> CredentialsProvider: """Adds refreshed Azure Active Directory (AAD) Service Principal OAuth tokens @@ -389,7 +421,7 @@ def token() -> oauth.Token: return OAuthCredentialsProvider(refreshed_headers, token) -@credentials_strategy("env-oidc", ["host"]) +@credentials_strategy("env-oidc", ["host"], supports_group=True) 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 @@ -401,7 +433,7 @@ def env_oidc(cfg) -> Optional[CredentialsProvider]: return oidc_credentials_provider(cfg, oidc.EnvIdTokenSource(env_var)) -@credentials_strategy("file-oidc", ["host", "oidc_token_filepath"]) +@credentials_strategy("file-oidc", ["host", "oidc_token_filepath"], supports_group=True) def file_oidc(cfg) -> Optional[CredentialsProvider]: return oidc_credentials_provider(cfg, oidc.FileIdTokenSource(cfg.oidc_token_filepath)) @@ -424,6 +456,7 @@ def oidc_credentials_provider(cfg, id_token_source: oidc.IdTokenSource) -> Optio id_token_source=id_token_source, disable_async=cfg.disable_async_token_refresh, scopes=cfg.get_scopes_as_string(), + group_id=cfg.group_id, ) def refreshed_headers() -> Dict[str, str]: @@ -476,15 +509,19 @@ def token_source_for(audience: str) -> oauth.TokenSource: # Should not happen, since we checked it above. raise Exception(f"Cannot get {provider_name} token") + 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", + } + if cfg.group_id: + endpoint_params["assume_group"] = cfg.group_id + 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", - }, + endpoint_params=endpoint_params, scopes=cfg.get_scopes_as_string(), use_params=True, disable_async=cfg.disable_async_token_refresh, @@ -501,7 +538,7 @@ def token() -> oauth.Token: return OAuthCredentialsProvider(refreshed_headers, token) -@oauth_credentials_strategy("github-oidc", ["host", "client_id"]) +@oauth_credentials_strategy("github-oidc", ["host", "client_id"], supports_group=True) def github_oidc(cfg: "Config") -> Optional[CredentialsProvider]: """ GitHub OIDC authentication uses a Token Supplier to get a JWT Token and exchanges @@ -516,7 +553,7 @@ def github_oidc(cfg: "Config") -> Optional[CredentialsProvider]: ) -@oauth_credentials_strategy("azure-devops-oidc", ["host", "client_id"]) +@oauth_credentials_strategy("azure-devops-oidc", ["host", "client_id"], supports_group=True) def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]: """ Azure DevOps OIDC authentication uses a Token Supplier to get a JWT Token @@ -533,7 +570,7 @@ def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]: # 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"]) +@oauth_credentials_strategy("github-oidc-azure", ["host", "azure_client_id"], supports_group=False) def github_oidc_azure(cfg: "Config") -> Optional[CredentialsProvider]: if "ACTIONS_ID_TOKEN_REQUEST_TOKEN" not in os.environ: # not in GitHub actions @@ -585,7 +622,7 @@ def token() -> oauth.Token: ] -@oauth_credentials_strategy("google-credentials", ["host", "google_credentials"]) +@oauth_credentials_strategy("google-credentials", ["host", "google_credentials"], supports_group=False) 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. @@ -624,7 +661,7 @@ def refreshed_headers() -> Dict[str, str]: return OAuthCredentialsProvider(refreshed_headers, token) -@oauth_credentials_strategy("google-id", ["host", "google_service_account"]) +@oauth_credentials_strategy("google-id", ["host", "google_service_account"], supports_group=False) def google_id(cfg: "Config") -> Optional[CredentialsProvider]: credentials, _project_id = google.auth.default() @@ -877,7 +914,7 @@ def get_subscription(cfg: "Config") -> Optional[str]: return components[2] -@credentials_strategy("azure-cli", ["effective_azure_login_app_id"]) +@credentials_strategy("azure-cli", ["effective_azure_login_app_id"], supports_group=False) def azure_cli(cfg: "Config") -> Optional[CredentialsProvider]: """Adds refreshed OAuth token granted by `az login` command to every request.""" cfg.load_azure_tenant_id() @@ -1209,7 +1246,7 @@ def _find_executable(name) -> str: raise err -@oauth_credentials_strategy("databricks-cli", ["host"]) +@oauth_credentials_strategy("databricks-cli", ["host"], supports_group=False) def databricks_cli(cfg: "Config") -> Optional[CredentialsProvider]: try: token_source = DatabricksCliTokenSource(cfg) @@ -1291,7 +1328,7 @@ def refresh(self) -> oauth.Token: return oauth.Token(access_token=access_token, token_type=token_type, expiry=expiry) -@credentials_strategy("metadata-service", ["host", "metadata_service_url"]) +@credentials_strategy("metadata-service", ["host", "metadata_service_url"], supports_group=False) def metadata_service(cfg: "Config") -> Optional[CredentialsProvider]: """Adds refreshed token granted by Databricks Metadata Service to every request.""" @@ -1410,7 +1447,7 @@ def inner() -> Dict[str, str]: return inner -@credentials_strategy("model-serving", []) +@credentials_strategy("model-serving", [], supports_group=False) 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") @@ -1504,6 +1541,8 @@ def auth_type(self): def __call__(self, cfg: "Config") -> CredentialsProvider: if ModelServingAuthProvider.should_fetch_model_serving_environment_oauth(): + if cfg.group_id: + raise _unsupported_group_role_assumption("model-serving") header_factory = model_serving_auth_visitor(cfg, self.credential_type) if not header_factory: raise ValueError( diff --git a/databricks/sdk/oauth.py b/databricks/sdk/oauth.py index e5d034ae2..af7eff789 100644 --- a/databricks/sdk/oauth.py +++ b/databricks/sdk/oauth.py @@ -762,6 +762,7 @@ def __init__( client_id: str, scopes: List[str] = None, client_secret: str = None, + group_id: str = None, ): if not scopes: # Default for direct OAuthClient users (e.g., via from_host()). @@ -775,6 +776,7 @@ def __init__( self._client_secret = client_secret self._oidc_endpoints = oidc_endpoints self._scopes = scopes + self._group_id = group_id @staticmethod def from_host( @@ -815,6 +817,8 @@ def initiate_consent(self) -> Consent: "code_challenge": challenge, "code_challenge_method": "S256", } + if self._group_id: + params["assume_group"] = self._group_id auth_url = f"{self._oidc_endpoints.authorization_endpoint}?{urllib.parse.urlencode(params)}" return Consent( state, @@ -896,6 +900,7 @@ class PATOAuthTokenExchange(Refreshable): host: str scopes: str authorization_details: str = None + group_id: str = None disable_async: bool = True def __post_init__(self): @@ -912,6 +917,8 @@ def refresh(self) -> Token: } if self.authorization_details: params["authorization_details"] = self.authorization_details + if self.group_id: + params["assume_group"] = self.group_id resp = requests.post(token_exchange_url, params) if not resp.ok: @@ -947,6 +954,7 @@ def __init__( client_secret: Optional[str] = None, scopes: Optional[List[str]] = None, profile: Optional[str] = None, + group_id: Optional[str] = None, ) -> None: self._host = host self._client_id = client_id @@ -955,6 +963,7 @@ def __init__( self._client_secret = client_secret self._scopes = scopes or [] self._profile = profile + self._group_id = group_id @property def filename(self) -> str: @@ -966,6 +975,8 @@ def filename(self) -> str: "scopes": self._scopes, "profile": self._profile or "", } + if self._group_id: + key["group_id"] = self._group_id h = hashlib.sha256(json.dumps(key, sort_keys=True).encode("utf-8")) return os.path.expanduser(os.path.join(self.__class__.BASE_PATH, h.hexdigest() + ".json")) diff --git a/databricks/sdk/oidc.py b/databricks/sdk/oidc.py index 6e37cbabc..1a3c8d3d2 100644 --- a/databricks/sdk/oidc.py +++ b/databricks/sdk/oidc.py @@ -163,6 +163,7 @@ def __init__( audience: Optional[str] = None, disable_async: bool = False, scopes: Optional[str] = None, + group_id: Optional[str] = None, ): self._host = host self._id_token_source = id_token_source @@ -171,6 +172,7 @@ def __init__( self._account_id = account_id self._audience = audience self._scopes = scopes + self._group_id = group_id # Refreshable.__init__ stores disable_async as self._disable_async, which # _exchange_id_token reads — no need to duplicate it here. super().__init__(disable_async=disable_async) @@ -206,15 +208,19 @@ def refresh(self) -> oauth.Token: # This function is used to create the OAuth client. # It exists to make it easier to test. def _exchange_id_token(self, id_token: IdToken) -> oauth.Token: + endpoint_params = { + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + "subject_token": id_token.jwt, + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + } + if self._group_id: + endpoint_params["assume_group"] = self._group_id + client = oauth.ClientCredentials( client_id=self._client_id, client_secret="", # there is no (rotatable) secrets in the OIDC flow token_url=self._token_endpoint, - endpoint_params={ - "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", - "subject_token": id_token.jwt, - "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", - }, + endpoint_params=endpoint_params, scopes=self._scopes, use_params=True, disable_async=self._disable_async, diff --git a/docs/auth-types-reference.md b/docs/auth-types-reference.md index 47837cc1c..b52054b2f 100644 --- a/docs/auth-types-reference.md +++ b/docs/auth-types-reference.md @@ -8,22 +8,22 @@ This document lists all authentication types (`auth_type`) supported by the Data |-----------|-------------|---------------------|---------------------|----------------------| | `pat` | Personal Access Token authentication - the most common method for programmatic access | `host`, `token` | - | `DATABRICKS_HOST`, `DATABRICKS_TOKEN` | | `basic` | Basic HTTP authentication using username and password (primarily for AWS) | `host`, `username`, `password` | `account_id` (for account-level operations) | `DATABRICKS_HOST`, `DATABRICKS_USERNAME`, `DATABRICKS_PASSWORD`, `DATABRICKS_ACCOUNT_ID` | -| `oauth-m2m` | OAuth 2.0 Machine-to-Machine (service principal) authentication | `host`, `client_id`, `client_secret` | `scopes`, `authorization_details` | `DATABRICKS_HOST`, `DATABRICKS_CLIENT_ID`, `DATABRICKS_CLIENT_SECRET` | -| `external-browser` | OAuth 2.0 authentication flow using local browser for user login | `host`, `auth_type='external-browser'` | `client_id`, `client_secret` | `DATABRICKS_HOST`, `DATABRICKS_AUTH_TYPE`, `DATABRICKS_CLIENT_ID` | +| `oauth-m2m` | OAuth 2.0 Machine-to-Machine (service principal) authentication | `host`, `client_id`, `client_secret` | `scopes`, `authorization_details`, `group_id` | `DATABRICKS_HOST`, `DATABRICKS_CLIENT_ID`, `DATABRICKS_CLIENT_SECRET`, `DATABRICKS_GROUP_ID` | +| `external-browser` | OAuth 2.0 authentication flow using local browser for user login | `host`, `auth_type='external-browser'` | `client_id`, `client_secret`, `group_id` | `DATABRICKS_HOST`, `DATABRICKS_AUTH_TYPE`, `DATABRICKS_CLIENT_ID`, `DATABRICKS_GROUP_ID` | | `databricks-cli` | Uses tokens from the Databricks CLI (`databricks auth login`) | `host` | `account_id` (for account-level), `databricks_cli_path` | `DATABRICKS_HOST`, `DATABRICKS_ACCOUNT_ID`, `DATABRICKS_CLI_PATH` | | `azure-client-secret` | Azure Active Directory (AAD) Service Principal authentication | `azure_client_id`, `azure_client_secret` | `azure_tenant_id` (auto-detected from `host` if not set), `host`, `azure_workspace_resource_id`, `azure_environment` | `ARM_CLIENT_ID`, `ARM_CLIENT_SECRET`, `ARM_TENANT_ID`, `DATABRICKS_HOST`, `DATABRICKS_AZURE_RESOURCE_ID`, `ARM_ENVIRONMENT` | | `azure-cli` | Uses credentials from Azure CLI (`az login`) | `host` (or `azure_workspace_resource_id`) | `azure_tenant_id` | `DATABRICKS_HOST`, `DATABRICKS_AZURE_RESOURCE_ID`, `ARM_TENANT_ID` | -| `github-oidc` | GitHub Actions OIDC authentication (workload identity federation) | `host`, `client_id` | `token_audience`, `account_id` | `DATABRICKS_HOST`, `DATABRICKS_CLIENT_ID`, `DATABRICKS_TOKEN_AUDIENCE`, `DATABRICKS_ACCOUNT_ID` | +| `github-oidc` | GitHub Actions OIDC authentication (workload identity federation) | `host`, `client_id` | `token_audience`, `account_id`, `group_id` | `DATABRICKS_HOST`, `DATABRICKS_CLIENT_ID`, `DATABRICKS_TOKEN_AUDIENCE`, `DATABRICKS_ACCOUNT_ID`, `DATABRICKS_GROUP_ID` | | `github-oidc-azure` | GitHub Actions OIDC for Azure Databricks workspaces | `host`, `azure_client_id` | `azure_tenant_id` | `DATABRICKS_HOST`, `ARM_CLIENT_ID`, `ARM_TENANT_ID` | -| `azure-devops-oidc` | Azure DevOps Pipelines OIDC authentication | `host`, `client_id` | `token_audience`, `account_id` | `DATABRICKS_HOST`, `DATABRICKS_CLIENT_ID`, `SYSTEM_ACCESSTOKEN` | +| `azure-devops-oidc` | Azure DevOps Pipelines OIDC authentication | `host`, `client_id` | `token_audience`, `account_id`, `group_id` | `DATABRICKS_HOST`, `DATABRICKS_CLIENT_ID`, `SYSTEM_ACCESSTOKEN`, `DATABRICKS_GROUP_ID` | | `google-credentials` | Google Cloud service account authentication using credentials JSON | `host`, `google_credentials` | - | `DATABRICKS_HOST`, `GOOGLE_CREDENTIALS` | | `google-id` | Google Cloud authentication using service account impersonation | `host`, `google_service_account` | - | `DATABRICKS_HOST`, `DATABRICKS_GOOGLE_SERVICE_ACCOUNT` | | `metadata-service` | Authentication using Databricks-hosted metadata service | `host`, `metadata_service_url` | - | `DATABRICKS_HOST`, `DATABRICKS_METADATA_SERVICE_URL` | | `runtime` | Auto-detected authentication when running in Databricks Runtime (notebooks, jobs) | _(auto-detected)_ | - | `DATABRICKS_RUNTIME_VERSION` (auto-set) | -| `runtime-oauth` | OAuth authentication for Databricks Runtime with fine-grained permissions | `scopes` | `authorization_details` | `DATABRICKS_RUNTIME_VERSION` (auto-set) | +| `runtime-oauth` | OAuth authentication for Databricks Runtime with fine-grained permissions | `scopes` | `authorization_details`, `group_id` | `DATABRICKS_RUNTIME_VERSION` (auto-set), `DATABRICKS_GROUP_ID` | | `model-serving` | Auto-detected authentication when running in Databricks Model Serving environment | _(auto-detected)_ | - | `IS_IN_DB_MODEL_SERVING_ENV` or `IS_IN_DATABRICKS_MODEL_SERVING_ENV` (auto-set) | -| `env-oidc` | OIDC token from environment variable | `host` | `oidc_token_env`, `client_id` | `DATABRICKS_HOST`, `DATABRICKS_OIDC_TOKEN`, `DATABRICKS_OIDC_TOKEN_ENV`, `DATABRICKS_CLIENT_ID` | -| `file-oidc` | OIDC token from file path | `host`, `oidc_token_filepath` | `client_id` | `DATABRICKS_HOST`, `DATABRICKS_OIDC_TOKEN_FILEPATH` (alias: `DATABRICKS_OIDC_TOKEN_FILE`), `DATABRICKS_CLIENT_ID` | +| `env-oidc` | OIDC token from environment variable | `host` | `oidc_token_env`, `client_id`, `group_id` | `DATABRICKS_HOST`, `DATABRICKS_OIDC_TOKEN`, `DATABRICKS_OIDC_TOKEN_ENV`, `DATABRICKS_CLIENT_ID`, `DATABRICKS_GROUP_ID` | +| `file-oidc` | OIDC token from file path | `host`, `oidc_token_filepath` | `client_id`, `group_id` | `DATABRICKS_HOST`, `DATABRICKS_OIDC_TOKEN_FILEPATH` (alias: `DATABRICKS_OIDC_TOKEN_FILE`), `DATABRICKS_CLIENT_ID`, `DATABRICKS_GROUP_ID` | For configuration options that apply to all authentication types (timeouts, debug settings, rate limits), see [Authentication](./authentication.md#additional-configuration-options). diff --git a/docs/authentication.md b/docs/authentication.md index 52d375ca4..719db9b79 100644 --- a/docs/authentication.md +++ b/docs/authentication.md @@ -131,6 +131,7 @@ For all authentication methods, you can override the default behavior in client | Argument | Description | Environment variable | |-------------------------|-------------|------------------------| | `auth_type` | _(String)_ When multiple auth attributes are available in the environment, use the auth type specified by this argument. This argument also holds the currently selected auth. When set explicitly, the SDK only attempts that authentication method, skipping automatic detection of others. See the [Authentication Types Reference](./auth-types-reference.md) for all valid values, required parameters, and usage examples. | `DATABRICKS_AUTH_TYPE` | +| `group_id` | _(String)_ ID of the group role to assume. Supported with `external-browser`, `oauth-m2m`, `github-oidc`, `azure-devops-oidc`, `env-oidc`, `file-oidc`, and `runtime-oauth`. | `DATABRICKS_GROUP_ID` | | `http_timeout_seconds` | _(Integer)_ Number of seconds for HTTP timeout. Default is _60_. | _(None)_ | | `retry_timeout_seconds` | _(Integer)_ Number of seconds to keep retrying HTTP requests. Default is _300 (5 minutes)_. | _(None)_ | | `debug_truncate_bytes` | _(Integer)_ Truncate JSON fields in debug logs above this limit. Default is 96. | `DATABRICKS_DEBUG_TRUNCATE_BYTES` | @@ -144,3 +145,23 @@ from databricks.sdk import WorkspaceClient w = WorkspaceClient(debug_headers=True) # Now call the Databricks workspace APIs as desired... ``` + +## Group role assumption + +Set `group_id`, `DATABRICKS_GROUP_ID`, or `group_id` in the selected configuration profile to use a group role. Explicit client configuration takes precedence over the environment, which takes precedence over the profile. The role replaces the caller's normal permissions for the lifetime of the client. + +```python +from databricks.sdk import WorkspaceClient + +w = WorkspaceClient( + host="https://your-workspace.cloud.databricks.com", + client_id="your-client-id", + client_secret="your-client-secret", + auth_type="oauth-m2m", + group_id="your-group-id", +) +``` + +Group role assumption requires a Databricks OAuth authentication method. PAT, basic, Databricks CLI, Azure or Google native authentication, metadata service, native runtime, and model serving credentials do not support it. Unified and account OAuth endpoints receive the group request and return an error when the server does not support it; the SDK never retries with normal permissions. + +An assumed group identity cannot currently call `/api/2.0/preview/scim/v2/Me`. This affects current-user lookup, workspace-ID fallback, and `Config.sql_http_path` when `cluster_id` is configured. Setting `workspace_id` avoids only the workspace-ID fallback; it does not make cluster SQL HTTP-path derivation work. The experimental Files storage-proxy probe also cannot resolve the workspace ID through `/Me` and falls back to presigned URLs. diff --git a/tests/integration/test_auth.py b/tests/integration/test_auth.py index 681c37f53..cc9fc01e7 100644 --- a/tests/integration/test_auth.py +++ b/tests/integration/test_auth.py @@ -7,6 +7,7 @@ import sys import typing import urllib.parse +from contextlib import ExitStack from functools import partial from pathlib import Path @@ -14,10 +15,11 @@ from databricks.sdk import AccountClient, WorkspaceClient from databricks.sdk.config import Config +from databricks.sdk.errors import NotFound, PermissionDenied from databricks.sdk.service import iam, oauth2 from databricks.sdk.service.compute import ClusterSpec, DataSecurityMode, Library, ResultType, SparkVersion from databricks.sdk.service.jobs import NotebookTask, Task, ViewType -from databricks.sdk.service.workspace import ImportFormat +from databricks.sdk.service.workspace import ImportFormat, Language @pytest.fixture @@ -270,6 +272,154 @@ def test_wif_workspace(ucacct, env_or_skip, random): ws.current_user.me() +def _ignore_not_found(action): + try: + action() + except NotFound: + pass + + +def test_wif_workspace_group_role_isolation(ucacct, env_or_skip, random): + """Verifies group-role WIF can read a resource that normal WIF credentials cannot.""" + # Use the GitHub Actions OIDC environment and an account administrator to arrange the test. + env_or_skip("ACTIONS_ID_TOKEN_REQUEST_URL") + workspace_id = int(env_or_skip("TEST_WORKSPACE_ID")) + workspace_url = env_or_skip("TEST_WORKSPACE_URL") + audience = "https://github.com/databricks-eng" + + # Use administrator credentials to create the workspace resource and set its permissions. + workspace_admin = WorkspaceClient(host=workspace_url) + + with ExitStack() as cleanup: + # Create the service principal whose normal and role-based WIF access will be compared. + service_principal = ucacct.service_principals_v2.create( + active=True, + display_name="py-sdk-wif-role-sp-" + random(), + ) + cleanup.callback( + _ignore_not_found, + lambda: ucacct.service_principals_v2.delete(service_principal.id), + ) + service_principal_id = int(service_principal.id) + + # Give the service principal basic workspace access without granting notebook access. + ucacct.workspace_assignment.update( + workspace_id, + service_principal_id, + permissions=[iam.WorkspacePermission.USER], + ) + cleanup.callback( + _ignore_not_found, + lambda: ucacct.workspace_assignment.delete(workspace_id, service_principal_id), + ) + + # Create the group that represents the temporary workspace role. + group = ucacct.groups_v2.create(display_name="py-sdk-wif-role-group-" + random()) + cleanup.callback( + _ignore_not_found, + lambda: ucacct.groups_v2.delete(group.id), + ) + group_id = int(group.id) + + # Assign the group to the workspace so that it can receive workspace permissions. + ucacct.workspace_assignment.update( + workspace_id, + group_id, + permissions=[iam.WorkspacePermission.USER], + ) + cleanup.callback( + _ignore_not_found, + lambda: ucacct.workspace_assignment.delete(workspace_id, group_id), + ) + + # Allow the service principal to assume the group role. + rule_set_name = f"accounts/{ucacct.config.account_id}/groups/{group.id}/ruleSets/default" + rule_set = ucacct.access_control.get_rule_set(rule_set_name, "") + ucacct.access_control.update_rule_set( + rule_set_name, + iam.RuleSetUpdateRequest( + name=rule_set_name, + etag=rule_set.etag, + grant_rules=[ + *(rule_set.grant_rules or []), + iam.GrantRule( + principals=[f"servicePrincipals/{service_principal.application_id}"], + role="roles/group.assumer", + ), + ], + ), + ) + + # Trust this repository's GitHub OIDC identity to authenticate as the service principal. + policy = ucacct.service_principal_federation_policy.create( + service_principal_id, + oauth2.FederationPolicy( + oidc_policy=oauth2.OidcFederationPolicy( + issuer="https://token.actions.githubusercontent.com", + audiences=[audience], + subject="repo:databricks-eng/eng-dev-ecosystem:environment:integration-tests", + ) + ), + ) + cleanup.callback( + _ignore_not_found, + lambda: ucacct.service_principal_federation_policy.delete( + service_principal_id, + policy.policy_id or policy.uid, + ), + ) + + # Create a private notebook that distinguishes normal access from role access. + workspace_admin_user = workspace_admin.current_user.me().user_name + notebook_directory = f"/Users/{workspace_admin_user}/.sdk/notebooks/py-sdk-wif-role-{random()}" + notebook_path = f"{notebook_directory}/notebook" + workspace_admin.workspace.mkdirs(notebook_directory) + cleanup.callback( + _ignore_not_found, + lambda: workspace_admin.workspace.delete(notebook_directory, recursive=True), + ) + workspace_admin.workspace.upload( + notebook_path, + b"print(1)", + format=ImportFormat.SOURCE, + language=Language.PYTHON, + overwrite=True, + ) + notebook = workspace_admin.workspace.get_status(notebook_path) + + # Grant only the group role permission to read the notebook. + workspace_admin.permissions.update( + "notebooks", + str(notebook.object_id), + access_control_list=[ + iam.AccessControlRequest( + group_name=group.display_name, + permission_level=iam.PermissionLevel.CAN_READ, + ) + ], + ) + + # Authenticate with the group role and verify that its notebook permission is usable. + role_client = WorkspaceClient( + host=workspace_url, + client_id=service_principal.application_id, + group_id=group.id, + auth_type="github-oidc", + token_audience=audience, + ) + role_client.workspace.get_status(notebook_path) + + # Authenticate normally as the same service principal and verify that access is denied. + normal_client = WorkspaceClient( + host=workspace_url, + client_id=service_principal.application_id, + auth_type="github-oidc", + token_audience=audience, + ) + with pytest.raises((PermissionDenied, NotFound)): + normal_client.workspace.get_status(notebook_path) + + def test_workspace_config_resolves_account_and_workspace_id(w, env_or_skip): """Test that Config resolves account_id and workspace_id from host metadata.""" env_or_skip("CLOUD_ENV") diff --git a/tests/test_client.py b/tests/test_client.py index c616662b4..19e6680cf 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -5,7 +5,9 @@ import pytest -from databricks.sdk import WorkspaceClient +from databricks.sdk import AccountClient, WorkspaceClient + +from .conftest import noop_credentials def test_autospec_fails_on_unknown_service(): @@ -37,6 +39,18 @@ def test_workspace_client_init_does_not_build_dbutils(config, mocker): spy.assert_not_called() +@pytest.mark.parametrize("client_type", [WorkspaceClient, AccountClient]) +def test_client_forwards_group_id_to_config(client_type): + """Verifies public clients preserve group_id in their underlying configuration.""" + client = client_type( + host="https://example.cloud.databricks.com", + group_id="test-group", + credentials_strategy=noop_credentials, + ) + + assert client.config.group_id == "test-group" + + def test_dbutils_first_access_builds_exactly_once(config, mocker): """First read of ``.dbutils`` calls ``_make_dbutils`` once; subsequent reads return the cached value without re-invoking.""" diff --git a/tests/test_config.py b/tests/test_config.py index 212228a25..348c620f6 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -603,6 +603,26 @@ def test_config_file_scopes_multiple_sorted(monkeypatch, mocker): assert config.get_scopes() == expected +def test_group_id_config_precedence(monkeypatch, mocker): + """Verifies constructor, environment, and profile group IDs use standard precedence.""" + mocker.patch("databricks.sdk.config.Config.init_auth") + set_home(monkeypatch, "/testdata") + monkeypatch.delenv("DATABRICKS_GROUP_ID", raising=False) + + config = Config(profile="scope-empty") + assert config.group_id is None + + config = Config(profile="group-role") + assert config.group_id == "profile-group" + + monkeypatch.setenv("DATABRICKS_GROUP_ID", "environment-group") + config = Config(profile="group-role") + assert config.group_id == "environment-group" + + config = Config(profile="group-role", group_id="configured-group") + assert config.group_id == "configured-group" + + def _get_scope_from_request(request_text: str) -> Optional[str]: """Extract the scope value from a URL-encoded request body.""" params = parse_qs(request_text) diff --git a/tests/test_core.py b/tests/test_core.py index 4c1b3a699..3c499f6d5 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -38,6 +38,21 @@ def test_parse_dsn(): assert "basic" == cfg.auth_type +def test_group_id_does_not_change_api_request_headers(requests_mock): + """Verifies role selection changes authentication without adding API request headers.""" + requests_mock.get("/test", json={}) + + normal_config = Config(host="http://localhost", credentials_strategy=noop_credentials) + ApiClient(normal_config).do("GET", "/test") + normal_headers = dict(requests_mock.last_request.headers) + + grouped_config = Config(host="http://localhost", group_id="group-id", credentials_strategy=noop_credentials) + ApiClient(grouped_config).do("GET", "/test") + grouped_headers = dict(requests_mock.last_request.headers) + + assert grouped_headers == normal_headers + + def test_databricks_cli_token_source_relative_path(config): config.databricks_cli_path = "./relative/path/to/cli" ts = DatabricksCliTokenSource(config) diff --git a/tests/test_credentials_provider.py b/tests/test_credentials_provider.py index 2536ee61b..1e8fd2df5 100644 --- a/tests/test_credentials_provider.py +++ b/tests/test_credentials_provider.py @@ -1,5 +1,6 @@ from datetime import datetime, timedelta from unittest.mock import Mock +from urllib.parse import parse_qs import pytest @@ -234,6 +235,7 @@ def test_external_browser_passes_profile_to_token_cache(mocker): mock_cfg.auth_type = "external-browser" mock_cfg.host = "https://test.databricks.com" mock_cfg.profile = "myprofile" + mock_cfg.group_id = "group-id" mock_cfg.client_id = "test-client-id" mock_cfg.client_secret = None mock_cfg.azure_client_id = None @@ -254,6 +256,354 @@ def test_external_browser_passes_profile_to_token_cache(mocker): credentials_provider.external_browser(mock_cfg) assert mock_token_cache_class.call_args.kwargs["profile"] == "myprofile" + assert mock_token_cache_class.call_args.kwargs["group_id"] == "group-id" + + +def test_external_browser_passes_group_to_oauth_client(mocker): + """Verifies external-browser auth forwards the requested group to its OAuth client.""" + mock_cfg = Mock() + mock_cfg.auth_type = "external-browser" + mock_cfg.host = "https://test.databricks.com" + mock_cfg.group_id = "group-id" + mock_cfg.profile = None + mock_cfg.client_id = "client-id" + mock_cfg.client_secret = None + mock_cfg.azure_client_id = None + mock_cfg.get_scopes.return_value = ["all-apis"] + mock_cfg.disable_oauth_refresh_token = True + + mocker.patch("databricks.sdk.credentials_provider.oauth.TokenCache").return_value.load.return_value = None + oauth_client_class = mocker.patch("databricks.sdk.credentials_provider.oauth.OAuthClient") + oauth_client_class.return_value.initiate_consent.return_value = None + + credentials_provider.external_browser(mock_cfg) + + assert oauth_client_class.call_args.kwargs["group_id"] == "group-id" + + +def test_external_browser_with_azure_entra_rejects_group(): + """Verifies Azure Entra browser auth rejects unsupported group-role requests.""" + cfg = Mock( + auth_type="external-browser", + group_id="group-id", + client_id=None, + azure_client_id="azure-client-id", + ) + + with pytest.raises(ValueError, match="external-browser with Azure Entra ID"): + credentials_provider.external_browser(cfg) + + +def test_external_browser_group_cache_isolation(mocker, monkeypatch, tmp_path): + """Verifies normal and grouped browser sessions persist in separate cache entries.""" + monkeypatch.setattr(oauth.TokenCache, "BASE_PATH", str(tmp_path)) + oidc_endpoints = oauth.OidcEndpoints( + "https://test.databricks.com/oidc/v1/authorize", + "https://test.databricks.com/oidc/v1/token", + ) + + def config(group_id): + cfg = Mock( + auth_type="external-browser", + host="https://test.databricks.com", + profile=None, + group_id=group_id, + client_id="client-id", + client_secret=None, + azure_client_id=None, + disable_oauth_refresh_token=False, + databricks_oidc_endpoints=oidc_endpoints, + ) + cfg.get_scopes.return_value = ["all-apis"] + return cfg + + def oauth_client(**kwargs): + group_id = kwargs["group_id"] or "normal" + credentials = oauth.SessionCredentials( + oauth.Token( + access_token=f"{group_id}-token", + token_type="Bearer", + expiry=datetime.now() + timedelta(hours=1), + ), + oidc_endpoints.token_endpoint, + "client-id", + ) + consent = Mock() + consent.launch_external_browser.return_value = credentials + client = Mock() + client.initiate_consent.return_value = consent + return client + + oauth_client_class = mocker.patch( + "databricks.sdk.credentials_provider.oauth.OAuthClient", + side_effect=oauth_client, + ) + + for group_id in [None, "group-a", "group-b"]: + provider = credentials_provider.external_browser(config(group_id)) + assert provider() == {"Authorization": f"Bearer {group_id or 'normal'}-token"} + + assert oauth_client_class.call_count == 3 + assert len(list(tmp_path.iterdir())) == 3 + + oauth_client_class.reset_mock() + for group_id in [None, "group-a", "group-b"]: + provider = credentials_provider.external_browser(config(group_id)) + assert provider() == {"Authorization": f"Bearer {group_id or 'normal'}-token"} + + oauth_client_class.assert_not_called() + + +@pytest.mark.parametrize( + "token_endpoint", + [ + "https://workspace.cloud.databricks.com/oidc/v1/token", + "https://accounts.cloud.databricks.com/oidc/accounts/account-id/v1/token", + "https://db.cloud.databricks.com/oidc/accounts/account-id/v1/token", + ], + ids=["workspace", "account", "unified"], +) +def test_oauth_m2m_sends_group(requests_mock, token_endpoint): + """Verifies M2M sends assume_group to workspace, account, and unified token endpoints.""" + requests_mock.post( + token_endpoint, + json={"access_token": "token", "token_type": "Bearer", "expires_in": 3600}, + ) + cfg = Mock( + group_id="group-id", + client_id="client-id", + client_secret="client-secret", + databricks_oidc_endpoints=oauth.OidcEndpoints("unused", token_endpoint), + disable_async_token_refresh=True, + authorization_details=None, + ) + cfg.get_scopes_as_string.return_value = "all-apis" + + provider = credentials_provider.oauth_service_principal(cfg) + provider() + + token_form = parse_qs(requests_mock.last_request.text) + assert token_form["assume_group"] == ["group-id"] + + +def test_oauth_m2m_without_group_preserves_token_form(requests_mock): + """Verifies ungrouped M2M requests do not add an assume_group parameter.""" + token_endpoint = "https://workspace.cloud.databricks.com/oidc/v1/token" + requests_mock.post( + token_endpoint, + json={"access_token": "token", "token_type": "Bearer", "expires_in": 3600}, + ) + cfg = Mock( + group_id=None, + client_id="client-id", + client_secret="client-secret", + databricks_oidc_endpoints=oauth.OidcEndpoints("unused", token_endpoint), + disable_async_token_refresh=True, + authorization_details=None, + ) + cfg.get_scopes_as_string.return_value = "all-apis" + + credentials_provider.oauth_service_principal(cfg)() + + token_form = parse_qs(requests_mock.last_request.text, keep_blank_values=True) + assert "assume_group" not in token_form + + +def test_oauth_m2m_group_rejection_does_not_retry_without_group(requests_mock): + """Verifies a rejected group request is surfaced instead of retried as normal access.""" + token_endpoint = "https://accounts.cloud.databricks.com/oidc/accounts/account-id/v1/token" + token_request = requests_mock.post( + token_endpoint, + status_code=400, + headers={"Content-Type": "application/json"}, + json={"error": "invalid_request", "error_description": "assume_group is not supported"}, + ) + cfg = Mock( + group_id="group-id", + client_id="client-id", + client_secret="client-secret", + databricks_oidc_endpoints=oauth.OidcEndpoints("unused", token_endpoint), + disable_async_token_refresh=True, + authorization_details=None, + ) + cfg.get_scopes_as_string.return_value = "all-apis" + provider = credentials_provider.oauth_service_principal(cfg) + + with pytest.raises(ValueError, match="invalid_request: assume_group is not supported"): + provider() + + assert token_request.call_count == 1 + assert parse_qs(token_request.last_request.text)["assume_group"] == ["group-id"] + + +def test_oauth_m2m_reexchange_sends_group(requests_mock): + """Verifies M2M retains assume_group when an expired token triggers re-exchange.""" + token_endpoint = "https://workspace.cloud.databricks.com/oidc/v1/token" + token_responses = iter( + [ + {"access_token": "expired-role-token", "token_type": "Bearer", "expires_in": -1}, + {"access_token": "fresh-role-token", "token_type": "Bearer", "expires_in": 3600}, + ] + ) + token_request = requests_mock.post( + token_endpoint, + json=lambda _request, _context: next(token_responses), + ) + cfg = Mock( + group_id="group-id", + client_id="client-id", + client_secret="client-secret", + databricks_oidc_endpoints=oauth.OidcEndpoints("unused", token_endpoint), + disable_async_token_refresh=True, + authorization_details=None, + ) + cfg.get_scopes_as_string.return_value = "all-apis" + provider = credentials_provider.oauth_service_principal(cfg) + + assert provider() == {"Authorization": "Bearer expired-role-token"} + assert provider() == {"Authorization": "Bearer fresh-role-token"} + + assert token_request.call_count == 2 + for request in token_request.request_history: + assert parse_qs(request.text)["assume_group"] == ["group-id"] + + +def test_oauth_m2m_group_caches_are_isolated_per_provider(requests_mock): + """Verifies normal and grouped M2M providers cache only their own tokens.""" + token_endpoint = "https://workspace.cloud.databricks.com/oidc/v1/token" + + def issue_token(request, _context): + group_id = parse_qs(request.text, keep_blank_values=True).get("assume_group", ["normal"])[0] + return {"access_token": f"{group_id}-token", "token_type": "Bearer", "expires_in": 3600} + + token_request = requests_mock.post(token_endpoint, json=issue_token) + providers = {} + for group_id in [None, "group-a", "group-b"]: + cfg = Mock( + group_id=group_id, + client_id="client-id", + client_secret="client-secret", + databricks_oidc_endpoints=oauth.OidcEndpoints("unused", token_endpoint), + disable_async_token_refresh=True, + authorization_details=None, + ) + cfg.get_scopes_as_string.return_value = "all-apis" + providers[group_id] = credentials_provider.oauth_service_principal(cfg) + + for group_id, provider in providers.items(): + expected_headers = {"Authorization": f"Bearer {group_id or 'normal'}-token"} + assert provider() == expected_headers + assert provider() == expected_headers + + assert token_request.call_count == 3 + + +def test_oidc_supplier_sends_group(requests_mock): + """Verifies the shared OIDC credentials provider forwards assume_group.""" + token_endpoint = "https://workspace.cloud.databricks.com/oidc/v1/token" + requests_mock.post( + token_endpoint, + json={"access_token": "token", "token_type": "Bearer", "expires_in": 3600}, + ) + supplier = Mock() + supplier.get_oidc_token.return_value = "id-token" + cfg = Mock( + group_id="group-id", + token_audience="audience", + client_id="client-id", + databricks_oidc_endpoints=oauth.OidcEndpoints("unused", token_endpoint), + disable_async_token_refresh=True, + authorization_details=None, + ) + cfg.get_scopes_as_string.return_value = "all-apis" + + provider = credentials_provider._oidc_credentials_provider(cfg, lambda: supplier, "test OIDC") + provider() + + token_form = parse_qs(requests_mock.last_request.text) + assert token_form["assume_group"] == ["group-id"] + + +@pytest.mark.parametrize( + "provider", + [ + credentials_provider.pat_auth, + credentials_provider.basic_auth, + credentials_provider.runtime_native_auth, + credentials_provider.azure_service_principal, + credentials_provider.github_oidc_azure, + credentials_provider.google_credentials, + credentials_provider.google_id, + credentials_provider.azure_cli, + credentials_provider.databricks_cli, + credentials_provider.metadata_service, + credentials_provider.model_serving_auth, + ], + ids=lambda provider: provider.auth_type(), +) +def test_explicit_unsupported_auth_rejects_group(provider): + """Verifies explicitly selected normal-access strategies reject group-role requests.""" + auth_type = provider.auth_type() + cfg = Mock(group_id="group-id", auth_type=auth_type) + + with pytest.raises(ValueError, match=f'auth type "{auth_type}" does not support group role assumption'): + provider(cfg) + + +def test_default_credentials_skips_unsupported_auth_for_group(): + """Verifies default discovery skips normal-access auth and continues to group-capable auth.""" + + @credentials_provider.credentials_strategy("unsupported", [], supports_group=False) + def unsupported(_): + return lambda: {"Authorization": "normal"} + + @credentials_provider.credentials_strategy("supported", [], supports_group=True) + def supported(_): + return lambda: {"Authorization": "role"} + + strategy = credentials_provider.DefaultCredentials() + strategy._auth_providers = [unsupported, supported] + + provider = strategy(Mock(group_id="group-id", auth_type=None)) + + assert provider() == {"Authorization": "role"} + + +def test_default_credentials_group_fallback_uses_oauth_m2m_not_pat(requests_mock): + """Verifies default discovery selects grouped M2M instead of available PAT credentials.""" + token_endpoint = "https://workspace.cloud.databricks.com/oidc/v1/token" + token_request = requests_mock.post( + token_endpoint, + json={"access_token": "role-token", "token_type": "Bearer", "expires_in": 3600}, + ) + cfg = Mock( + group_id="group-id", + auth_type=None, + host="https://workspace.cloud.databricks.com", + token="normal-pat", + client_id="client-id", + client_secret="client-secret", + databricks_oidc_endpoints=oauth.OidcEndpoints("unused", token_endpoint), + disable_async_token_refresh=True, + authorization_details=None, + ) + cfg.get_scopes_as_string.return_value = "all-apis" + strategy = credentials_provider.DefaultCredentials() + strategy._auth_providers = [credentials_provider.pat_auth, credentials_provider.oauth_service_principal] + + provider = strategy(cfg) + + assert provider() == {"Authorization": "Bearer role-token"} + assert strategy.auth_type() == "oauth-m2m" + assert parse_qs(token_request.last_request.text)["assume_group"] == ["group-id"] + + +def test_default_credentials_group_exhaustion_keeps_generic_error(): + """Verifies exhausting default discovery retains its established generic error.""" + cfg = Mock(group_id="group-id", auth_type=None, host=None, scopes=None) + + with pytest.raises(ValueError, match="cannot configure default credentials"): + credentials_provider.DefaultCredentials()(cfg) def test_oidc_credentials_provider_invalid_id_token_source(): @@ -703,7 +1053,7 @@ class TestCloudAgnosticHosts: def test_azure_service_principal_with_cloud_agnostic_host(self, mocker): """Test that azure_service_principal works with cloud-agnostic hosts after removing is_azure requirement.""" # Mock Config with cloud-agnostic host - mock_cfg = Mock() + mock_cfg = Mock(group_id=None) mock_cfg.host = "https://api.databricks.com" # Cloud-agnostic host mock_cfg.azure_client_id = "test-azure-client-id" mock_cfg.azure_client_secret = "test-azure-secret" @@ -739,7 +1089,7 @@ def test_azure_service_principal_with_cloud_agnostic_host(self, mocker): def test_google_credentials_with_cloud_agnostic_host(self, mocker): """Test that google_credentials works with cloud-agnostic hosts after removing is_gcp check.""" # Mock Config with cloud-agnostic host - mock_cfg = Mock() + mock_cfg = Mock(group_id=None) mock_cfg.host = "https://api.databricks.com" # Cloud-agnostic host mock_cfg.google_credentials = '{"type": "service_account", "project_id": "test"}' mock_cfg.client_type = ClientType.WORKSPACE @@ -769,7 +1119,7 @@ def test_google_credentials_with_cloud_agnostic_host(self, mocker): def test_google_credentials_includes_sa_token_on_success(self, mocker): """Test that google_credentials includes GCP SA access token when refresh succeeds.""" - mock_cfg = Mock() + mock_cfg = Mock(group_id=None) mock_cfg.host = "https://api.databricks.com" mock_cfg.google_credentials = '{"type": "service_account", "project_id": "test"}' mock_cfg.disable_async_token_refresh = True @@ -796,7 +1146,7 @@ def test_google_credentials_includes_sa_token_on_success(self, mocker): def test_google_credentials_warns_on_sa_token_failure(self, mocker): """Test that google_credentials logs warning and omits SA token when refresh fails.""" - mock_cfg = Mock() + mock_cfg = Mock(group_id=None) mock_cfg.host = "https://api.databricks.com" mock_cfg.google_credentials = '{"type": "service_account", "project_id": "test"}' mock_cfg.disable_async_token_refresh = True @@ -827,7 +1177,7 @@ def test_google_credentials_warns_on_sa_token_failure(self, mocker): def test_google_id_with_cloud_agnostic_host(self, mocker): """Test that google_id works with cloud-agnostic hosts after removing is_gcp check.""" # Mock Config with cloud-agnostic host - mock_cfg = Mock() + mock_cfg = Mock(group_id=None) mock_cfg.host = "https://api.databricks.com" # Cloud-agnostic host mock_cfg.google_service_account = "test-sa@project.iam.gserviceaccount.com" mock_cfg.client_type = ClientType.WORKSPACE @@ -864,7 +1214,7 @@ def test_google_id_with_cloud_agnostic_host(self, mocker): def test_google_id_includes_sa_token_on_success(self, mocker): """Test that google_id includes GCP SA access token when refresh succeeds.""" - mock_cfg = Mock() + mock_cfg = Mock(group_id=None) mock_cfg.host = "https://api.databricks.com" mock_cfg.google_service_account = "test-sa@project.iam.gserviceaccount.com" @@ -896,7 +1246,7 @@ def test_google_id_includes_sa_token_on_success(self, mocker): def test_google_id_warns_on_sa_token_failure(self, mocker): """Test that google_id logs warning and omits SA token when refresh fails.""" - mock_cfg = Mock() + mock_cfg = Mock(group_id=None) mock_cfg.host = "https://api.databricks.com" mock_cfg.google_service_account = "test-sa@project.iam.gserviceaccount.com" @@ -935,7 +1285,7 @@ def test_github_oidc_azure_with_cloud_agnostic_host(self, mocker): mocker.patch.dict("os.environ", {"ACTIONS_ID_TOKEN_REQUEST_TOKEN": "test-token"}) # Mock Config with cloud-agnostic host - mock_cfg = Mock() + mock_cfg = Mock(group_id=None) mock_cfg.host = "https://api.databricks.com" # Cloud-agnostic host mock_cfg.azure_client_id = "test-azure-client-id" mock_cfg.azure_tenant_id = None # Will be auto-detected @@ -982,7 +1332,7 @@ def test_github_oidc_azure_with_cloud_agnostic_host(self, mocker): def test_azure_cli_requires_effective_azure_login_app_id(self, mocker): """Test that azure_cli now requires effective_azure_login_app_id instead of is_azure.""" # Mock Config with cloud-agnostic host - mock_cfg = Mock() + mock_cfg = Mock(group_id=None) mock_cfg.host = "https://api.databricks.com" # Cloud-agnostic host mock_cfg.azure_tenant_id = "test-tenant-id" mock_cfg.azure_workspace_resource_id = None @@ -1017,7 +1367,7 @@ def test_azure_cli_requires_effective_azure_login_app_id(self, mocker): def test_azure_cli_returns_none_without_effective_azure_login_app_id(self): """Test that azure_cli returns None when effective_azure_login_app_id is not set.""" # Mock Config without effective_azure_login_app_id - mock_cfg = Mock() + mock_cfg = Mock(group_id=None) mock_cfg.host = "https://api.databricks.com" mock_cfg.effective_azure_login_app_id = None # Not set diff --git a/tests/test_notebook_oauth.py b/tests/test_notebook_oauth.py index cffc502c9..e12740bde 100644 --- a/tests/test_notebook_oauth.py +++ b/tests/test_notebook_oauth.py @@ -5,6 +5,7 @@ import types from datetime import datetime, timedelta from typing import Dict +from urllib.parse import parse_qs import pytest @@ -104,6 +105,46 @@ def test_runtime_oauth_success_scenarios( assert headers["Authorization"] == "Bearer exchanged-oauth-token" +def test_runtime_oauth_with_group_reexchange_uses_fresh_notebook_pat( + mock_runtime_env, mock_runtime_native_auth, monkeypatch, requests_mock +): + """Verifies grouped runtime OAuth retains its role and fetches a fresh PAT per exchange.""" + token_endpoint = "https://test.cloud.databricks.com/oidc/v1/token" + token_responses = iter( + [ + {"access_token": "expired-role-token", "token_type": "Bearer", "expires_in": -1}, + {"access_token": "fresh-role-token", "token_type": "Bearer", "expires_in": 3600}, + ] + ) + token_request = requests_mock.post( + token_endpoint, + json=lambda _request, _context: next(token_responses), + ) + # runtime_oauth verifies PAT availability before the refreshable source fetches one per exchange. + notebook_pats = iter(["preflight-pat", "first-exchange-pat", "second-exchange-pat"]) + + def init_runtime_native_auth(): + return "https://test.cloud.databricks.com", lambda: {"Authorization": f"Bearer {next(notebook_pats)}"} + + monkeypatch.setattr(sys.modules["databricks.sdk.runtime"], "init_runtime_native_auth", init_runtime_native_auth) + cfg = Config( + host="https://test.cloud.databricks.com", + scopes="all-apis", + group_id="group-id", + credentials_strategy=MockCredentialsStrategy(), + ) + + provider = runtime_oauth(cfg) + + assert provider() == {"Authorization": "Bearer expired-role-token"} + assert provider() == {"Authorization": "Bearer fresh-role-token"} + + assert token_request.call_count == 2 + token_forms = [parse_qs(request.text) for request in token_request.request_history] + assert [form["subject_token"] for form in token_forms] == [["first-exchange-pat"], ["second-exchange-pat"]] + assert [form["assume_group"] for form in token_forms] == [["group-id"], ["group-id"]] + + @pytest.mark.parametrize( "scopes", [ diff --git a/tests/test_oauth.py b/tests/test_oauth.py index be15d1a1b..2a9023f3d 100644 --- a/tests/test_oauth.py +++ b/tests/test_oauth.py @@ -1,9 +1,14 @@ +import pathlib +from urllib.parse import parse_qs, urlparse + import pytest from databricks.sdk._base_client import _BaseClient from databricks.sdk.oauth import ( HostMetadata, + OAuthClient, OidcEndpoints, + PATOAuthTokenExchange, TokenCache, get_account_endpoints, get_endpoints_from_url, @@ -79,6 +84,114 @@ def test_token_cache_filename_no_delimiter_collision(): ) +def test_token_cache_group_isolation_preserves_normal_cache_key(): + """Verifies grouped caches are distinct without changing the legacy normal cache key.""" + common_args = dict( + host="http://localhost:", + client_id="abc", + oidc_endpoints=OidcEndpoints("http://localhost:1234", "http://localhost:1234"), + ) + normal = TokenCache(**common_args).filename + + # Changing this hash would orphan existing ungrouped credentials on disk. + assert pathlib.Path(normal).name == "676cce09b3b66924475b6ad807598b3550b6731a3debca900d42ed88f462e366.json" + assert TokenCache(group_id=None, **common_args).filename == normal + assert TokenCache(group_id="group-a", **common_args).filename != normal + assert ( + TokenCache(group_id="group-a", **common_args).filename != TokenCache(group_id="group-b", **common_args).filename + ) + + +@pytest.mark.parametrize( + "authorization_endpoint", + [ + "https://workspace.cloud.databricks.com/oidc/v1/authorize", + "https://accounts.cloud.databricks.com/oidc/accounts/account-id/v1/authorize", + "https://db.cloud.databricks.com/oidc/accounts/account-id/v1/authorize", + ], + ids=["workspace", "account", "unified"], +) +def test_oauth_client_adds_group_to_authorization_only(requests_mock, authorization_endpoint): + """Verifies browser OAuth sends assume_group only on the authorization request.""" + token_endpoint = authorization_endpoint.replace("authorize", "token") + requests_mock.post( + token_endpoint, + json={"access_token": "token", "token_type": "Bearer", "expires_in": 3600}, + ) + oauth_client = OAuthClient( + OidcEndpoints(authorization_endpoint, token_endpoint), + "http://localhost:8020", + "client-id", + group_id="group-id", + ) + + consent = oauth_client.initiate_consent() + authorization_query = parse_qs(urlparse(consent.authorization_url).query) + assert authorization_query["assume_group"] == ["group-id"] + + consent.exchange("code", consent._state) + token_form = parse_qs(requests_mock.last_request.text, keep_blank_values=True) + assert "assume_group" not in token_form + + +def test_grouped_oauth_session_refresh_omits_group(requests_mock): + """Verifies a grouped browser session refreshes without resending assume_group.""" + token_endpoint = "https://workspace.cloud.databricks.com/oidc/v1/token" + token_responses = iter( + [ + { + "access_token": "expired-role-token", + "refresh_token": "role-refresh-token", + "token_type": "Bearer", + "expires_in": -1, + }, + {"access_token": "refreshed-role-token", "token_type": "Bearer", "expires_in": 3600}, + ] + ) + token_request = requests_mock.post( + token_endpoint, + json=lambda _request, _context: next(token_responses), + ) + oauth_client = OAuthClient( + OidcEndpoints("https://workspace.cloud.databricks.com/oidc/v1/authorize", token_endpoint), + "http://localhost:8020", + "client-id", + group_id="group-id", + ) + + consent = oauth_client.initiate_consent() + credentials = consent.exchange("code", consent._state) + token = credentials.token() + + assert token.access_token == "refreshed-role-token" + assert token_request.call_count == 2 + authorization_code_form = parse_qs(token_request.request_history[0].text, keep_blank_values=True) + refresh_form = parse_qs(token_request.request_history[1].text, keep_blank_values=True) + assert "assume_group" not in authorization_code_form + assert refresh_form["grant_type"] == ["refresh_token"] + assert refresh_form["refresh_token"] == ["role-refresh-token"] + assert "assume_group" not in refresh_form + + +def test_pat_oauth_exchange_sends_group_in_token_form(requests_mock): + """Verifies PAT-to-OAuth exchange includes the requested group in its token form.""" + requests_mock.post( + "https://workspace.cloud.databricks.com/oidc/v1/token", + json={"access_token": "token", "token_type": "Bearer", "expires_in": 3600}, + ) + source = PATOAuthTokenExchange( + get_original_token=lambda: "pat", + host="https://workspace.cloud.databricks.com", + scopes="all-apis", + group_id="group-id", + ) + + source.token() + + token_form = parse_qs(requests_mock.last_request.text) + assert token_form["assume_group"] == ["group-id"] + + def test_account_oidc_endpoints(requests_mock): requests_mock.get( "https://accounts.cloud.databricks.com/oidc/accounts/abc-123/.well-known/oauth-authorization-server", diff --git a/tests/test_oidc.py b/tests/test_oidc.py index 4aad82f4c..80f26f5a0 100644 --- a/tests/test_oidc.py +++ b/tests/test_oidc.py @@ -1,6 +1,7 @@ from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Optional, Tuple +from urllib.parse import parse_qs import pytest @@ -187,3 +188,85 @@ def test_databricks_oidc_token_source_missing_host_raises(): ) with pytest.raises(ValueError, match="missing Host"): ts.token() + + +@pytest.mark.parametrize( + "token_endpoint", + [ + "https://workspace.cloud.databricks.com/oidc/v1/token", + "https://accounts.cloud.databricks.com/oidc/accounts/account-id/v1/token", + "https://db.cloud.databricks.com/oidc/accounts/account-id/v1/token", + ], + ids=["workspace", "account", "unified"], +) +def test_databricks_oidc_token_source_sends_group(requests_mock, token_endpoint): + """Verifies WIF sends assume_group to workspace, account, and unified token endpoints.""" + requests_mock.post( + token_endpoint, + json={"access_token": "token", "token_type": "Bearer", "expires_in": 3600}, + ) + source = oidc.DatabricksOidcTokenSource( + host="https://example.cloud.databricks.com", + token_endpoint=token_endpoint, + id_token_source=_CountingIdTokenSource(), + client_id="client-id", + group_id="group-id", + ) + + source.token() + + token_form = parse_qs(requests_mock.last_request.text) + assert token_form["assume_group"] == ["group-id"] + + +def test_databricks_oidc_token_source_reexchange_sends_group(requests_mock): + """Verifies WIF retains assume_group when an expired token triggers re-exchange.""" + token_endpoint = "https://workspace.cloud.databricks.com/oidc/v1/token" + token_request = requests_mock.post( + token_endpoint, + json={"access_token": "expired-token", "token_type": "Bearer", "expires_in": -1}, + ) + source = oidc.DatabricksOidcTokenSource( + host="https://workspace.cloud.databricks.com", + token_endpoint=token_endpoint, + id_token_source=_CountingIdTokenSource(), + client_id="client-id", + group_id="group-id", + disable_async=True, + ) + + source.token() + source.token() + + assert token_request.call_count == 2 + for request in token_request.request_history: + assert parse_qs(request.text)["assume_group"] == ["group-id"] + + +def test_databricks_oidc_token_source_group_caches_are_isolated(requests_mock): + """Verifies normal and grouped WIF token sources cache only their own tokens.""" + token_endpoint = "https://workspace.cloud.databricks.com/oidc/v1/token" + + def issue_token(request, _context): + group_id = parse_qs(request.text, keep_blank_values=True).get("assume_group", ["normal"])[0] + return {"access_token": f"{group_id}-token", "token_type": "Bearer", "expires_in": 3600} + + token_request = requests_mock.post(token_endpoint, json=issue_token) + sources = { + group_id: oidc.DatabricksOidcTokenSource( + host="https://workspace.cloud.databricks.com", + token_endpoint=token_endpoint, + id_token_source=_CountingIdTokenSource(), + client_id="client-id", + group_id=group_id, + disable_async=True, + ) + for group_id in [None, "group-a", "group-b"] + } + + for group_id, source in sources.items(): + expected_token = f"{group_id or 'normal'}-token" + assert source.token().access_token == expected_token + assert source.token().access_token == expected_token + + assert token_request.call_count == 3 diff --git a/tests/testdata/.databrickscfg b/tests/testdata/.databrickscfg index 2ffb627ae..07ddcae0a 100644 --- a/tests/testdata/.databrickscfg +++ b/tests/testdata/.databrickscfg @@ -50,3 +50,7 @@ scopes = clusters [scope-multiple] host = https://example.cloud.databricks.com scopes = clusters, jobs, pipelines, iam:read, files:read, mlflow, model-serving:read + +[group-role] +host = https://example.cloud.databricks.com +group_id = profile-group