Skip to content

Resolve token_audience from host metadata for account hosts #1321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hectorcast-db merged 1 commit into from
Mar 12, 2026
Merged

Resolve token_audience from host metadata for account hosts #1321

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion databricks/sdk/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ class Config:

# Audience for OIDC ID token source accepting an audience as a parameter.
# For example, the GitHub action ID token source.
token_audience: str = ConfigAttribute(env="DATABRICKS_TOKEN_AUDIENCE", auth="github-oidc")
Copy link
Copy Markdown
Contributor Author

@hectorcast-db hectorcast-db Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not really specific to github-oidc. But having the auth annotation makes it that the validation fails if token_audience is set for any other auth type.

token_audience: str = ConfigAttribute(env="DATABRICKS_TOKEN_AUDIENCE")

# Environment variable for OIDC token.
oidc_token_env: str = ConfigAttribute(env="DATABRICKS_OIDC_TOKEN_ENV", auth="env-oidc")
Expand Down Expand Up @@ -675,6 +675,13 @@ 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
# 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

def _fix_host_if_needed(self):
updated_host = _fix_host_if_needed(self.host)
Expand Down
62 changes: 62 additions & 0 deletions tests/test_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1011,3 +1011,65 @@ def test_resolve_host_metadata_cloud_missing_in_response(mocker):
)
config = Config(host=_DUMMY_WS_HOST, token="t", experimental_is_unified_host=True)
assert config.cloud is None


# ---------------------------------------------------------------------------
# token_audience resolution from host metadata
# ---------------------------------------------------------------------------


def test_resolve_host_metadata_sets_token_audience_for_account_host(mocker):
"""When metadata has no workspace_id, token_audience is set to account_id."""
mocker.patch(
"databricks.sdk.config.get_host_metadata",
return_value=oauth.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,
experimental_is_unified_host=True,
)
assert config.token_audience == _DUMMY_ACCOUNT_ID


def test_resolve_host_metadata_no_token_audience_for_workspace_host(mocker):
"""When metadata contains workspace_id, token_audience is not set from metadata."""
mocker.patch(
"databricks.sdk.config.get_host_metadata",
return_value=oauth.HostMetadata.from_dict(
{
"oidc_endpoint": f"{_DUMMY_WS_HOST}/oidc",
"account_id": _DUMMY_ACCOUNT_ID,
"workspace_id": _DUMMY_WORKSPACE_ID,
}
),
)
config = Config(host=_DUMMY_WS_HOST, token="t", experimental_is_unified_host=True)
assert config.token_audience is None


def test_resolve_host_metadata_does_not_overwrite_token_audience(mocker):
"""An explicitly set token_audience is never overwritten by metadata resolution."""
mocker.patch(
"databricks.sdk.config.get_host_metadata",
return_value=oauth.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,
experimental_is_unified_host=True,
token_audience="custom-audience",
)
assert config.token_audience == "custom-audience"
Loading