diff --git a/acceptance/cmd/auth/login/device-code/unsupported-host/output.txt b/acceptance/cmd/auth/login/device-code/unsupported-host/output.txt index 5d327ce2db2..7875f49ed6b 100644 --- a/acceptance/cmd/auth/login/device-code/unsupported-host/output.txt +++ b/acceptance/cmd/auth/login/device-code/unsupported-host/output.txt @@ -2,7 +2,7 @@ === No host >>> musterr [CLI] auth login --profile new-profile --use-device-code -Error: --use-device-code requires a workspace host. Pass it as an argument or with --host +Error: the command is being run in a non-interactive environment, please specify a host using --host === Account host diff --git a/acceptance/cmd/auth/login/device-code/unsupported-host/script b/acceptance/cmd/auth/login/device-code/unsupported-host/script index dd8a9ddca47..d744aec1709 100644 --- a/acceptance/cmd/auth/login/device-code/unsupported-host/script +++ b/acceptance/cmd/auth/login/device-code/unsupported-host/script @@ -3,7 +3,8 @@ sethome "./home" title "No host\n" # The profile does not exist yet and no host is given. auth login does not read # DATABRICKS_HOST, so the usual login would fall back to the browser-based -# discovery flow, which device code login does not support. +# discovery flow. Device code login skips it and asks for the host instead, +# which fails without a terminal. trace musterr $CLI auth login --profile new-profile --use-device-code envsubst < "./home/.databrickscfg.tmpl" > "./home/.databrickscfg" diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 13a5a8a61e0..43463c6f53f 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -225,9 +225,15 @@ a new profile is created. // the user re-login to an existing profile, create a new one, or enter // a host URL. With no profiles configured the picker still shows the // two action entries so the user can choose between web-based discovery - // (Create a new profile) and a manual host URL. + // (Create a new profile) and a manual host URL. Device code login skips + // discovery, so for it "Create a new profile" prompts for the host after + // the profile name. if profileName == "" && authArguments.Host == "" && len(args) == 0 && cmdio.IsPromptSupported(ctx) { - allProfiles, err := profile.DefaultProfiler.LoadProfiles(ctx, profile.MatchAllProfiles) + matchProfiles := profile.MatchAllProfiles + if useDeviceCode { + matchProfiles = matchDeviceCodeProfiles + } + allProfiles, err := profile.DefaultProfiler.LoadProfiles(ctx, matchProfiles) if err != nil && !errors.Is(err, profile.ErrNoConfiguration) { return err } @@ -306,11 +312,10 @@ a new profile is created. } // If no host is available from any source, use the discovery flow - // via login.databricks.com. - if shouldUseDiscovery(authArguments.Host, args, existingProfile) { - if useDeviceCode { - return errDeviceCodeNeedsHost - } + // via login.databricks.com. Discovery opens a browser on this machine, + // and device code login is meant for machines without one, so it skips + // discovery and setHost resolves the host instead. + if !useDeviceCode && shouldUseDiscovery(authArguments.Host, args, existingProfile) { if err := validateDiscoveryFlagCompatibility(cmd); err != nil { return err } @@ -332,11 +337,10 @@ a new profile is created. return err } // Check the host once it is resolved, before setAccountID prompts for - // an account ID for account and unified hosts. - if useDeviceCode { - if err := validateDeviceCodeHost(authArguments); err != nil { - return err - } + // an account ID for account and unified hosts. A workspace without + // device code login enabled is rejected later by DeviceCodeChallenge. + if useDeviceCode && !isDeviceCodeSupportedForHost(authArguments.Host, authArguments.DiscoveryURL) { + return errDeviceCodeWorkspaceOnly } err = setAccountID(ctx, existingProfile, authArguments) if err != nil { diff --git a/cmd/auth/login_device_code.go b/cmd/auth/login_device_code.go index 93455ad9dd7..971f31ea71b 100644 --- a/cmd/auth/login_device_code.go +++ b/cmd/auth/login_device_code.go @@ -6,9 +6,9 @@ import ( "fmt" "strings" - "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/auth/u2m" "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/env" "github.com/spf13/cobra" ) @@ -22,19 +22,20 @@ var deviceCodeIncompatibleFlags = []string{ "skip-workspace", } -var ( - errDeviceCodeNeedsHost = errors.New("--use-device-code requires a workspace host. Pass it as an argument or with --host") - errDeviceCodeWorkspaceOnly = errors.New("--use-device-code is only supported for workspace hosts") -) +var errDeviceCodeWorkspaceOnly = errors.New("--use-device-code is only supported for workspace hosts") -// validateDeviceCodeHost returns an error if the host resolved by setHost is -// an account or unified host. A workspace without device code login enabled -// is rejected by DeviceCodeChallenge before it creates a device code. -func validateDeviceCodeHost(authArguments *auth.AuthArguments) error { - if isAccountOrUnifiedHost(authArguments.Host, authArguments.DiscoveryURL) { - return errDeviceCodeWorkspaceOnly - } - return nil +// isDeviceCodeSupportedForHost reports whether host is a workspace host. A +// unified host is only recognized by its account-scoped discoveryURL, so with +// an empty discoveryURL only classic account hosts are rejected. +func isDeviceCodeSupportedForHost(host, discoveryURL string) bool { + return !isAccountOrUnifiedHost(host, discoveryURL) +} + +// matchDeviceCodeProfiles matches the profiles that device code login does not +// always reject. Profiles on unified hosts still match, because the discovery +// URL that identifies them is not saved on the profile. +func matchDeviceCodeProfiles(p profile.Profile) bool { + return isDeviceCodeSupportedForHost(p.Host, "") && len(u2mResourcesFromProfile(&p)) == 0 } // validateDeviceCodeLogin returns an error if device code login is not enabled diff --git a/cmd/auth/login_device_code_test.go b/cmd/auth/login_device_code_test.go index 84fdc3f0bbd..f6231d20b1a 100644 --- a/cmd/auth/login_device_code_test.go +++ b/cmd/auth/login_device_code_test.go @@ -12,6 +12,8 @@ import ( "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/auth/u2m" "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/cmdio/cmdiotest/termtest" + "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/flags" "github.com/spf13/cobra" @@ -167,7 +169,7 @@ func TestFormatCode(t *testing.T) { assert.Equal(t, "ABC", formatCode("ABC")) } -func TestValidateDeviceCodeHost(t *testing.T) { +func TestIsDeviceCodeSupportedForHost(t *testing.T) { workspace := newDiscoveryServer(t, map[string]any{ "account_id": "acct-123", "workspace_id": "12345", @@ -184,23 +186,104 @@ func TestValidateDeviceCodeHost(t *testing.T) { }) tests := []struct { - name string - args auth.AuthArguments - wantErr error + name string + args auth.AuthArguments + want bool }{ - {name: "workspace host", args: auth.AuthArguments{Host: workspace.URL}}, - {name: "workspace host with query params", args: auth.AuthArguments{Host: workspace.URL + "?o=12345"}}, - {name: "classic account host", args: auth.AuthArguments{Host: "https://accounts.cloud.databricks.test"}, wantErr: errDeviceCodeWorkspaceOnly}, - {name: "unified host", args: auth.AuthArguments{Host: unified.URL}, wantErr: errDeviceCodeWorkspaceOnly}, - {name: "unified host with query params", args: auth.AuthArguments{Host: unified.URL + "?o=12345"}, wantErr: errDeviceCodeWorkspaceOnly}, - {name: "unified host with account ID in query params", args: auth.AuthArguments{Host: unifiedWithoutAccountID.URL + "?a=acct-123"}, wantErr: errDeviceCodeWorkspaceOnly}, - {name: "unified host with account ID flag", args: auth.AuthArguments{Host: unifiedWithoutAccountID.URL, AccountID: "acct-123"}, wantErr: errDeviceCodeWorkspaceOnly}, + {name: "workspace host", args: auth.AuthArguments{Host: workspace.URL}, want: true}, + {name: "workspace host with query params", args: auth.AuthArguments{Host: workspace.URL + "?o=12345"}, want: true}, + {name: "classic account host", args: auth.AuthArguments{Host: "https://accounts.cloud.databricks.test"}}, + {name: "unified host", args: auth.AuthArguments{Host: unified.URL}}, + {name: "unified host with query params", args: auth.AuthArguments{Host: unified.URL + "?o=12345"}}, + {name: "unified host with account ID in query params", args: auth.AuthArguments{Host: unifiedWithoutAccountID.URL + "?a=acct-123"}}, + {name: "unified host with account ID flag", args: auth.AuthArguments{Host: unifiedWithoutAccountID.URL, AccountID: "acct-123"}}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { args := tt.args require.NoError(t, setHost(t.Context(), nil, &args, nil)) - assert.ErrorIs(t, validateDeviceCodeHost(&args), tt.wantErr) + assert.Equal(t, tt.want, isDeviceCodeSupportedForHost(args.Host, args.DiscoveryURL)) }) } } + +// startDeviceCodeLogin runs "auth login --use-device-code" with args in a +// terminal emulator, using config as the contents of .databrickscfg. +func startDeviceCodeLogin(t *testing.T, config string, args ...string) *termtest.Term[struct{}] { + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("USERPROFILE", home) + t.Setenv("DATABRICKS_CONFIG_FILE", writeTempConfig(t, config)) + t.Setenv("DATABRICKS_AUTH_STORAGE", "plaintext") + t.Setenv(deviceCodeEnvVar, "true") + + return termtest.Run(t, func(ctx context.Context) (struct{}, error) { + cmd := newLoginCommand(&auth.AuthArguments{}) + cmd.Flags().String("profile", "", "") + cmd.SetContext(ctx) + cmd.SetArgs(append([]string{"--use-device-code"}, args...)) + cmd.SetOut(io.Discard) + cmd.SetErr(io.Discard) + return struct{}{}, cmd.Execute() + }) +} + +func TestLoginDeviceCodeNewProfilePromptsForHost(t *testing.T) { + tm := startDeviceCodeLogin(t, "", "--profile", "newprof") + tm.WaitFor("Databricks host") + tm.Type(termtest.KeyCtrlC) + _, err := tm.Result() + assert.EqualError(t, err, "^C") +} + +func TestMatchDeviceCodeProfiles(t *testing.T) { + tests := []struct { + name string + profile profile.Profile + want bool + }{ + {name: "workspace host", profile: profile.Profile{Host: "https://workspace.test"}, want: true}, + // setHost prompts for the host of a profile without one. + {name: "no host", profile: profile.Profile{}, want: true}, + {name: "classic account host", profile: profile.Profile{Host: "https://accounts.cloud.databricks.test", AccountID: "acct-123"}}, + {name: "classic account host for DoD", profile: profile.Profile{Host: "https://accounts-dod.cloud.databricks.test", AccountID: "acct-123"}}, + { + name: "scoped to resources", + profile: profile.Profile{Host: "https://workspace.test", AuthType: authTypeDatabricksCLI, Resources: "https://workspace.test/ai-gateway/mcp-services/system.ai.github"}, + }, + // Login only reuses the resources of a databricks-cli profile. + { + name: "resources on another auth type", + profile: profile.Profile{Host: "https://workspace.test", AuthType: "pat", Resources: "https://workspace.test/ai-gateway/mcp-services/system.ai.github"}, + want: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, matchDeviceCodeProfiles(tt.profile)) + }) + } +} + +func TestLoginDeviceCodePickerHidesUnsupportedProfiles(t *testing.T) { + tm := startDeviceCodeLogin(t, `[workspace-prof] +host = https://workspace.test + +[account-prof] +host = https://accounts.cloud.databricks.test +account_id = acct-123 + +[scoped-prof] +host = https://workspace.test +auth_type = databricks-cli +resources = https://workspace.test/ai-gateway/mcp-services/system.ai.github +`) + tm.WaitFor(profilePickerEnterHostLabel) + screen := tm.Snapshot() + assert.Contains(t, screen, "workspace-prof") + assert.NotContains(t, screen, "account-prof") + assert.NotContains(t, screen, "scoped-prof") + tm.Type(termtest.KeyCtrlC) + _, err := tm.Result() + assert.EqualError(t, err, "^C") +} diff --git a/docs/device-code-login.md b/docs/device-code-login.md index 87c96c03c16..08f1fcab358 100644 --- a/docs/device-code-login.md +++ b/docs/device-code-login.md @@ -33,7 +33,11 @@ Profile my-profile was successfully saved ``` 1. Run `databricks auth login` with `--use-device-code`. The host can be passed as - an argument, with `--host`, or come from an existing profile. + an argument, with `--host`, or come from an existing profile. Otherwise the + CLI prompts for it. The profile picker leaves out profiles that device code + login cannot use: profiles for an account host and profiles saved with + `--resource`. Profiles for a unified host are still listed and are rejected + after you pick one. 2. On another device, open the URL that the CLI shows, sign in to the workspace if needed, and enter the code. The code expires after about 10 minutes.