Fix CI OIDC token caching

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)

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()

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",
     [

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"