Cache tokens minted by DatabricksOidcTokenSource (databricks#1460)

## What

`DatabricksOidcTokenSource` (Workload Identity Federation / account-wide
token federation) minted a fresh `/oidc/v1/token` exchange on **every**
authenticated API call. Its `token()` always called
`_exchange_id_token()`, which constructed a brand-new
`oauth.ClientCredentials` each time — discarding the `Refreshable`
caching/locking/refresh machinery already built into it.

This change makes `DatabricksOidcTokenSource` extend `oauth.Refreshable`
and moves the exchange logic into `refresh()`, so the minted token is
cached and reused until it is stale or expired (mirroring the existing
M2M/PAT flows). Each refresh re-fetches the ID token from the
`id_token_source`, so short-lived **rotating** ID tokens are handled
transparently.

## Why

The previous behaviour:
- added an extra RTT to every request,
- amplified transient federation-policy errors by the number of API
calls,
- and hit OIDC token-endpoint rate limits.

A real-world 470 MiB Volume upload minted dozens of tokens in under 30s,
with bursts of 6–7 concurrent mints.

## Tests

Added unit tests in `tests/test_oidc.py`:
- token is minted once and reused across calls (caching),
- an expired token triggers a re-exchange that fetches a **fresh** ID
token (rotation),
- missing host still raises `ValueError`.

All existing auth tests pass (`test_oauth`, `test_oidc`,
`test_credentials_provider`, `test_oidc_token_supplier`). Added a
`NEXT_CHANGELOG.md` Bug Fixes entry.

Signed-off-by: Hector Castejon Diaz <hector.castejon@databricks.com>

Author: hectorcast-db

NEXT_CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bug Fixes
 
+* Cache tokens minted by `DatabricksOidcTokenSource` (Workload Identity Federation / account-wide token federation). Previously a fresh `/oidc/v1/token` exchange was performed on every authenticated API call, adding latency, amplifying transient federation-policy errors, and hitting OIDC token-endpoint rate limits. The token source now reuses the cached token until it is stale or expired, fetching a fresh ID token on each refresh to support rotation.
+
 ### Documentation
 
 ### Breaking Changes

databricks/sdk/oidc.py

@@ -123,9 +123,15 @@ def id_token(self) -> IdToken:
         return IdToken(jwt=token)
 
 
-class DatabricksOidcTokenSource(oauth.TokenSource):
+class DatabricksOidcTokenSource(oauth.Refreshable):
     """A TokenSource which exchanges a token using Workload Identity Federation.
 
+    The exchanged token is cached and reused across calls. It is only refreshed
+    when it is stale or expired, mirroring the M2M/PAT auth flows. This avoids
+    minting a fresh token on every authenticated API call. Each refresh fetches a
+    fresh ID token from the ``id_token_source``, so short-lived (rotating) ID
+    tokens are handled transparently.
+
     Parameters
     ----------
     host : str
@@ -164,15 +170,20 @@ def __init__(
         self._client_id = client_id
         self._account_id = account_id
         self._audience = audience
-        self._disable_async = disable_async
         self._scopes = scopes
+        # 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)
+
+    def refresh(self) -> oauth.Token:
+        """Mint a fresh token by exchanging the ID token.
 
-    def token(self) -> oauth.Token:
-        """Get a token by exchanging the ID token.
+        Called by the base :class:`oauth.Refreshable` only when the cached token
+        is missing, stale, or expired. The result is cached by the base class.
 
         Returns
         -------
-        dict
+        oauth.Token
             The exchanged token.
 
         Raises

tests/test_oidc.py

@@ -1,9 +1,10 @@
 from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
 from typing import Optional, Tuple
 
 import pytest
 
-from databricks.sdk import oidc
+from databricks.sdk import oauth, oidc
 
 
 @dataclass
@@ -107,3 +108,82 @@ def test_file_id_token_source(test_case: FileTestCase, tmp_path):
             source.id_token()
     else:
         assert source.id_token() == test_case.want
+
+
+class _CountingIdTokenSource(oidc.IdTokenSource):
+    """An IdTokenSource that counts how many times it is invoked and returns a
+    fresh jwt each time, so we can assert ID-token rotation on refresh."""
+
+    def __init__(self):
+        self.calls = 0
+
+    def id_token(self) -> oidc.IdToken:
+        self.calls += 1
+        return oidc.IdToken(jwt=f"id-token-{self.calls}")
+
+
+def _make_token_source(id_token_source, exchanges, expiry):
+    """Build a DatabricksOidcTokenSource whose exchange step is stubbed to record
+    the ID token it received and return a token with the given expiry."""
+    ts = oidc.DatabricksOidcTokenSource(
+        host="https://test.cloud.databricks.com",
+        token_endpoint="https://test.cloud.databricks.com/oidc/v1/token",
+        id_token_source=id_token_source,
+        client_id="test-client-id",
+    )
+
+    def _exchange(id_token: oidc.IdToken) -> oauth.Token:
+        exchanges.append(id_token.jwt)
+        return oauth.Token(access_token=f"exchanged-{id_token.jwt}", token_type="Bearer", expiry=expiry)
+
+    ts._exchange_id_token = _exchange
+    return ts
+
+
+def test_databricks_oidc_token_source_caches_token():
+    # A valid token (1h TTL) should be minted once and reused across calls,
+    # rather than exchanging the ID token on every token() call.
+    id_token_source = _CountingIdTokenSource()
+    exchanges = []
+    ts = _make_token_source(
+        id_token_source,
+        exchanges,
+        expiry=datetime.now(tz=timezone.utc) + timedelta(hours=1),
+    )
+
+    first = ts.token()
+    for _ in range(5):
+        again = ts.token()
+        assert again.access_token == first.access_token
+
+    assert exchanges == ["id-token-1"]
+    assert id_token_source.calls == 1
+
+
+def test_databricks_oidc_token_source_refreshes_when_expired():
+    # An already-expired token must trigger a re-exchange, and the refresh must
+    # fetch a *fresh* ID token (rotation), not reuse the original jwt.
+    id_token_source = _CountingIdTokenSource()
+    exchanges = []
+    ts = _make_token_source(
+        id_token_source,
+        exchanges,
+        expiry=datetime.now(tz=timezone.utc) - timedelta(seconds=1),
+    )
+
+    ts.token()
+    ts.token()
+
+    assert exchanges == ["id-token-1", "id-token-2"]
+    assert id_token_source.calls == 2
+
+
+def test_databricks_oidc_token_source_missing_host_raises():
+    ts = oidc.DatabricksOidcTokenSource(
+        host="",
+        token_endpoint="https://test.cloud.databricks.com/oidc/v1/token",
+        id_token_source=_CountingIdTokenSource(),
+        client_id="test-client-id",
+    )
+    with pytest.raises(ValueError, match="missing Host"):
+        ts.token()