Skip to content

Commit fdf335d

Browse files
Remove unified flag usage and fix tests for server-side OIDC resolution (#1331)
## 🥞 Stacked PR Use this [link](https://github.com/databricks/databricks-sdk-py/pull/1331/files) to review incremental changes. - [**stack/unified-for-real**](#1331) [[Files changed](https://github.com/databricks/databricks-sdk-py/pull/1331/files)] --------- ## Summary - Remove `experimental_is_unified_host` flag and `ClientType.UNIFIED` usages — all hosts now go through the same metadata resolution path - Simplify `databricks_oidc_endpoints` to use `discovery_url` from host metadata instead of client-side URL construction - Add per-test `HostMetadata` overrides for tests that need a valid `discovery_url` - Remove 5 redundant OIDC endpoint tests that only round-tripped mocked metadata - Remove 3 tests for deleted unified flag behavior - Update remaining tests to match new behavior (metadata always resolved, `is_account_client` no longer raises on unified hosts) NOTE: This PR does not remove yet the flag itself, since the feature is still on experimental mode. ## Test plan - [x] All 165 remaining tests pass - [x] Verified removed tests were redundant (covered by `test_databricks_oidc_endpoints_uses_discovery_url`) This pull request was AI-assisted by Isaac.
1 parent a8e6ef1 commit fdf335d

39 files changed

Lines changed: 945 additions & 1194 deletions

.github/workflows/tagging.yml

100755100644
File mode changed.

NEXT_CHANGELOG.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,4 +16,4 @@
1616
* Add `databricks.sdk.service.environments` package.
1717
* Add [w.environments](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/environments/environments.html) workspace-level service.
1818
* Add `parent_path` field for `databricks.sdk.service.dashboards.GenieSpace`.
19-
* Add `can_create_app` enum value for `databricks.sdk.service.iam.PermissionLevel`.
19+
* Add `can_create_app` enum value for `databricks.sdk.service.iam.PermissionLevel`.

databricks/sdk/config.py

Lines changed: 21 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,9 @@
1919
OAuthCredentialsProvider)
2020
from .environments import (ALL_ENVS, AzureEnvironment, Cloud,
2121
DatabricksEnvironment, get_environment_for_hostname)
22-
from .oauth import (OidcEndpoints, Token, get_account_endpoints,
22+
from .oauth import (OidcEndpoints, Token,
2323
get_azure_entra_id_workspace_endpoints,
24-
get_endpoints_from_url, get_host_metadata,
25-
get_unified_endpoints, get_workspace_endpoints)
24+
get_endpoints_from_url, get_host_metadata)
2625

2726
logger = logging.getLogger("databricks.sdk")
2827

@@ -415,14 +414,14 @@ def is_aws(self) -> bool:
415414

416415
@property
417416
def host_type(self) -> HostType:
418-
"""Determine the type of host based on the configuration.
419-
420-
Returns the HostType which can be ACCOUNTS, WORKSPACE, or UNIFIED.
421417
"""
422-
# Check if explicitly marked as unified host
423-
if self.experimental_is_unified_host:
424-
return HostType.UNIFIED
418+
[DEPRECATED]
419+
Host type and client type are deprecated. Some hosts can now support both workspace and account APIs.
420+
This method returns the HostType based on the host pattern, which is not accurate.
421+
For example, a unified host can support both workspace and account APIs, but WORKSPACE is returned.
425422
423+
This method still returns the correct value for legacy hosts which only support either workspace or account APIs.
424+
"""
426425
if not self.host:
427426
return HostType.WORKSPACE
428427

@@ -434,15 +433,13 @@ def host_type(self) -> HostType:
434433

435434
@property
436435
def client_type(self) -> ClientType:
437-
"""Determine the type of client configuration.
438-
439-
This is separate from host_type. For example, a unified host can support both
440-
workspace and account client types.
441-
442-
Returns ClientType.ACCOUNT or ClientType.WORKSPACE based on the configuration.
436+
"""
437+
[DEPRECATED]
438+
Host type and client type are deprecated. Some hosts can now support both workspace and account APIs.
439+
This method returns the ClientType based on the host pattern, which is not accurate.
440+
For example, a unified host can support both workspace and account APIs, but WORKSPACE is returned.
443441
444-
For unified hosts, account_id must be set. If workspace_id is also set,
445-
returns WORKSPACE, otherwise returns ACCOUNT.
442+
This method still returns the correct value for legacy hosts which only support either workspace or account APIs.
446443
"""
447444
host_type = self.host_type
448445

@@ -452,25 +449,18 @@ def client_type(self) -> ClientType:
452449
if host_type == HostType.WORKSPACE:
453450
return ClientType.WORKSPACE
454451

