Skip to content

Commit fe485cc

Browse files
authored
Improve auth token UX: profile selection and better empty-state handling (#4584)
## Summary When `databricks auth token` is invoked with no arguments, resolve through env vars (`DATABRICKS_HOST`, `DATABRICKS_CONFIG_PROFILE`) first, then show an interactive profile picker or a clear non-interactive error. This is part of us moving away from a host-focused auth experience, to a profile-focused auth experience ### Interactive profile picker ``` Search: █ Select a profile DEFAULT (https://e2-dogfood.staging.cloud.databricks.com) e2-demo-west (https://e2-demo-west.cloud.databricks.com) tilefood (https://eng-ml-agent-bricks-us-east-1.staging.cloud.databricks.com/) logfood (https://adb-2548836972759138.18.azuredatabricks.net/) Create a new profile Enter a host URL manually ``` **Item ordering rationale:** - Profiles are listed first — this is the primary use case (selecting an existing profile) - "Create a new profile" appears right after profiles — the natural next action when no existing profile fits - "Enter a host URL manually" is last — a legacy escape hatch for power users, de-prioritized in the "profile == identity" model **Search mode (`>5` profiles):** - Auto-activates search to let users find profiles quickly by name or host URL - Pasting a host URL from muscle memory still works — search matches against host strings - Typing "create" or "host" surfaces the action items **`<=5` profiles:** All items visible without scrolling, no search needed. ### "Create a new profile" runs login inline Selecting "Create a new profile" now runs the full `auth login` flow inline (prompt for profile name, prompt for host, OAuth browser challenge, save to `.databrickscfg`) and then returns the token — instead of printing an error telling the user to run a separate command. ### Non-interactive mode - `"no profile specified. Use --profile <name>"` (profiles exist) - `"no profiles configured. Run 'databricks auth login'"` (no profiles) ### Unchanged - Existing `--profile`, `--host`, and positional arg behavior ## Test plan - [x] Unit tests: 6 new test cases covering all non-interactive paths and env var resolution - [x] Acceptance tests: `no-args-with-profiles` and `no-args-no-profiles` - [x] All existing tests pass unchanged - [x] `make lintfull` and `make checks` pass - [x] Manual testing: interactive profile picker, env var resolution, existing flags - [x] Manual testing: "Create a new profile" inline login flow
1 parent 3542987 commit fe485cc

File tree

10 files changed

+381
-14
lines changed

10 files changed

+381
-14
lines changed

acceptance/cmd/auth/token/no-args-no-profiles/out.test.toml

Lines changed: 5 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Error: no profiles configured. Run 'databricks auth login' to create a profile
2+
3+
Exit code: 1
Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
sethome "./home"
2+
3+
unset DATABRICKS_HOST
4+
unset DATABRICKS_TOKEN
5+
unset DATABRICKS_CONFIG_PROFILE
6+
7+
# No config file, non-interactive: should error with login hint
8+
errcode $CLI auth token
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Ignore = [
2+
"home"
3+
]

acceptance/cmd/auth/token/no-args-with-profiles/out.test.toml

Lines changed: 5 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Error: no profile specified. Use --profile <name> to specify which profile to use
2+
3+
Exit code: 1
Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
sethome "./home"
2+
3+
unset DATABRICKS_HOST
4+
unset DATABRICKS_TOKEN
5+
unset DATABRICKS_CONFIG_PROFILE
6+
7+
# Create a .databrickscfg with a profile
8+
cat > "./home/.databrickscfg" <<'ENDCFG'
9+
[myprofile]
10+
host = https://myworkspace.cloud.databricks.com
11+
auth_type = databricks-cli
12+
ENDCFG
13+
14+
# No arguments, non-interactive: should error with profile hint
15+
errcode $CLI auth token
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Ignore = [
2+
"home"
3+
]

cmd/auth/token.go

Lines changed: 221 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,15 @@ import (
55
"encoding/json"
66
"errors"
77
"fmt"
8+
"os"
89
"strings"
910
"time"
1011

1112
"github.com/databricks/cli/libs/auth"
1213
"github.com/databricks/cli/libs/cmdio"
14+
"github.com/databricks/cli/libs/databrickscfg"
1315
"github.com/databricks/cli/libs/databrickscfg/profile"
16+
"github.com/databricks/cli/libs/env"
1417
"github.com/databricks/databricks-sdk-go/config"
1518
"github.com/databricks/databricks-sdk-go/credentials/u2m"
1619
"github.com/databricks/databricks-sdk-go/credentials/u2m/cache"
@@ -24,6 +27,30 @@ func helpfulError(ctx context.Context, profile string, persistentAuth u2m.OAuthA
2427
return fmt.Sprintf("Try logging in again with `%s` before retrying. If this fails, please report this issue to the Databricks CLI maintainers at https://github.com/databricks/cli/issues/new", loginMsg)
2528
}
2629

30+
// profileSelectionResult represents the user's choice from the interactive
31+
// profile picker.
32+
type profileSelectionResult int
33+
34+
const (
35+
profileSelected profileSelectionResult = iota // User picked a profile
36+
enterHostSelected // User chose "Enter a host URL manually"
37+
createNewSelected // User chose "Create a new profile"
38+
)
39+
40+
// applyUnifiedHostFlags copies unified host fields from the profile to the
41+
// auth arguments when they are not already set.
42+
func applyUnifiedHostFlags(p *profile.Profile, args *auth.AuthArguments) {
43+
if p == nil {
44+
return
45+
}
46+
if !args.IsUnifiedHost && p.IsUnifiedHost {
47+
args.IsUnifiedHost = p.IsUnifiedHost
48+
}
49+
if args.WorkspaceID == "" && p.WorkspaceID != "" {
50+
args.WorkspaceID = p.WorkspaceID
51+
}
52+
}
53+
2754
func newTokenCommand(authArguments *auth.AuthArguments) *cobra.Command {
2855
cmd := &cobra.Command{
2956
Use: "token [HOST_OR_PROFILE]",
@@ -115,14 +142,18 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
115142
return nil, err
116143
}
117144

118-
// Load unified host flags from the profile if available
119-
if existingProfile != nil {
120-
if !args.authArguments.IsUnifiedHost && existingProfile.IsUnifiedHost {
121-
args.authArguments.IsUnifiedHost = existingProfile.IsUnifiedHost
122-
}
123-
if args.authArguments.WorkspaceID == "" && existingProfile.WorkspaceID != "" {
124-
args.authArguments.WorkspaceID = existingProfile.WorkspaceID
145+
applyUnifiedHostFlags(existingProfile, args.authArguments)
146+
147+
// When no explicit profile, host, or positional args are provided, attempt to
148+
// resolve the target through environment variables or interactive profile selection.
149+
if args.profileName == "" && args.authArguments.Host == "" && len(args.args) == 0 {
150+
var resolvedProfile string
151+
resolvedProfile, existingProfile, err = resolveNoArgsToken(ctx, args.profiler, args.authArguments)
152+
if err != nil {
153+
return nil, err
125154
}
155+
args.profileName = resolvedProfile
156+
applyUnifiedHostFlags(existingProfile, args.authArguments)
126157
}
127158

128159
err = setHostAndAccountId(ctx, existingProfile, args.authArguments, args.args)
@@ -228,3 +259,186 @@ func askForMatchingProfile(ctx context.Context, profiles profile.Profiles, host
228259
}
229260
return profiles[i].Name, nil
230261
}
262+
263+
// resolveNoArgsToken resolves a profile or host when `auth token` is invoked
264+
// with no explicit profile, host, or positional arguments. It checks environment
265+
// variables first, then falls back to interactive profile selection or a clear
266+
// non-interactive error.
267+
//
268+
// Returns the resolved profile name and profile (if any). The host and related
269+
// fields on authArgs are updated in place when resolved via environment variables.
270+
func resolveNoArgsToken(ctx context.Context, profiler profile.Profiler, authArgs *auth.AuthArguments) (string, *profile.Profile, error) {
271+
// Step 1: Try DATABRICKS_HOST env var (highest priority).
272+
if envHost := env.Get(ctx, "DATABRICKS_HOST"); envHost != "" {
273+
authArgs.Host = envHost
274+
if v := env.Get(ctx, "DATABRICKS_ACCOUNT_ID"); v != "" {
275+
authArgs.AccountID = v
276+
}
277+
if v := env.Get(ctx, "DATABRICKS_WORKSPACE_ID"); v != "" {
278+
authArgs.WorkspaceID = v
279+
}
280+
if ok, _ := env.GetBool(ctx, "DATABRICKS_EXPERIMENTAL_IS_UNIFIED_HOST"); ok {
281+
authArgs.IsUnifiedHost = true
282+
}
283+
return "", nil, nil
284+
}
285+
286+
// Step 2: Try DATABRICKS_CONFIG_PROFILE env var.
287+
if envProfile := env.Get(ctx, "DATABRICKS_CONFIG_PROFILE"); envProfile != "" {
288+
p, err := loadProfileByName(ctx, envProfile, profiler)
289+
if err != nil {
290+
return "", nil, err
291+
}
292+
return envProfile, p, nil
293+
}
294+
295+
// Step 3: No env vars resolved. Load all profiles for interactive selection
296+
// or non-interactive error.
297+
allProfiles, err := profiler.LoadProfiles(ctx, profile.MatchAllProfiles)
298+
if err != nil && !errors.Is(err, profile.ErrNoConfiguration) {
299+
return "", nil, err
300+
}
301+
302+
if !cmdio.IsPromptSupported(ctx) {
303+
if len(allProfiles) > 0 {
304+
return "", nil, errors.New("no profile specified. Use --profile <name> to specify which profile to use")
305+
}
306+
return "", nil, errors.New("no profiles configured. Run 'databricks auth login' to create a profile")
307+
}
308+
309+
// Interactive: show profile picker.
310+
result, selectedName, err := promptForProfileSelection(ctx, allProfiles)
311+
if err != nil {
312+
return "", nil, err
313+
}
314+
switch result {
315+
case enterHostSelected:
316+
// Fall through — setHostAndAccountId will prompt for the host.
317+
return "", nil, nil
318+
case createNewSelected:
319+
return runInlineLogin(ctx, profiler)
320+
default:
321+
p, err := loadProfileByName(ctx, selectedName, profiler)
322+
if err != nil {
323+
return "", nil, err
324+
}
325+
return selectedName, p, nil
326+
}
327+
}
328+
329+
// profileSelectItem is used by promptForProfileSelection to render both
330+
// regular profiles and special action options in the same select list.
331+
type profileSelectItem struct {
332+
Name string
333+
Host string
334+
}
335+
336+
// promptForProfileSelection shows a promptui select list with all configured
337+
// profiles plus "Enter a host URL" and "Create a new profile" options.
338+
// Returns the selection type and, when a profile is selected, its name.
339+
func promptForProfileSelection(ctx context.Context, profiles profile.Profiles) (profileSelectionResult, string, error) {
340+
items := make([]profileSelectItem, 0, len(profiles)+2)
341+
for _, p := range profiles {
342+
items = append(items, profileSelectItem{Name: p.Name, Host: p.Host})
343+
}
344+
createProfileIdx := len(items)
345+
items = append(items, profileSelectItem{Name: "Create a new profile"})
346+
enterHostIdx := len(items)
347+
items = append(items, profileSelectItem{Name: "Enter a host URL manually"})
348+
349+
i, _, err := cmdio.RunSelect(ctx, &promptui.Select{
350+
Label: "Select a profile",
351+
Items: items,
352+
StartInSearchMode: len(profiles) > 5,
353+
Searcher: func(input string, index int) bool {
354+
input = strings.ToLower(input)
355+
name := strings.ToLower(items[index].Name)
356+
host := strings.ToLower(items[index].Host)
357+
return strings.Contains(name, input) || strings.Contains(host, input)
358+
},
359+
Templates: &promptui.SelectTemplates{
360+
Label: "{{ . | faint }}",
361+
Active: `{{.Name | bold}}{{if .Host}} ({{.Host|faint}}){{end}}`,
362+
Inactive: `{{.Name}}{{if .Host}} ({{.Host}}){{end}}`,
363+
Selected: `{{ "Using profile" | faint }}: {{ .Name | bold }}`,
364+
},
365+
})
366+
if err != nil {
367+
return 0, "", err
368+
}
369+
370+
switch i {
371+
case enterHostIdx:
372+
return enterHostSelected, "", nil
373+
case createProfileIdx:
374+
return createNewSelected, "", nil
375+
default:
376+
return profileSelected, profiles[i].Name, nil
377+
}
378+
}
379+
380+
// runInlineLogin runs a minimal interactive login flow: prompts for a profile
381+
// name and host, performs the OAuth challenge, saves the profile to
382+
// .databrickscfg, and returns the new profile name and profile.
383+
func runInlineLogin(ctx context.Context, profiler profile.Profiler) (string, *profile.Profile, error) {
384+
profileName, err := promptForProfile(ctx, "DEFAULT")
385+
if err != nil {
386+
return "", nil, err
387+
}
388+
389+
existingProfile, err := loadProfileByName(ctx, profileName, profiler)
390+
if err != nil {
391+
return "", nil, err
392+
}
393+
394+
loginArgs := &auth.AuthArguments{}
395+
applyUnifiedHostFlags(existingProfile, loginArgs)
396+
397+
err = setHostAndAccountId(ctx, existingProfile, loginArgs, nil)
398+
if err != nil {
399+
return "", nil, err
400+
}
401+
402+
loginArgs.Profile = profileName
403+
404+
oauthArgument, err := loginArgs.ToOAuthArgument()
405+
if err != nil {
406+
return "", nil, err
407+
}
408+
persistentAuth, err := u2m.NewPersistentAuth(ctx,
409+
u2m.WithOAuthArgument(oauthArgument),
410+
u2m.WithBrowser(openURLSuppressingStderr),
411+
)
412+
if err != nil {
413+
return "", nil, err
414+
}
415+
defer persistentAuth.Close()
416+
417+
ctx, cancel := context.WithTimeout(ctx, defaultTimeout)
418+
defer cancel()
419+
420+
if err = persistentAuth.Challenge(); err != nil {
421+
return "", nil, err
422+
}
423+
424+
err = databrickscfg.SaveToProfile(ctx, &config.Config{
425+
Profile: profileName,
426+
Host: loginArgs.Host,
427+
AuthType: authTypeDatabricksCLI,
428+
AccountID: loginArgs.AccountID,
429+
WorkspaceID: loginArgs.WorkspaceID,
430+
Experimental_IsUnifiedHost: loginArgs.IsUnifiedHost,
431+
ConfigFile: os.Getenv("DATABRICKS_CONFIG_FILE"),
432+
})
433+
if err != nil {
434+
return "", nil, err
435+
}
436+
437+
cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully saved", profileName))
438+
439+
p, err := loadProfileByName(ctx, profileName, profiler)
440+
if err != nil {
441+
return "", nil, err
442+
}
443+
return profileName, p, nil
444+
}

0 commit comments

Comments
 (0)