From f48e10283426bda51b2ef18c218ec019d51de3c9 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Mon, 30 Mar 2026 16:54:57 +0200 Subject: [PATCH 1/7] Add HostType to HostMetadata Signed-off-by: Tanmay Rustagi --- databricks/sdk/client_types.py | 19 ++++++ databricks/sdk/config.py | 6 ++ databricks/sdk/oauth.py | 3 + tests/test_config.py | 116 +++++++++++++++++++++++++++++++++ 4 files changed, 144 insertions(+) diff --git a/databricks/sdk/client_types.py b/databricks/sdk/client_types.py index 3937cdde0..729a454af 100644 --- a/databricks/sdk/client_types.py +++ b/databricks/sdk/client_types.py @@ -1,4 +1,5 @@ from enum import Enum +from typing import Optional class HostType(Enum): @@ -8,6 +9,24 @@ class HostType(Enum): WORKSPACE = "workspace" UNIFIED = "unified" + @staticmethod + def from_api_value(value: str) -> Optional["HostType"]: + """Normalize a host_type string from the API to a HostType enum value. + + Maps "workspace" -> WORKSPACE, "account" -> ACCOUNTS, "unified" -> UNIFIED. + Returns None for unrecognized or empty values. + """ + if not value: + return None + normalized = value.lower() + if normalized == "workspace": + return HostType.WORKSPACE + if normalized == "account": + return HostType.ACCOUNTS + if normalized == "unified": + return HostType.UNIFIED + return None + class ClientType(Enum): """Enum representing the type of client configuration.""" diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index fc0f722b8..473ef38e0 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -294,6 +294,7 @@ def __init__( self._header_factory = None self._inner = {} self._user_agent_other_info = [] + self._resolved_host_type = None self._custom_headers = custom_headers or {} if credentials_strategy and credentials_provider: raise ValueError("When providing `credentials_strategy` field, `credential_provider` cannot be specified.") @@ -665,6 +666,11 @@ 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 self._resolved_host_type is None and meta.host_type: + resolved = HostType.from_api_value(meta.host_type) + if resolved is not None: + logger.debug(f"Resolved host_type from host metadata: {meta.host_type}") + self._resolved_host_type = resolved # 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. diff --git a/databricks/sdk/oauth.py b/databricks/sdk/oauth.py index 9f3656fda..85ac05b2a 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 + host_type: Optional[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", "")), + host_type=d.get("host_type"), ) 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, + "host_type": self.host_type, } diff --git a/tests/test_config.py b/tests/test_config.py index 14ea4ca69..7d536e89a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -981,3 +981,119 @@ def test_resolve_host_metadata_does_not_overwrite_token_audience(mocker): token_audience="custom-audience", ) assert config.token_audience == "custom-audience" + + +# --------------------------------------------------------------------------- +# HostType.from_api_value tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "api_value,expected", + [ + ("workspace", HostType.WORKSPACE), + ("Workspace", HostType.WORKSPACE), + ("WORKSPACE", HostType.WORKSPACE), + ("account", HostType.ACCOUNTS), + ("Account", HostType.ACCOUNTS), + ("ACCOUNT", HostType.ACCOUNTS), + ("unified", HostType.UNIFIED), + ("Unified", HostType.UNIFIED), + ("UNIFIED", HostType.UNIFIED), + ("unknown", None), + ("", None), + (None, None), + ], +) +def test_host_type_from_api_value(api_value, expected): + assert HostType.from_api_value(api_value) == expected + + +# --------------------------------------------------------------------------- +# HostMetadata.from_dict with host_type field +# --------------------------------------------------------------------------- + + +def test_host_metadata_from_dict_with_host_type(): + """HostMetadata.from_dict parses the host_type field.""" + d = { + "oidc_endpoint": "https://host/oidc", + "account_id": "acc-1", + "host_type": "workspace", + } + meta = HostMetadata.from_dict(d) + assert meta.host_type == "workspace" + + +def test_host_metadata_from_dict_without_host_type(): + """HostMetadata.from_dict returns None for missing host_type.""" + d = {"oidc_endpoint": "https://host/oidc"} + meta = HostMetadata.from_dict(d) + assert meta.host_type is None + + +# --------------------------------------------------------------------------- +# _resolve_host_metadata populates _resolved_host_type +# --------------------------------------------------------------------------- + + +def test_resolve_host_metadata_populates_resolved_host_type(mocker): + """_resolved_host_type is populated from metadata host_type.""" + mocker.patch( + "databricks.sdk.config.get_host_metadata", + return_value=HostMetadata.from_dict( + { + "oidc_endpoint": f"{_DUMMY_WS_HOST}/oidc", + "host_type": "unified", + } + ), + ) + config = Config(host=_DUMMY_WS_HOST, token="t") + assert config._resolved_host_type == HostType.UNIFIED + + +def test_resolve_host_metadata_does_not_overwrite_existing_resolved_host_type(mocker): + """An already-set _resolved_host_type is not overwritten by metadata.""" + mocker.patch( + "databricks.sdk.config.get_host_metadata", + return_value=HostMetadata.from_dict( + { + "oidc_endpoint": f"{_DUMMY_WS_HOST}/oidc", + "host_type": "account", + } + ), + ) + config = Config(host=_DUMMY_WS_HOST, token="t") + # Manually set resolved host type then re-resolve + config._resolved_host_type = HostType.UNIFIED + config._resolve_host_metadata() + assert config._resolved_host_type == HostType.UNIFIED + + +def test_resolve_host_metadata_resolved_host_type_none_when_missing(mocker): + """_resolved_host_type stays None when metadata has no host_type.""" + mocker.patch( + "databricks.sdk.config.get_host_metadata", + return_value=HostMetadata.from_dict( + { + "oidc_endpoint": f"{_DUMMY_WS_HOST}/oidc", + } + ), + ) + config = Config(host=_DUMMY_WS_HOST, token="t") + assert config._resolved_host_type is None + + +def test_resolve_host_metadata_resolved_host_type_unknown_string(mocker): + """_resolved_host_type stays None for unrecognized host_type strings.""" + mocker.patch( + "databricks.sdk.config.get_host_metadata", + return_value=HostMetadata.from_dict( + { + "oidc_endpoint": f"{_DUMMY_WS_HOST}/oidc", + "host_type": "some_future_type", + } + ), + ) + config = Config(host=_DUMMY_WS_HOST, token="t") + assert config._resolved_host_type is None From 142c1695aff1c945c8abdbf7d7369c7065ba34f0 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Mon, 30 Mar 2026 18:39:31 +0200 Subject: [PATCH 2/7] Update NEXT_CHANGELOG.md with host_type metadata entry Signed-off-by: Tanmay Rustagi --- NEXT_CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 0c81bdfbf..5a3fb4bf5 100755 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -14,6 +14,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. +* Add `host_type` to `HostMetadata` and `HostType.from_api_value()` for normalizing host type strings from the discovery endpoint. ### 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. From 493fe3855812e16dcecf521b3dc52d896845cee4 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Mon, 30 Mar 2026 16:56:07 +0200 Subject: [PATCH 3/7] Resolve token_audience from default_oidc_audience in host metadata Signed-off-by: Tanmay Rustagi --- databricks/sdk/config.py | 4 ++- databricks/sdk/oauth.py | 3 ++ tests/test_config.py | 72 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 78 insertions(+), 1 deletion(-) diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index 473ef38e0..b581e2c3d 100644 --- a/databricks/sdk/config.py +++ b/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 diff --git a/databricks/sdk/oauth.py b/databricks/sdk/oauth.py index 85ac05b2a..b1ede5381 100644 --- a/databricks/sdk/oauth.py +++ b/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, } diff --git a/tests/test_config.py b/tests/test_config.py index 7d536e89a..2330dffde 100644 --- a/tests/test_config.py +++ b/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 From 5b452430782ebd5ad6f016e741172a1da11e4cb7 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Mon, 30 Mar 2026 16:57:36 +0200 Subject: [PATCH 4/7] Use resolved host type from host metadata in host_type property Signed-off-by: Tanmay Rustagi --- databricks/sdk/config.py | 12 ++--- tests/test_config.py | 100 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 106 insertions(+), 6 deletions(-) diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index b581e2c3d..025da1b93 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -421,14 +421,14 @@ def is_aws(self) -> bool: @property def host_type(self) -> HostType: - """ - [DEPRECATED] - Host type and client type are deprecated. Some hosts can now support both workspace and account APIs. - This method returns the HostType based on the host pattern, which is not accurate. - For example, a unified host can support both workspace and account APIs, but WORKSPACE is returned. + """Returns the type of host that the client is configured for. - This method still returns the correct value for legacy hosts which only support either workspace or account APIs. + When available, uses the host type resolved from the /.well-known/databricks-config + discovery endpoint. Falls back to URL-based pattern matching when metadata is unavailable. """ + if self._resolved_host_type is not None: + return self._resolved_host_type + if not self.host: return HostType.WORKSPACE diff --git a/tests/test_config.py b/tests/test_config.py index 2330dffde..2ea69170f 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1169,3 +1169,103 @@ def test_resolve_host_metadata_falls_back_to_account_id_when_no_default_oidc_aud 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 + + +# --------------------------------------------------------------------------- +# host_type property uses resolved host type from metadata +# --------------------------------------------------------------------------- + + +def test_host_type_uses_metadata_workspace(mocker): + """Metadata host_type=workspace overrides an account-like URL pattern.""" + acc_host = "https://accounts.cloud.databricks.com" + mocker.patch( + "databricks.sdk.config.get_host_metadata", + return_value=HostMetadata.from_dict( + { + "oidc_endpoint": f"{acc_host}/oidc/accounts/{_DUMMY_ACCOUNT_ID}", + "account_id": _DUMMY_ACCOUNT_ID, + "host_type": "workspace", + } + ), + ) + config = Config(host=acc_host, token="t", account_id=_DUMMY_ACCOUNT_ID) + assert config.host_type == HostType.WORKSPACE + + +def test_host_type_uses_metadata_account(mocker): + """Metadata host_type=account overrides a workspace-like URL pattern.""" + mocker.patch( + "databricks.sdk.config.get_host_metadata", + return_value=HostMetadata.from_dict( + { + "oidc_endpoint": f"{_DUMMY_WS_HOST}/oidc", + "host_type": "account", + } + ), + ) + config = Config(host=_DUMMY_WS_HOST, token="t") + assert config.host_type == HostType.ACCOUNTS + + +def test_host_type_uses_metadata_unified(mocker): + """Metadata host_type=unified is returned by host_type property.""" + mocker.patch( + "databricks.sdk.config.get_host_metadata", + return_value=HostMetadata.from_dict( + { + "oidc_endpoint": f"{_DUMMY_WS_HOST}/oidc", + "host_type": "unified", + } + ), + ) + config = Config(host=_DUMMY_WS_HOST, token="t") + assert config.host_type == HostType.UNIFIED + + +def test_host_type_falls_back_to_url_when_metadata_unknown(mocker): + """Falls back to URL-based matching when metadata host_type is unrecognized.""" + acc_host = "https://accounts.cloud.databricks.com" + mocker.patch( + "databricks.sdk.config.get_host_metadata", + return_value=HostMetadata.from_dict( + { + "oidc_endpoint": f"{acc_host}/oidc/accounts/{_DUMMY_ACCOUNT_ID}", + "account_id": _DUMMY_ACCOUNT_ID, + } + ), + ) + config = Config(host=acc_host, token="t", account_id=_DUMMY_ACCOUNT_ID) + assert config.host_type == HostType.ACCOUNTS + + +def test_host_type_falls_back_to_url_workspace_when_metadata_unknown(mocker): + """Falls back to URL-based WORKSPACE when metadata has no host_type.""" + mocker.patch( + "databricks.sdk.config.get_host_metadata", + return_value=HostMetadata.from_dict( + { + "oidc_endpoint": f"{_DUMMY_WS_HOST}/oidc", + } + ), + ) + config = Config(host=_DUMMY_WS_HOST, token="t") + assert config.host_type == HostType.WORKSPACE + + +def test_host_type_end_to_end_metadata_resolves(mocker): + """End-to-end: resolve → host_type returns metadata value over URL matching.""" + 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, + "host_type": "account", + } + ), + ) + # Workspace-like URL but account host_type from metadata + config = Config(host=_DUMMY_WS_HOST, token="t") + assert config.host_type == HostType.ACCOUNTS + assert config._resolved_host_type == HostType.ACCOUNTS From 0406f86b23beade8e659ba6af572594d520b5527 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Mon, 13 Apr 2026 11:46:48 +0200 Subject: [PATCH 5/7] Rename default_oidc_audience to token_federation_default_oidc_audiences Align the field name with the API contract for the well-known databricks-config endpoint. Co-authored-by: Isaac --- databricks/sdk/config.py | 6 +++--- databricks/sdk/oauth.py | 6 +++--- tests/test_config.py | 28 ++++++++++++++-------------- 3 files changed, 20 insertions(+), 20 deletions(-) diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index 6eec5258a..933bc589a 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -661,9 +661,9 @@ 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 + if not self.token_audience and meta.token_federation_default_oidc_audiences: + logger.debug(f"Resolved token_audience from host metadata token_federation_default_oidc_audiences: {meta.token_federation_default_oidc_audiences}") + self.token_audience = meta.token_federation_default_oidc_audiences # 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. diff --git a/databricks/sdk/oauth.py b/databricks/sdk/oauth.py index b1ede5381..b399f454f 100644 --- a/databricks/sdk/oauth.py +++ b/databricks/sdk/oauth.py @@ -449,7 +449,7 @@ class HostMetadata: workspace_id: Optional[str] = None cloud: Optional[Cloud] = None host_type: Optional[str] = None - default_oidc_audience: Optional[str] = None + token_federation_default_oidc_audiences: Optional[str] = None @staticmethod def from_dict(d: dict) -> "HostMetadata": @@ -459,7 +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"), + token_federation_default_oidc_audiences=d.get("token_federation_default_oidc_audiences"), ) def as_dict(self) -> dict: @@ -469,7 +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, + "token_federation_default_oidc_audiences": self.token_federation_default_oidc_audiences, } diff --git a/tests/test_config.py b/tests/test_config.py index f17b8fd41..2925879d8 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1092,12 +1092,12 @@ def test_resolve_host_metadata_resolved_host_type_unknown_string(mocker): # --------------------------------------------------------------------------- -# default_oidc_audience resolution from host metadata +# token_federation_default_oidc_audiences 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.""" +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( @@ -1105,7 +1105,7 @@ def test_resolve_host_metadata_sets_token_audience_from_default_oidc_audience(mo "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", + "token_federation_default_oidc_audiences": f"{_DUMMY_WS_HOST}/oidc/v1/token", } ), ) @@ -1113,25 +1113,25 @@ def test_resolve_host_metadata_sets_token_audience_from_default_oidc_audience(mo 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.""" +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, - "default_oidc_audience": "custom-audience-from-server", + "token_federation_default_oidc_audiences": "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 + # 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_default_oidc_audience_does_not_override_existing_token_audience(mocker): - """An explicitly set token_audience is not overwritten by default_oidc_audience.""" +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( @@ -1139,7 +1139,7 @@ def test_resolve_host_metadata_default_oidc_audience_does_not_override_existing_ "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", + "token_federation_default_oidc_audiences": f"{_DUMMY_WS_HOST}/oidc/v1/token", } ), ) @@ -1147,8 +1147,8 @@ def test_resolve_host_metadata_default_oidc_audience_does_not_override_existing_ 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.""" +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( @@ -1159,5 +1159,5 @@ def test_resolve_host_metadata_falls_back_to_account_id_when_no_default_oidc_aud ), ) 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 + # No token_federation_default_oidc_audiences and no workspace_id → falls back to account_id assert config.token_audience == _DUMMY_ACCOUNT_ID From 9d7973e630ab099d3762be2423cb95082083c73e Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Mon, 13 Apr 2026 15:02:55 +0200 Subject: [PATCH 6/7] Change token_federation_default_oidc_audiences from str to list type The API returns this field as an array of strings. Use the first element when resolving token_audience. Co-authored-by: Isaac --- databricks/sdk/config.py | 5 +++-- databricks/sdk/oauth.py | 2 +- tests/test_config.py | 6 +++--- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index 933bc589a..65f45ef42 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -662,8 +662,9 @@ def _resolve_host_metadata(self) -> 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.token_federation_default_oidc_audiences: - logger.debug(f"Resolved token_audience from host metadata token_federation_default_oidc_audiences: {meta.token_federation_default_oidc_audiences}") - self.token_audience = 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. diff --git a/databricks/sdk/oauth.py b/databricks/sdk/oauth.py index b399f454f..18856b730 100644 --- a/databricks/sdk/oauth.py +++ b/databricks/sdk/oauth.py @@ -449,7 +449,7 @@ class HostMetadata: workspace_id: Optional[str] = None cloud: Optional[Cloud] = None host_type: Optional[str] = None - token_federation_default_oidc_audiences: Optional[str] = None + token_federation_default_oidc_audiences: Optional[List[str]] = None @staticmethod def from_dict(d: dict) -> "HostMetadata": diff --git a/tests/test_config.py b/tests/test_config.py index 2925879d8..d77738fbd 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1105,7 +1105,7 @@ def test_resolve_host_metadata_sets_token_audience_from_token_federation_default "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", + "token_federation_default_oidc_audiences": [f"{_DUMMY_WS_HOST}/oidc/v1/token"], } ), ) @@ -1121,7 +1121,7 @@ def test_resolve_host_metadata_token_federation_default_oidc_audiences_takes_pri { "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", + "token_federation_default_oidc_audiences": ["custom-audience-from-server"], } ), ) @@ -1139,7 +1139,7 @@ def test_resolve_host_metadata_token_federation_default_oidc_audiences_does_not_ "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", + "token_federation_default_oidc_audiences": [f"{_DUMMY_WS_HOST}/oidc/v1/token"], } ), ) From 9fd8f3723bdfd9b86fc3d8c0957c236e06965bc7 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Mon, 13 Apr 2026 15:18:13 +0200 Subject: [PATCH 7/7] Run formatters (black line length) on config.py and test_config.py Co-authored-by: Isaac --- databricks/sdk/config.py | 4 +++- tests/test_config.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index 65f45ef42..d5ca4568b 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -663,7 +663,9 @@ def _resolve_host_metadata(self) -> None: self._resolved_host_type = resolved 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}") + 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 diff --git a/tests/test_config.py b/tests/test_config.py index d77738fbd..e58817597 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1130,7 +1130,9 @@ def test_resolve_host_metadata_token_federation_default_oidc_audiences_takes_pri 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): +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",