455-
if host_type == HostType.UNIFIED:
456-
if not self.account_id:
457-
raise ValueError("Unified host requires account_id to be set")
458-
if self.workspace_id:
459-
return ClientType.WORKSPACE
460-
return ClientType.ACCOUNT
461-
462452
# Default to workspace for backward compatibility
463453
return ClientType.WORKSPACE
464454

465455
@property
466456
def is_account_client(self) -> bool:
467457
"""[Deprecated]
468-
Host type and client type are deprecated. Clients can now support both workspace and account APIs.
458+
Host type and client type are deprecated. Some hosts can now support both workspace and account APIs.
459+
This method returns True if the host is an accounts host, which is not accurate.
460+
For example, a unified host can support both workspace and account APIs, but False is returned.
461+
462+
This method still returns the correct value for legacy hosts which only support either workspace or account APIs.
469463
"""
470-
if self.experimental_is_unified_host:
471-
raise ValueError(
472-
"is_account_client cannot be used with unified hosts; use host_type or client_type instead"
473-
)
474464
if not self.host:
475465
return False
476466
return self.host.startswith("https://accounts.") or self.host.startswith("https://accounts-dod.")
@@ -535,21 +525,7 @@ def databricks_oidc_endpoints(self) -> Optional[OidcEndpoints]:
535525
if not self.host:
536526
return None
537527

538-
if self.discovery_url:
539-
return get_endpoints_from_url(self.discovery_url)
540-
541-
# Handle unified hosts
542-
if self.host_type == HostType.UNIFIED:
543-
if not self.account_id:
544-
raise ValueError("Unified host requires account_id to be set for OAuth endpoints")
545-
return get_unified_endpoints(self.host, self.account_id)
546-
547-
# Handle traditional account hosts
548-
if self.host_type == HostType.ACCOUNTS and self.account_id:
549-
return get_account_endpoints(self.host, self.account_id)
550-
551-
# Default to workspace endpoints
552-
return get_workspace_endpoints(self.host)
528+
return get_endpoints_from_url(self.discovery_url)
553529

554530
@property
555531
def oidc_endpoints(self) -> Optional[OidcEndpoints]:
@@ -647,15 +623,12 @@ def attributes(cls) -> Iterable[ConfigAttribute]:
647623
return cls._attributes
648624

649625
def _resolve_host_metadata(self) -> None:
650-
"""[Experimental] Populate missing config fields from the host's
626+
"""Populate missing config fields from the host's
651627
/.well-known/databricks-config discovery endpoint.
652628
653629
Fills in account_id, workspace_id, and discovery_url (derived from oidc_endpoint,
654630
with any {account_id} placeholder substituted) if not already set.
655631
"""
656-
# TODO: Enable this everywhere
657-
if not self.host_type == HostType.UNIFIED:
658-
return
659632
if not self.host:
660633
return
661634
try:

databricks/sdk/credentials_provider.py

Lines changed: 4 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -440,9 +440,7 @@ def _oidc_credentials_provider(
440440

441441
# Determine the audience for token exchange
442442
audience = cfg.token_audience
443-
if audience is None and cfg.client_type == ClientType.ACCOUNT:
444-
audience = cfg.account_id
445-
if audience is None and cfg.client_type != ClientType.ACCOUNT:
443+
if audience is None:
446444
audience = cfg.databricks_oidc_endpoints.token_endpoint
447445

448446
# Try to get an OIDC token. If no supplier returns a token, we cannot use this authentication mode.
@@ -986,14 +984,9 @@ def _validate_token_scopes(self, token: oauth.Token):
986984
def _build_host_args(cfg: "Config") -> List[str]:
987985
"""Build CLI arguments using --host (legacy path)."""
988986
args = ["auth", "token", "--host", cfg.host]
989-
if cfg.experimental_is_unified_host:
990-
# For unified hosts, pass account_id, workspace_id, and experimental flag
991-
args += ["--experimental-is-unified-host"]
992-
if cfg.account_id:
993-
args += ["--account-id", cfg.account_id]
994-
if cfg.workspace_id:
995-
args += ["--workspace-id", str(cfg.workspace_id)]
996-
elif cfg.client_type == ClientType.ACCOUNT:
987+
# This is here to support older versions of the Databricks CLI, so we need to keep the client type check.
988+
# This won't work for unified hosts, but it is not supposed to.
989+
if cfg.client_type == ClientType.ACCOUNT:
997990
args += ["--account-id", cfg.account_id]
998991
return args
999992

databricks/sdk/service/agentbricks.py

Lines changed: 6 additions & 7 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)