[Internal] Resolve token_audience from token_federation_default_oidc_audiences in host metadata

NEXT_CHANGELOG.md

@@ -18,6 +18,7 @@
 ### Internal Changes
 * Replace the async-disabling mechanism on token refresh failure with a 1-minute retry backoff. Previously, a single failed async refresh would disable proactive token renewal until the token expired. Now, the SDK waits a short cooldown period and retries, improving resilience to transient errors.
 * Extract `_resolve_profile` to simplify config file loading and improve `__settings__` error messages.
+* Resolve `token_audience` from the `token_federation_default_oidc_audiences` field in the host metadata discovery endpoint, removing the need for explicit audience configuration.
 
 ### API Changes
 * Add `create_catalog()`, `create_synced_table()`, `delete_catalog()`, `delete_synced_table()`, `get_catalog()` and `get_synced_table()` methods for [w.postgres](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/postgres/postgres.html) workspace-level service.

databricks/sdk/config.py

@@ -655,10 +655,15 @@ def _resolve_host_metadata(self) -> None:
         if not self.cloud and meta.cloud:
             logger.debug(f"Resolved cloud from host metadata: {meta.cloud.value}")
             self.cloud = meta.cloud
+        if not self.token_audience and meta.token_federation_default_oidc_audiences:
+            audience = meta.token_federation_default_oidc_audiences[0]
+            logger.debug(
+                f"Resolved token_audience from host metadata token_federation_default_oidc_audiences: {audience}"
+            )
+            self.token_audience = audience
         # Account hosts use account_id as the OIDC token audience instead of the token endpoint.
         # This is a special case: when the metadata has no workspace_id, the host is acting as an
         # account-level endpoint and the audience must be scoped to the account.
-        # TODO: Add explicit audience to the metadata discovery endpoint.
         if not self.token_audience and not meta.workspace_id and self.account_id:
             logger.debug(f"Setting token_audience to account_id for account host: {self.account_id}")
             self.token_audience = self.account_id

databricks/sdk/oauth.py

@@ -448,6 +448,7 @@ class HostMetadata:
     account_id: Optional[str] = None
     workspace_id: Optional[str] = None
     cloud: Optional[Cloud] = None
+    token_federation_default_oidc_audiences: Optional[List[str]] = None
 
     @staticmethod
     def from_dict(d: dict) -> "HostMetadata":
@@ -456,6 +457,7 @@ def from_dict(d: dict) -> "HostMetadata":
             account_id=d.get("account_id"),
             workspace_id=d.get("workspace_id"),
             cloud=Cloud.parse(d.get("cloud", "")),
+            token_federation_default_oidc_audiences=d.get("token_federation_default_oidc_audiences"),
         )
 
     def as_dict(self) -> dict:
@@ -464,6 +466,7 @@ def as_dict(self) -> dict:
             "account_id": self.account_id,
             "workspace_id": self.workspace_id,
             "cloud": self.cloud.value if self.cloud else None,
+            "token_federation_default_oidc_audiences": self.token_federation_default_oidc_audiences,
         }
 
 

tests/test_config.py

@@ -973,3 +973,77 @@ def test_resolve_host_metadata_does_not_overwrite_token_audience(mocker):
         token_audience="custom-audience",
     )
     assert config.token_audience == "custom-audience"
+
+
+# ---------------------------------------------------------------------------
+# token_federation_default_oidc_audiences resolution from host metadata
+# ---------------------------------------------------------------------------
+
+
+def test_resolve_host_metadata_sets_token_audience_from_token_federation_default_oidc_audiences(mocker):
+    """token_audience is set from token_federation_default_oidc_audiences in host metadata."""
+    mocker.patch(
+        "databricks.sdk.config.get_host_metadata",
+        return_value=HostMetadata.from_dict(
+            {
+                "oidc_endpoint": f"{_DUMMY_WS_HOST}/oidc",
+                "account_id": _DUMMY_ACCOUNT_ID,
+                "workspace_id": _DUMMY_WORKSPACE_ID,
+                "token_federation_default_oidc_audiences": [f"{_DUMMY_WS_HOST}/oidc/v1/token"],
+            }
+        ),
+    )
+    config = Config(host=_DUMMY_WS_HOST, token="t")
+    assert config.token_audience == f"{_DUMMY_WS_HOST}/oidc/v1/token"
+
+
+def test_resolve_host_metadata_token_federation_default_oidc_audiences_takes_priority_over_account_id_fallback(mocker):
+    """token_federation_default_oidc_audiences takes priority over the account_id fallback."""
+    mocker.patch(
+        "databricks.sdk.config.get_host_metadata",
+        return_value=HostMetadata.from_dict(
+            {
+                "oidc_endpoint": f"{_DUMMY_ACC_HOST}/oidc/accounts/{_DUMMY_ACCOUNT_ID}",
+                "account_id": _DUMMY_ACCOUNT_ID,
+                "token_federation_default_oidc_audiences": ["custom-audience-from-server"],
+            }
+        ),
+    )
+    config = Config(host=_DUMMY_ACC_HOST, token="t", account_id=_DUMMY_ACCOUNT_ID)
+    # token_federation_default_oidc_audiences should take priority over the account_id fallback
+    assert config.token_audience == "custom-audience-from-server"
+
+
+def test_resolve_host_metadata_token_federation_default_oidc_audiences_does_not_override_existing_token_audience(
+    mocker,
+):
+    """An explicitly set token_audience is not overwritten by token_federation_default_oidc_audiences."""
+    mocker.patch(
+        "databricks.sdk.config.get_host_metadata",
+        return_value=HostMetadata.from_dict(
+            {
+                "oidc_endpoint": f"{_DUMMY_WS_HOST}/oidc",
+                "account_id": _DUMMY_ACCOUNT_ID,
+                "workspace_id": _DUMMY_WORKSPACE_ID,
+                "token_federation_default_oidc_audiences": [f"{_DUMMY_WS_HOST}/oidc/v1/token"],
+            }
+        ),
+    )
+    config = Config(host=_DUMMY_WS_HOST, token="t", token_audience="my-custom-audience")
+    assert config.token_audience == "my-custom-audience"
+
+
+def test_resolve_host_metadata_falls_back_to_account_id_when_no_token_federation_default_oidc_audiences(mocker):
+    """When no token_federation_default_oidc_audiences and no workspace_id, falls back to account_id."""
+    mocker.patch(
+        "databricks.sdk.config.get_host_metadata",
+        return_value=HostMetadata.from_dict(
+            {
+                "oidc_endpoint": f"{_DUMMY_ACC_HOST}/oidc/accounts/{_DUMMY_ACCOUNT_ID}",
+                "account_id": _DUMMY_ACCOUNT_ID,
+            }
+        ),
+    )
+    config = Config(host=_DUMMY_ACC_HOST, token="t", account_id=_DUMMY_ACCOUNT_ID)
+    # No token_federation_default_oidc_audiences and no workspace_id → falls back to account_id
+    assert config.token_audience == _DUMMY_ACCOUNT_ID