Resolve token_audience from default_oidc_audience in host metadata

Signed-off-by: Tanmay Rustagi <tanmay.rustagi@databricks.com>

Author: tanmay-db

databricks/sdk/config.py

@@ -671,10 +671,12 @@ def _resolve_host_metadata(self) -> None:
             if resolved is not None:
                 logger.debug(f"Resolved host_type from host metadata: {meta.host_type}")
                 self._resolved_host_type = resolved
+        if not self.token_audience and meta.default_oidc_audience:
+            logger.debug(f"Resolved token_audience from host metadata default_oidc_audience: {meta.default_oidc_audience}")
+            self.token_audience = meta.default_oidc_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

@@ -449,6 +449,7 @@ class HostMetadata:
     workspace_id: Optional[str] = None
     cloud: Optional[Cloud] = None
     host_type: Optional[str] = None
+    default_oidc_audience: Optional[str] = None
 
     @staticmethod
     def from_dict(d: dict) -> "HostMetadata":
@@ -458,6 +459,7 @@ def from_dict(d: dict) -> "HostMetadata":
             workspace_id=d.get("workspace_id"),
             cloud=Cloud.parse(d.get("cloud", "")),
             host_type=d.get("host_type"),
+            default_oidc_audience=d.get("default_oidc_audience"),
         )
 
     def as_dict(self) -> dict:
@@ -467,6 +469,7 @@ def as_dict(self) -> dict:
             "workspace_id": self.workspace_id,
             "cloud": self.cloud.value if self.cloud else None,
             "host_type": self.host_type,
+            "default_oidc_audience": self.default_oidc_audience,
         }
 
 

tests/test_config.py

@@ -1097,3 +1097,75 @@ def test_resolve_host_metadata_resolved_host_type_unknown_string(mocker):
     )
     config = Config(host=_DUMMY_WS_HOST, token="t")
     assert config._resolved_host_type is None
+
+
+# ---------------------------------------------------------------------------
+# default_oidc_audience resolution from host metadata
+# ---------------------------------------------------------------------------
+
+
+def test_resolve_host_metadata_sets_token_audience_from_default_oidc_audience(mocker):
+    """token_audience is set from default_oidc_audience 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,
+                "default_oidc_audience": 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_default_oidc_audience_takes_priority_over_account_id_fallback(mocker):
+    """default_oidc_audience 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,
+                "default_oidc_audience": "custom-audience-from-server",
+            }
+        ),
+    )
+    config = Config(host=_DUMMY_ACC_HOST, token="t", account_id=_DUMMY_ACCOUNT_ID)
+    # default_oidc_audience should take priority over the account_id fallback
+    assert config.token_audience == "custom-audience-from-server"
+
+
+def test_resolve_host_metadata_default_oidc_audience_does_not_override_existing_token_audience(mocker):
+    """An explicitly set token_audience is not overwritten by default_oidc_audience."""
+    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,
+                "default_oidc_audience": 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_default_oidc_audience(mocker):
+    """When no default_oidc_audience 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 default_oidc_audience and no workspace_id → falls back to account_id
+    assert config.token_audience == _DUMMY_ACCOUNT_ID