From ddfa00c1696b5cd51f4a45ea7cf6a8a21d487ecf Mon Sep 17 00:00:00 2001 From: simon Date: Wed, 27 May 2026 15:32:45 +0200 Subject: [PATCH] auth: --skip-workspace lands on the account selector in discovery flow `databricks auth login --skip-workspace` (no --host) now lands the user on the account selector at login.databricks.com instead of the workspace selector. The workspace pick was a wasted step when the caller already said they only want account-level access. Wired through u2m.WithDiscoveryAccountTarget() from databricks-sdk-go, which sets target=ACCOUNT on the discovery authorize URL. Co-authored-by: Isaac --- NEXT_CHANGELOG.md | 1 + cmd/auth/login.go | 8 +++++++ cmd/auth/login_test.go | 52 ++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 59 insertions(+), 2 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 004cd4da4c..325835da0b 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -5,6 +5,7 @@ ### Notable Changes ### CLI +* `databricks auth login --skip-workspace` (without `--host`) now lands the user on the account selector at `login.databricks.com` instead of the workspace selector, skipping a step that is wasted for account-only logins. ### Bundles * The error reported when a direct-only resource (catalogs, external locations, vector search endpoints) is used with the terraform engine now also suggests setting `bundle.engine: direct` in `databricks.yml`, in addition to the `DATABRICKS_BUNDLE_ENGINE` environment variable ([#5295](https://github.com/databricks/cli/pull/5295)). diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 475a72ed36..dc428ddf0d 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -251,6 +251,7 @@ a new profile is created. profileName: profileName, timeout: loginTimeout, scopes: scopes, + skipWorkspace: skipWorkspace, existingProfile: existingProfile, browserFunc: getBrowserFunc(cmd), tokenCache: tokenCache, @@ -600,6 +601,7 @@ type discoveryLoginInputs struct { profileName string timeout time.Duration scopes string + skipWorkspace bool existingProfile *profile.Profile browserFunc func(string) error tokenCache cache.TokenCache @@ -629,6 +631,12 @@ func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error { if len(scopesList) > 0 { opts = append(opts, u2m.WithScopes(scopesList)) } + // --skip-workspace lands the user on the account selector at + // login.databricks.com instead of the workspace selector, which is a + // wasted step for account-only logins. + if in.skipWorkspace { + opts = append(opts, u2m.WithDiscoveryAccountTarget()) + } discoveryHost := env.Get(ctx, discoveryHostEnvVar) if discoveryHost != "" { opts = append(opts, u2m.WithDiscoveryHost(discoveryHost)) diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 101f2e92e5..0dfcd77fb1 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -87,8 +87,9 @@ type fakeDiscoveryClient struct { introspection *auth.IntrospectionResult introspectionErr error // For assertions - introspectHost string - introspectToken string + introspectHost string + introspectToken string + persistentAuthOpts []u2m.PersistentAuthOption } func (f *fakeDiscoveryClient) NewOAuthArgument(profileName string) (*u2m.BasicDiscoveryOAuthArgument, error) { @@ -99,6 +100,7 @@ func (f *fakeDiscoveryClient) NewOAuthArgument(profileName string) (*u2m.BasicDi } func (f *fakeDiscoveryClient) NewPersistentAuth(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { + f.persistentAuthOpts = opts if f.persistentAuthErr != nil { return nil, f.persistentAuthErr } @@ -1116,6 +1118,52 @@ func TestDiscoveryLogin_OverridesHostFromEnv(t *testing.T) { assert.NotContains(t, stderr.String(), "Opening login.databricks.com in your browser...") } +func TestDiscoveryLogin_SkipWorkspaceAddsAccountTargetOption(t *testing.T) { + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, ".databrickscfg") + require.NoError(t, os.WriteFile(configPath, []byte(""), 0o600)) + t.Setenv("DATABRICKS_CONFIG_FILE", configPath) + + newDiscoveryClient := func(t *testing.T) *fakeDiscoveryClient { + oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY") + require.NoError(t, err) + oauthArg.SetDiscoveredHost("https://workspace.example.com") + return &fakeDiscoveryClient{ + oauthArg: oauthArg, + persistentAuth: &fakeDiscoveryPersistentAuth{ + token: &oauth2.Token{AccessToken: "test-token"}, + }, + introspectionErr: errors.New("not asserted here"), + } + } + + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + + // Baseline: no --skip-workspace. + baselineDC := newDiscoveryClient(t) + require.NoError(t, discoveryLogin(ctx, discoveryLoginInputs{ + dc: baselineDC, + profileName: "DISCOVERY", + timeout: time.Second, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + })) + + // With --skip-workspace: expect exactly one extra option. + skipDC := newDiscoveryClient(t) + require.NoError(t, discoveryLogin(ctx, discoveryLoginInputs{ + dc: skipDC, + profileName: "DISCOVERY", + timeout: time.Second, + skipWorkspace: true, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + })) + + assert.Equal(t, len(baselineDC.persistentAuthOpts)+1, len(skipDC.persistentAuthOpts), + "--skip-workspace should add exactly one extra PersistentAuthOption (WithDiscoveryAccountTarget)") +} + func TestLoginRejectsPositionalArgWithHostFlag(t *testing.T) { ctx := cmdio.MockDiscard(t.Context()) authArgs := &auth.AuthArguments{Host: "https://example.com"}