Skip to content

Discovery-driven login with workspace selection for SPOG hosts#4809

Merged
andrewnester merged 23 commits intomainfrom
simonfaltum/spog-discovery-login
Mar 26, 2026
Merged

Discovery-driven login with workspace selection for SPOG hosts#4809
andrewnester merged 23 commits intomainfrom
simonfaltum/spog-discovery-login

Conversation

@simonfaltum
Copy link
Copy Markdown
Member

@simonfaltum simonfaltum commented Mar 22, 2026

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-config discovery endpoint, so it can't detect SPOG hosts. This means account_id and workspace_id don't get populated for SPOG profiles, and users of multi-workspace SPOG accounts have no way to set workspace_id through 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:

  • URL query param parsing in setHostAndAccountId(): extracts ?o= (workspace_id) and ?a= (account_id) from host URLs, strips query params from host. Explicit flags take precedence over URL params.
  • Host metadata discovery: calls 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 on DiscoveryURL from EnsureResolved() instead of Experimental_IsUnifiedHost flag. Account-scoped OIDC endpoints (/oidc/accounts/) trigger unified OAuth. Classic workspace hosts are never misrouted.
  • Profile matching simplified: MatchWorkspaceProfiles() uses WorkspaceID != "" presence. MatchAccountProfiles() uses AccountID != "" && WorkspaceID == "". Both still handle legacy IsUnifiedHost profiles.
  • Post-auth workspace selection: for SPOG hosts with account_id but no workspace_id, prompts user to select from Workspaces.List(). Auto-selects single workspace. Capped at 50 entries. Skippable. Non-interactive mode proceeds without selection.
  • Token path updated: loadToken() uses workspace_id for disambiguation when matching multi-workspace SPOG profiles.

Note on discovery in ToOAuthArgument(): This function now calls EnsureResolved() 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 (like discovery_url or experimental_is_unified_host) to .databrickscfg so 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-config responses with a 1-hour TTL, making repeated calls essentially free.

Backward compatibility:

  • Existing profiles with experimental_is_unified_host = true still work (legacy fallback in ToOAuthArgument())
  • Classic accounts.* login flow unchanged
  • Regular workspace login unchanged
  • --experimental-is-unified-host flag preserved in re-auth error messages

Test plan

  • Table-driven tests for URL query param parsing (all param combinations, explicit flag precedence, invalid URLs)
  • Discovery tests with mock HTTP server (SPOG host, classic workspace host, discovery failure)
  • ToOAuthArgument() routing tests (SPOG -> unified, classic workspace not misrouted, env var contamination prevented, legacy IsUnifiedHost fallback)
  • Profile matching tests (all profile types: regular workspace, regular account, SPOG workspace, SPOG account, legacy unified)
  • Token path profile disambiguation tests (workspace_id matching for multi-workspace SPOG)
  • IsAccountsHost() helper tests
  • make checks passes
  • All existing auth tests pass

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
@eng-dev-ecosystem-bot
Copy link
Copy Markdown
Collaborator

eng-dev-ecosystem-bot commented Mar 22, 2026

Commit: e7e316c

Run: 23599480353

