From 8259278b98634623e45fe158d8d67c8819840bd3 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Tue, 14 Apr 2026 11:57:55 +0200 Subject: [PATCH 1/2] Resolve token_audience from host metadata token_federation_default_oidc_audiences Read the token_federation_default_oidc_audiences list from the well-known discovery endpoint and use the first element as the token_audience when not explicitly configured. Falls back to account_id for account-level hosts without workspace_id. Co-authored-by: Isaac --- NEXT_CHANGELOG.md | 1 + databricks/sdk/config.py | 23 +++++++----- databricks/sdk/oauth.py | 3 ++ tests/test_config.py | 77 ++++++++++++++++++++++++++++++++++++++-- 4 files changed, 94 insertions(+), 10 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 7b7310b02..1b9a3ae70 100755 --- a/NEXT_CHANGELOG.md +++ b/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. diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index 599dc0efe..ee7641f6c 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -14,13 +14,15 @@ from ._base_client import _fix_host_if_needed from .client_types import ClientType, HostType from .clock import Clock, RealClock -from .credentials_provider import (CredentialsStrategy, DefaultCredentials, - OAuthCredentialsProvider) -from .environments import (ALL_ENVS, AzureEnvironment, Cloud, - DatabricksEnvironment, get_environment_for_hostname) -from .oauth import (OidcEndpoints, Token, - get_azure_entra_id_workspace_endpoints, - get_endpoints_from_url, get_host_metadata) +from .credentials_provider import CredentialsStrategy, DefaultCredentials, OAuthCredentialsProvider +from .environments import ALL_ENVS, AzureEnvironment, Cloud, DatabricksEnvironment, get_environment_for_hostname +from .oauth import ( + OidcEndpoints, + Token, + get_azure_entra_id_workspace_endpoints, + get_endpoints_from_url, + get_host_metadata, +) logger = logging.getLogger("databricks.sdk") @@ -655,10 +657,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 diff --git a/databricks/sdk/oauth.py b/databricks/sdk/oauth.py index 9f3656fda..e5d034ae2 100644 --- a/databricks/sdk/oauth.py +++ b/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, } diff --git a/tests/test_config.py b/tests/test_config.py index 0130fab88..64f4caac8 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -10,8 +10,7 @@ import pytest from databricks.sdk import AccountClient, WorkspaceClient, oauth, useragent -from databricks.sdk.config import (ClientType, Config, HostType, with_product, - with_user_agent_extra) +from databricks.sdk.config import ClientType, Config, HostType, with_product, with_user_agent_extra from databricks.sdk.environments import Cloud from databricks.sdk.oauth import HostMetadata from databricks.sdk.version import __version__ @@ -973,3 +972,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 From 998e2636b36eb31f8a453586f6fd9adab3ac935c Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Tue, 14 Apr 2026 12:13:25 +0200 Subject: [PATCH 2/2] Fix import formatting to match project isort/yapf style Revert black-style imports back to the project's continuation-line indent style used by yapf/isort. Co-authored-by: Isaac --- databricks/sdk/config.py | 16 +++++++--------- tests/test_config.py | 3 ++- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index ee7641f6c..0a6b8a56a 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -14,15 +14,13 @@ from ._base_client import _fix_host_if_needed from .client_types import ClientType, HostType from .clock import Clock, RealClock -from .credentials_provider import CredentialsStrategy, DefaultCredentials, OAuthCredentialsProvider -from .environments import ALL_ENVS, AzureEnvironment, Cloud, DatabricksEnvironment, get_environment_for_hostname -from .oauth import ( - OidcEndpoints, - Token, - get_azure_entra_id_workspace_endpoints, - get_endpoints_from_url, - get_host_metadata, -) +from .credentials_provider import (CredentialsStrategy, DefaultCredentials, + OAuthCredentialsProvider) +from .environments import (ALL_ENVS, AzureEnvironment, Cloud, + DatabricksEnvironment, get_environment_for_hostname) +from .oauth import (OidcEndpoints, Token, + get_azure_entra_id_workspace_endpoints, + get_endpoints_from_url, get_host_metadata) logger = logging.getLogger("databricks.sdk") diff --git a/tests/test_config.py b/tests/test_config.py index 64f4caac8..db872c8a7 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -10,7 +10,8 @@ import pytest from databricks.sdk import AccountClient, WorkspaceClient, oauth, useragent -from databricks.sdk.config import ClientType, Config, HostType, with_product, with_user_agent_extra +from databricks.sdk.config import (ClientType, Config, HostType, with_product, + with_user_agent_extra) from databricks.sdk.environments import Cloud from databricks.sdk.oauth import HostMetadata from databricks.sdk.version import __version__