Skip to content

Commit b999e77

Browse files
authored
auth: suggest ACCOUNT-<id> as default profile name for --skip-workspace (#5339)
## Why When you run \`databricks auth login --host X --skip-workspace\`, the profile-name prompt suggests the host's first DNS label (e.g. \`db-deco-test\`) even though the resulting profile is account-only. ## Changes \`getProfileName\` already returns \`ACCOUNT-<account-id>\` when \`AccountID\` is populated, but \`setHostAndAccountId\` (which fills \`AccountID\` from \`.well-known/databricks-config\`) only runs *after* the profile-name prompt. So at prompt time, \`AccountID\` was always empty for the no-\`--profile\` flow. - Before: \`auth login --host https://db-deco-test... --skip-workspace\` → prompt suggests \`db-deco-test\`. - Now: same command → URL-param extraction + host discovery run eagerly before the prompt, so the suggestion is \`ACCOUNT-<account-id>\`. Other login paths are unchanged. The discovery call is redundant with the one in \`setHostAndAccountId\` later in the flow; \`runHostDiscovery\` is idempotent for the auth-arguments fields it touches, so the cost is at most one extra \`.well-known\` round-trip (≤5s timeout) and only on the \`--skip-workspace\` path. ## Test plan - [x] New unit tests: \`TestGetProfileName\`, \`TestSkipWorkspaceProfileNameUsesDiscoveredAccountID\` (uses a mock \`.well-known\` server) - [x] \`go test ./cmd/auth/...\` (full package) - [x] \`./task checks\` - [x] \`./task lint-q\`
1 parent a57fe14 commit b999e77

2 files changed

Lines changed: 66 additions & 0 deletions

File tree

cmd/auth/login.go

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,23 @@ a new profile is created.
221221
}
222222
}
223223

224+
// If --skip-workspace is set and we already know the host, eagerly run
225+
// URL-param extraction and host discovery so getProfileName can suggest
226+
// ACCOUNT-<account-id> as the default name. Without this, AccountID is
227+
// only populated by setHostAndAccountId further down, which runs after
228+
// the profile-name prompt.
229+
if skipWorkspace && profileName == "" && authArguments.Host != "" && authArguments.AccountID == "" {
230+
params := auth.ExtractHostQueryParams(authArguments.Host)
231+
authArguments.Host = params.Host
232+
if authArguments.AccountID == "" {
233+
authArguments.AccountID = params.AccountID
234+
}
235+
if authArguments.WorkspaceID == "" {
236+
authArguments.WorkspaceID = params.WorkspaceID
237+
}
238+
runHostDiscovery(ctx, authArguments)
239+
}
240+
224241
// If the user has not specified a profile name, prompt for one.
225242
if profileName == "" {
226243
var err error

cmd/auth/login_test.go

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -564,6 +564,55 @@ func TestSetHostAndAccountId_URLParamsOverrideProfile(t *testing.T) {
564564
assert.Equal(t, "99999", args.WorkspaceID)
565565
}
566566

567+
func TestGetProfileName(t *testing.T) {
568+
tests := []struct {
569+
name string
570+
args *auth.AuthArguments
571+
want string
572+
}{
573+
{
574+
name: "account id set",
575+
args: &auth.AuthArguments{Host: "https://db-deco-test.databricks.com", AccountID: "abc-123"},
576+
want: "ACCOUNT-abc-123",
577+
},
578+
{
579+
name: "no account id falls back to host",
580+
args: &auth.AuthArguments{Host: "https://db-deco-test.databricks.com"},
581+
want: "db-deco-test",
582+
},
583+
}
584+
for _, tt := range tests {
585+
t.Run(tt.name, func(t *testing.T) {
586+
assert.Equal(t, tt.want, getProfileName(tt.args))
587+
})
588+
}
589+
}
590+
591+
// TestSkipWorkspaceProfileNameUsesDiscoveredAccountID verifies that the
592+
// pre-naming discovery block populates AccountID from .well-known so the
593+
// profile-name prompt suggests ACCOUNT-<account-id> instead of the host-based
594+
// default.
595+
func TestSkipWorkspaceProfileNameUsesDiscoveredAccountID(t *testing.T) {
596+
server := newDiscoveryServer(t, map[string]any{
597+
"account_id": "abc-123",
598+
"oidc_endpoint": "https://spog.example.com/oidc/accounts/abc-123",
599+
})
600+
601+
ctx := t.Context()
602+
args := &auth.AuthArguments{Host: server.URL}
603+
604+
// Mirrors the pre-naming block in newLoginCommand's RunE for --skip-workspace.
605+
params := auth.ExtractHostQueryParams(args.Host)
606+
args.Host = params.Host
607+
if args.AccountID == "" {
608+
args.AccountID = params.AccountID
609+
}
610+
runHostDiscovery(ctx, args)
611+
612+
assert.Equal(t, "abc-123", args.AccountID)
613+
assert.Equal(t, "ACCOUNT-abc-123", getProfileName(args))
614+
}
615+
567616
func TestValidateDiscoveryFlagCompatibility(t *testing.T) {
568617
tests := []struct {
569618
name string

0 commit comments

Comments
 (0)