diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index 7382d77bc..b4fb884ec 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -457,6 +457,7 @@ def oidc_credentials_provider(cfg, id_token_source: oidc.IdTokenSource) -> Optio disable_async=cfg.disable_async_token_refresh, scopes=cfg.get_scopes_as_string(), group_id=cfg.group_id, + authorization_details=cfg.authorization_details, ) def refreshed_headers() -> Dict[str, str]: @@ -469,6 +470,19 @@ def token() -> oauth.Token: return OAuthCredentialsProvider(refreshed_headers, token) +class _OidcTokenSupplierIdTokenSource(oidc.IdTokenSource): + def __init__(self, supplier: Any, audience: str, provider_name: str): + self._supplier = supplier + self._audience = audience + self._provider_name = provider_name + + def id_token(self) -> oidc.IdToken: + token = self._supplier.get_oidc_token(self._audience) + if not token: + raise ValueError(f"Cannot get {self._provider_name} token") + return oidc.IdToken(jwt=token) + + def _oidc_credentials_provider( cfg: "Config", supplier_factory: Callable[[], Any], provider_name: str ) -> Optional[CredentialsProvider]: @@ -495,47 +509,11 @@ def _oidc_credentials_provider( if audience is None: audience = cfg.databricks_oidc_endpoints.token_endpoint - # Try to get an OIDC token. If no supplier returns a token, we cannot use this authentication mode. - id_token = supplier.get_oidc_token(audience) - if not id_token: - logger.debug(f"{provider_name}: no token available, skipping authentication method") - return None - - logger.info(f"Configured {provider_name} authentication") - - def token_source_for(audience: str) -> oauth.TokenSource: - id_token = supplier.get_oidc_token(audience) - if not id_token: - # Should not happen, since we checked it above. - raise Exception(f"Cannot get {provider_name} token") - - 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=endpoint_params, - 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) + id_token_source = _OidcTokenSupplierIdTokenSource(supplier, audience, provider_name) + provider = oidc_credentials_provider(cfg, id_token_source) + if provider is not None: + logger.info(f"Configured {provider_name} authentication") + return provider @oauth_credentials_strategy("github-oidc", ["host", "client_id"], supports_group=True) diff --git a/databricks/sdk/oidc.py b/databricks/sdk/oidc.py index 1a3c8d3d2..8a6379725 100644 --- a/databricks/sdk/oidc.py +++ b/databricks/sdk/oidc.py @@ -151,6 +151,8 @@ class DatabricksOidcTokenSource(oauth.Refreshable): audience : Optional[str], optional The audience of the Databricks OIDC application. Only used for Workspace level tokens. + authorization_details : Optional[str], optional + JSON-encoded authorization details to include in the token exchange. """ def __init__( @@ -164,6 +166,7 @@ def __init__( disable_async: bool = False, scopes: Optional[str] = None, group_id: Optional[str] = None, + authorization_details: Optional[str] = None, ): self._host = host self._id_token_source = id_token_source @@ -173,6 +176,7 @@ def __init__( self._audience = audience self._scopes = scopes self._group_id = group_id + self._authorization_details = authorization_details # 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) @@ -224,6 +228,7 @@ def _exchange_id_token(self, id_token: IdToken) -> oauth.Token: scopes=self._scopes, use_params=True, disable_async=self._disable_async, + authorization_details=self._authorization_details, ) return client.token() diff --git a/tests/test_credentials_provider.py b/tests/test_credentials_provider.py index 1e8fd2df5..e14236301 100644 --- a/tests/test_credentials_provider.py +++ b/tests/test_credentials_provider.py @@ -524,6 +524,45 @@ def test_oidc_supplier_sends_group(requests_mock): assert token_form["assume_group"] == ["group-id"] +@pytest.mark.parametrize( + ("provider", "supplier_class"), + [ + (credentials_provider.github_oidc, "GitHubOIDCTokenSupplier"), + (credentials_provider.azure_devops_oidc, "AzureDevOpsOIDCTokenSupplier"), + ], +) +def test_ci_oidc_provider_caches_exchanged_token(mocker, requests_mock, provider, supplier_class): + token_endpoint = "https://workspace.cloud.databricks.com/oidc/v1/token" + token_request = 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" + mocker.patch.object(credentials_provider.oidc_token_supplier, supplier_class, return_value=supplier) + cfg = Mock( + auth_type=provider.auth_type(), + host="https://workspace.cloud.databricks.com", + group_id=None, + token_audience="audience", + client_id="client-id", + account_id=None, + 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(cfg) + + assert credentials() == {"Authorization": "Bearer token"} + assert credentials() == {"Authorization": "Bearer token"} + assert credentials.oauth_token().access_token == "token" + assert token_request.call_count == 1 + assert supplier.get_oidc_token.call_count == 2 + supplier.get_oidc_token.assert_called_with("audience") + + @pytest.mark.parametrize( "provider", [ diff --git a/tests/test_oidc.py b/tests/test_oidc.py index 80f26f5a0..2d81dd6e2 100644 --- a/tests/test_oidc.py +++ b/tests/test_oidc.py @@ -219,6 +219,26 @@ def test_databricks_oidc_token_source_sends_group(requests_mock, token_endpoint) assert token_form["assume_group"] == ["group-id"] +def test_databricks_oidc_token_source_sends_authorization_details(requests_mock): + 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}, + ) + source = oidc.DatabricksOidcTokenSource( + host="https://workspace.cloud.databricks.com", + token_endpoint=token_endpoint, + id_token_source=_CountingIdTokenSource(), + client_id="client-id", + authorization_details='[{"type":"example"}]', + ) + + source.token() + + token_form = parse_qs(requests_mock.last_request.text) + assert token_form["authorization_details"] == ['[{"type":"example"}]'] + + 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"