Env 🟨​KNOWN 🔄​flaky 💚​RECOVERED 🙈​SKIP ✅​pass 🙈​skip Time
🟨​ aws linux 7 10 270 802 7:26
🟨​ aws windows 7 10 272 800 6:10
💚​ aws-ucws linux 7 10 366 718 10:39
🔄​ aws-ucws windows 2 7 10 366 716 9:02
💚​ azure linux 1 12 273 800 6:33
💚​ azure windows 1 12 275 798 5:39
💚​ azure-ucws linux 1 12 371 714 10:12
💚​ azure-ucws windows 1 12 373 712 8:37
💚​ gcp linux 1 12 269 803 5:39
💚​ gcp windows 1 12 271 801 5:22
19 interesting tests: 10 SKIP, 7 KNOWN, 2 flaky
Test Name aws linux aws windows aws-ucws linux aws-ucws windows azure linux azure windows azure-ucws linux azure-ucws windows gcp linux gcp windows
🟨​ TestAccept 🟨​K 🟨​K 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R
🙈​ TestAccept/bundle/resources/permissions 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
🟨​ TestAccept/bundle/resources/permissions/jobs/destroy_without_mgmtperms/with_permissions 🟨​K 🟨​K 💚​R 💚​R 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
🟨​ TestAccept/bundle/resources/permissions/jobs/destroy_without_mgmtperms/with_permissions/DATABRICKS_BUNDLE_ENGINE=direct 🟨​K 🟨​K 💚​R 💚​R
🟨​ TestAccept/bundle/resources/permissions/jobs/destroy_without_mgmtperms/with_permissions/DATABRICKS_BUNDLE_ENGINE=terraform 🟨​K 🟨​K 💚​R 💚​R
🟨​ TestAccept/bundle/resources/permissions/jobs/destroy_without_mgmtperms/without_permissions 🟨​K 🟨​K 💚​R 💚​R 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
🟨​ TestAccept/bundle/resources/permissions/jobs/destroy_without_mgmtperms/without_permissions/DATABRICKS_BUNDLE_ENGINE=direct 🟨​K 🟨​K 💚​R 💚​R
🟨​ TestAccept/bundle/resources/permissions/jobs/destroy_without_mgmtperms/without_permissions/DATABRICKS_BUNDLE_ENGINE=terraform 🟨​K 🟨​K 💚​R 💚​R
🙈​ TestAccept/bundle/resources/postgres_branches/basic 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
🙈​ TestAccept/bundle/resources/postgres_branches/recreate 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
🙈​ TestAccept/bundle/resources/postgres_branches/update_protected 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
🙈​ TestAccept/bundle/resources/postgres_branches/without_branch_id 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
🙈​ TestAccept/bundle/resources/postgres_endpoints/basic 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
🙈​ TestAccept/bundle/resources/postgres_endpoints/recreate 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
🙈​ TestAccept/bundle/resources/postgres_projects/update_display_name 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
🙈​ TestAccept/bundle/resources/synced_database_tables/basic 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
🙈​ TestAccept/ssh/connection 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
🔄​ TestFilerWorkspaceNotebook ✅​p ✅​p ✅​p 🔄​f ✅​p ✅​p ✅​p ✅​p ✅​p ✅​p
🔄​ TestFilerWorkspaceNotebook/rNb.r ✅​p ✅​p ✅​p 🔄​f ✅​p ✅​p ✅​p ✅​p ✅​p ✅​p
Top 23 slowest tests (at least 2 minutes):
duration env testname
4:13 gcp windows TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=direct
4:05 aws-ucws windows TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=terraform
3:56 azure linux TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=terraform
3:49 gcp linux TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=direct
3:49 gcp windows TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=terraform
3:43 azure-ucws windows TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=direct
3:29 aws-ucws linux TestAccept/bundle/resources/volumes/recreate/DATABRICKS_BUNDLE_ENGINE=direct
3:20 gcp linux TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=terraform
3:14 aws linux TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=direct
3:10 azure-ucws linux TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=terraform
3:04 aws-ucws linux TestAccept/bundle/resources/volumes/recreate/DATABRICKS_BUNDLE_ENGINE=terraform
3:00 azure-ucws linux TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=direct
2:58 azure-ucws windows TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=terraform
2:57 azure-ucws linux TestAccept/bundle/resources/volumes/recreate/DATABRICKS_BUNDLE_ENGINE=terraform
2:53 aws-ucws linux TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=direct
2:51 aws-ucws linux TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=terraform
2:48 azure windows TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=terraform
2:47 aws windows TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=terraform
2:46 aws linux TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=terraform
2:44 aws-ucws windows TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=direct
2:37 aws windows TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=direct
2:14 azure windows TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=direct
2:13 azure linux TestAccept/bundle/resources/apps/inline_config/DATABRICKS_BUNDLE_ENGINE=direct

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
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
@simonfaltum simonfaltum marked this pull request as ready for review March 22, 2026 09:38
@github-actions
Copy link
Copy Markdown

Suggested reviewers

Based on git history of the changed files, these people are best suited to review:

  • @tanmay-db -- recent work in cmd/auth/, libs/auth/, libs/databrickscfg/profile/

Confidence: high

Eligible reviewers

Based 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.

Copy link
Copy Markdown
Contributor

@hectorcast-db hectorcast-db left a comment

Choose a reason for hiding this comment

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

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
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
@simonfaltum simonfaltum enabled auto-merge March 24, 2026 13:39
@simonfaltum simonfaltum disabled auto-merge March 25, 2026 22:07
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
@andrewnester andrewnester added this pull request to the merge queue Mar 26, 2026
Merged via the queue into main with commit 39891ac Mar 26, 2026
19 checks passed
@andrewnester andrewnester deleted the simonfaltum/spog-discovery-login branch March 26, 2026 15:38
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants