Discovery-driven login with workspace selection for SPOG hosts#4809
Merged
andrewnester merged 23 commits intomainfrom Mar 26, 2026
Merged
Discovery-driven login with workspace selection for SPOG hosts#4809andrewnester merged 23 commits intomainfrom
andrewnester merged 23 commits intomainfrom
Conversation
Call .well-known/databricks-config during login to detect SPOG hosts and auto-populate account_id/workspace_id. Refactor ToOAuthArgument() to route based on DiscoveryURL from EnsureResolved() instead of the Experimental_IsUnifiedHost flag. Key changes: - Parse ?o= and ?a= query params from host URLs in auth login - Run host metadata discovery via EnsureResolved() in setHostAndAccountId() - Refactor ToOAuthArgument() to use DiscoveryURL for OAuth flow routing - Simplify profile matching to use WorkspaceID/AccountID presence - Add post-auth workspace selection for multi-workspace SPOG accounts - Remove --experimental-is-unified-host from generated login commands - Keep backward compat with existing profiles that have IsUnifiedHost
- URL query params (?o=) now override stale profile workspace_id values by deferring profile inheritance to after URL parsing - ToOAuthArgument() checks /oidc/accounts/ in DiscoveryURL and uses caller-provided AccountID to prevent env var contamination and classic workspace misrouting - Token profile matching uses host+account_id+workspace_id for SPOG workspace disambiguation - Re-add --experimental-is-unified-host to BuildLoginCommand for legacy profile compatibility - Cap workspace selection prompt at 50 entries - Strip trailing slash in extractHostQueryParams after removing query params - Add tests for discovery routing, URL param precedence, and profile matching Co-authored-by: Isaac
Collaborator
|
Commit: e7e316c
19 interesting tests: 10 SKIP, 7 KNOWN, 2 flaky
Top 23 slowest tests (at least 2 minutes):
|
Remove context.Background() usage in ToOAuthArgument by silently discarding resolution errors (discovery failure is expected for non-SPOG hosts). Replace require.NoError in HTTP handlers with http.Error to avoid cross-goroutine FailNow calls. Fix gofmt alignment in login_test.go. Co-authored-by: Isaac
Co-authored-by: Isaac
1. Replace IsAccountsHost with SDK's HostType() to avoid divergence risk for new sovereign clouds. Remove IsAccountsHost and normalizeHost helpers. 2. Validate ?o= and ?workspace_id= query params are numeric, skipping non-numeric values to avoid confusing downstream errors. 3. Tighten writeReauthSteps condition to only append --workspace-id for SPOG/unified hosts (both account_id and workspace_id present). 4. Cache discovery results via DiscoveryURL field on AuthArguments to avoid duplicate .well-known/databricks-config HTTP requests during login. Co-authored-by: Isaac
…ount admins SPOG account admins who don't need a workspace were prompted on every login. This adds two ways to persist that choice: a --skip-workspace flag on `auth login`, and persisting workspace_id=none when the user selects "Skip" from the interactive prompt. On subsequent logins the sentinel is recognized and the prompt is suppressed. Co-authored-by: Isaac
Suggested reviewersBased on git history of the changed files, these people are best suited to review:
Confidence: high Eligible reviewersBased on CODEOWNERS, these people or teams could also review: @andrewnester, @anton-107, @denik, @pietern, @shreyas-goenka Suggestions based on git history of 21 changed files (9 scored). See CODEOWNERS for path-specific ownership rules. |
Contributor
hectorcast-db
left a comment
There was a problem hiding this comment.
Added some comments.
Code looks good, but I am missing some context on CLI to fully understand the implications of this change. Please get also a review from DABs team.
Replace HostType-based heuristics with discovery-aware detection: - login.go: Use IsUnifiedHost flag instead of HostType() for workspace selection prompt. HostType is unreliable for SPOG hosts. - arguments.go: Add comment explaining HostType is only used for classic accounts.* URLs where URL pattern matching is reliable. - error.go: Use Experimental_IsUnifiedHost instead of AccountID+WorkspaceID heuristic to detect SPOG hosts in re-auth commands. Legacy workspaces can also have both IDs set. - error_test.go: Add test for legacy workspace with both IDs to verify --workspace-id is not incorrectly appended. Co-authored-by: Isaac
Replace HostType() call with direct URL prefix check for classic accounts.* hosts. HostType() was only checking the URL prefix anyway, so a direct strings.HasPrefix is clearer and avoids the dependency. Co-authored-by: Isaac
andrewnester
approved these changes
Mar 24, 2026
Call auth.ExtractHostQueryParams directly at the call site instead of going through a private wrapper. Move unit tests to libs/auth where the shared function lives. Co-authored-by: Isaac
hectorcast-db
approved these changes
Mar 25, 2026
The post-OAuth workspace selection was gated on IsUnifiedHost, which is only set via --experimental-is-unified-host or from an existing profile. When a user passes --host pointing to a SPOG-enabled workspace, discovery fetches account_id from .well-known/databricks-config but never sets IsUnifiedHost, so workspace selection was silently skipped. Now the condition checks only for account_id + missing workspace_id, regardless of IsUnifiedHost. Also adds a message in non-interactive mode telling the user to set workspace_id. Co-authored-by: Isaac
The UnifiedHost case in setHostAndAccountId prompted for workspace_id via a text input before OAuth. This is redundant now that workspace selection happens post-OAuth via the workspace list selector, and it created a split path where unified hosts got a worse UX than discovery. Co-authored-by: Isaac
…est.go Co-authored-by: Isaac
github-merge-queue bot
pushed a commit
that referenced
this pull request
Mar 26, 2026
## Why The discovery login flow via `login.databricks.com` (PR #4702) is a separate code path from the regular `--host` login. After #4809 added SPOG host detection via `.well-known/databricks-config` to the regular login path, the discovery flow was missing this behavior. Profiles created via `login.databricks.com` for SPOG hosts had no `account_id` and no discovery metadata, which breaks re-authentication (because `ToOAuthArgument()` needs `account_id` to route to unified OAuth). ## Changes **Before:** `discoveryLogin()` only used token introspection for `workspace_id` and deliberately skipped saving `account_id`. **Now:** After getting the host from `login.databricks.com`, `discoveryLogin()` calls `runHostDiscovery()` on the discovered host to populate `account_id`, `workspace_id`, and `DiscoveryURL` from `.well-known/databricks-config`. Token introspection is kept as a fallback for hosts where discovery is unavailable (e.g. classic workspace hosts). `account_id` is now saved to the profile. **Note:** This is not the full SPOG story for discovery login. Most SPOG workspaces will need additional handling during the login.databricks.com flow (e.g. workspace selection after discovery detects a multi-workspace account). A follow-up PR will address this. ## Test plan - New test: discovery login with SPOG host (mock `.well-known/databricks-config`) verifies `account_id` and `workspace_id` come from discovery - New test: discovery login where host discovery fails verifies fallback to introspection - Updated existing tests to assert `account_id` is now saved - All `go test ./cmd/auth/... ./libs/auth/...` pass - `make checks` passes This pull request was AI-assisted by Isaac.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
SPOG (Single Pane of Glass) hosts use one URL for both account-level and workspace-level APIs. The CLI's login flow doesn't call the
.well-known/databricks-configdiscovery endpoint, so it can't detect SPOG hosts. This meansaccount_idandworkspace_iddon't get populated for SPOG profiles, and users of multi-workspace SPOG accounts have no way to setworkspace_idthrough the interactive login flow.Changes
The CLI now calls discovery during login to detect SPOG hosts and auto-populate fields. Users can also paste browser URLs with query params directly into
auth login.Before: SPOG hosts classified as regular workspace hosts. No account_id/workspace_id discovery. Users must manually edit
.databrickscfg.Now: Discovery runs automatically during login. SPOG fields are populated from
.well-known/databricks-config. Multi-workspace accounts get an interactive workspace picker.Concrete changes:
setHostAndAccountId(): extracts?o=(workspace_id) and?a=(account_id) from host URLs, strips query params from host. Explicit flags take precedence over URL params.EnsureResolved()with a temporary config during login to fetch.well-known/databricks-config. Populates account_id/workspace_id if not already set. Best-effort, never blocks login on failure.ToOAuthArgument()refactored: routes OAuth flow based onDiscoveryURLfromEnsureResolved()instead ofExperimental_IsUnifiedHostflag. Account-scoped OIDC endpoints (/oidc/accounts/) trigger unified OAuth. Classic workspace hosts are never misrouted.MatchWorkspaceProfiles()usesWorkspaceID != ""presence.MatchAccountProfiles()usesAccountID != "" && WorkspaceID == "". Both still handle legacyIsUnifiedHostprofiles.Workspaces.List(). Auto-selects single workspace. Capped at 50 entries. Skippable. Non-interactive mode proceeds without selection.loadToken()uses workspace_id for disambiguation when matching multi-workspace SPOG profiles.Note on discovery in
ToOAuthArgument(): This function now callsEnsureResolved()on every invocation, which makes an unauthenticated HTTP request to{host}/.well-known/databricks-config. This is a deliberate tradeoff. The alternative was to persist additional state (likediscovery_urlorexperimental_is_unified_host) to.databrickscfgso we could route the OAuth flow without a network call. We chose not to do that because we want to move away from persisting host-type markers to profiles and instead derive behavior from the host at runtime. The added latency (~100ms per call) will be addressed by a separate discovery caching PR that caches.well-known/databricks-configresponses with a 1-hour TTL, making repeated calls essentially free.Backward compatibility:
experimental_is_unified_host = truestill work (legacy fallback inToOAuthArgument())accounts.*login flow unchanged--experimental-is-unified-hostflag preserved in re-auth error messagesTest plan
ToOAuthArgument()routing tests (SPOG -> unified, classic workspace not misrouted, env var contamination prevented, legacy IsUnifiedHost fallback)IsAccountsHost()helper testsmake checkspasses