diff --git a/.nextchanges/cli/auth-workspace-search.md b/.nextchanges/cli/auth-workspace-search.md new file mode 100644 index 00000000000..37b0cf7df1d --- /dev/null +++ b/.nextchanges/cli/auth-workspace-search.md @@ -0,0 +1 @@ +* Remove the 50-workspace limit in `auth login` and support searching workspaces by name or ID. ([#7071](https://github.com/databricks/cli/pull/7071)) diff --git a/cmd/auth/export_test.go b/cmd/auth/export_test.go new file mode 100644 index 00000000000..ad363a491ac --- /dev/null +++ b/cmd/auth/export_test.go @@ -0,0 +1,6 @@ +package auth + +var ( + NewWorkspaceSearcher = newWorkspaceSearcher + PromptForWorkspaceSelection = promptForWorkspaceSelection +) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 1f0ddb9b707..a5e39d77f8d 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -882,12 +882,6 @@ func promptForWorkspaceSelection(ctx context.Context, authArguments *auth.AuthAr return "", nil } - const maxWorkspaces = 50 - if len(workspaces) > maxWorkspaces { - cmdio.LogString(ctx, fmt.Sprintf("Account has %d workspaces. Showing first %d. Use --workspace-id to specify directly.", len(workspaces), maxWorkspaces)) - workspaces = workspaces[:maxWorkspaces] - } - if len(workspaces) == 1 { wsID := strconv.FormatInt(workspaces[0].WorkspaceId, 10) cmdio.LogString(ctx, fmt.Sprintf("Auto-selected workspace %q (%s)", workspaces[0].WorkspaceName, wsID)) @@ -907,11 +901,38 @@ func promptForWorkspaceSelection(ctx context.Context, authArguments *auth.AuthAr Id: "", }) - selected, err := cmdio.SelectOrdered(ctx, items, "Select a workspace") + idx, err := cmdio.RunSelect(ctx, cmdio.SelectOptions{ + Label: "Select a workspace (search by name or ID)", + Items: items, + HideSelected: true, + StartInSearchMode: true, + Searcher: newWorkspaceSearcher(items), + Active: `▸ {{if .Name}}{{.Name | bold}}{{if .Id}} ({{.Id | faint}}){{end}}{{else}}{{.Id | bold}}{{end}}`, + Inactive: ` {{if .Name}}{{.Name}}{{if .Id}} ({{.Id | faint}}){{end}}{{else}}{{.Id}}{{end}}`, + }) if err != nil { return "", err } - return selected, nil + return items[idx].Id, nil +} + +// newWorkspaceSearcher gives exact workspace IDs precedence over substring matches. +func newWorkspaceSearcher(items []cmdio.Tuple) func(string, int) bool { + ids := make(map[string]struct{}, len(items)) + names := make([]string, len(items)) + for idx, item := range items { + names[idx] = strings.ToLower(item.Name) + if item.Id != "" { + ids[item.Id] = struct{}{} + } + } + return func(input string, idx int) bool { + input = strings.TrimSpace(input) + if _, exact := ids[input]; exact { + return items[idx].Id == input + } + return strings.Contains(names[idx], strings.ToLower(input)) || strings.Contains(items[idx].Id, input) + } } // promptForWorkspaceID asks the user to manually enter a workspace ID. diff --git a/cmd/auth/login_workspace_test.go b/cmd/auth/login_workspace_test.go new file mode 100644 index 00000000000..b3b83ce1d19 --- /dev/null +++ b/cmd/auth/login_workspace_test.go @@ -0,0 +1,264 @@ +package auth_test + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + authcmd "github.com/databricks/cli/cmd/auth" + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/cmdio/cmdiotest/termtest" + "github.com/databricks/databricks-sdk-go/service/provisioning" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" +) + +const ( + workspaceTestAccountID = "test-account" + workspaceTestToken = "test-token" + workspaceSkipLabel = "Skip (account-level access only)" + workspacePickerLabel = "Select a workspace" + workspaceIDPrompt = "Enter workspace ID (empty to skip)" +) + +func TestWorkspaceSearcher(t *testing.T) { + items := []cmdio.Tuple{ + {Name: "Analytics", Id: "1234"}, + {Name: "Report 123", Id: "99"}, + {Name: "Production", Id: "123"}, + {Name: "Shared", Id: "887766"}, + {Name: "Shared", Id: "778899"}, + {Name: "", Id: "40"}, + {Name: "\u00c5ngstr\u00f6m", Id: "9007199254740993"}, + {Name: "Skip Analytics", Id: "0"}, + {Name: workspaceSkipLabel, Id: ""}, + } + search := authcmd.NewWorkspaceSearcher(items) + for _, tc := range []struct { + name string + query string + want []int + }{ + {name: "empty", want: []int{0, 1, 2, 3, 4, 5, 6, 7, 8}}, + {name: "whitespace", query: " \t ", want: []int{0, 1, 2, 3, 4, 5, 6, 7, 8}}, + {name: "name", query: "aNaLyT", want: []int{0, 7}}, + {name: "trim name", query: " Production ", want: []int{2}}, + {name: "partial id", query: "887", want: []int{3}}, + {name: "partial name or id", query: "12", want: []int{0, 1, 2}}, + {name: "exact id beats longer id and numeric name", query: "123", want: []int{2}}, + {name: "trim exact id", query: "\t123 ", want: []int{2}}, + {name: "duplicate names", query: "shared", want: []int{3, 4}}, + {name: "duplicate name by id", query: "778899", want: []int{4}}, + {name: "empty name", query: "40", want: []int{5}}, + {name: "non ascii name", query: "\u00e5NGSTR\u00d6M", want: []int{6}}, + {name: "large id", query: "9007199254740993", want: []int{6}}, + {name: "zero id is not skip", query: "0", want: []int{7}}, + {name: "skip stays last", query: "sKiP", want: []int{7, 8}}, + {name: "skip by label", query: "account-level", want: []int{8}}, + {name: "missing name", query: "not-a-workspace"}, + {name: "missing id", query: "555555"}, + } { + t.Run(tc.name, func(t *testing.T) { + var matches []int + for i := range items { + if search(tc.query, i) { + matches = append(matches, i) + } + } + assert.Equal(t, tc.want, matches) + }) + } +} + +func TestPromptForWorkspaceSelection(t *testing.T) { + workspaces := make([]provisioning.Workspace, 60) + for i := range workspaces { + workspaces[i] = provisioning.Workspace{ + WorkspaceId: int64(1000 + i), + WorkspaceName: fmt.Sprintf("Workspace %c%c", 'A'+i/26, 'A'+i%26), + } + } + workspaces[59] = provisioning.Workspace{WorkspaceId: 9007199254740993, WorkspaceName: "Target Workspace"} + + for _, tc := range []struct { + name string + workspaces []provisioning.Workspace + input string + want string + }{ + {name: "id beyond first fifty", workspaces: workspaces, input: "9007199254740993", want: "9007199254740993"}, + {name: "name beyond first fifty", workspaces: workspaces, input: "tArGeT", want: "9007199254740993"}, + { + name: "exact id precedence", + workspaces: []provisioning.Workspace{ + {WorkspaceId: 1234, WorkspaceName: "Longer ID"}, + {WorkspaceId: 99, WorkspaceName: "Report 123"}, + {WorkspaceId: 123, WorkspaceName: "Exact ID"}, + }, + input: " 123 ", want: "123", + }, + { + name: "duplicate names", + workspaces: []provisioning.Workspace{ + {WorkspaceId: 111, WorkspaceName: "Shared"}, + {WorkspaceId: 222, WorkspaceName: "Shared"}, + }, + input: "222", want: "222", + }, + { + name: "empty name", + workspaces: []provisioning.Workspace{ + {WorkspaceId: 111, WorkspaceName: "Named"}, + {WorkspaceId: 222}, + }, + input: "222", want: "222", + }, + { + name: "workspace before skip", + workspaces: []provisioning.Workspace{ + {WorkspaceId: 111, WorkspaceName: "Named"}, + {WorkspaceId: 222, WorkspaceName: "Skip Analytics"}, + }, + input: "skip", want: "222", + }, + {name: "skip", workspaces: workspaces, input: "skip", want: ""}, + {name: "whitespace preserves API order", workspaces: workspaces, input: " ", want: "1000"}, + } { + t.Run(tc.name, func(t *testing.T) { + args, requests := newWorkspaceListServer(t, http.StatusOK, tc.workspaces) + tm := termtest.Run(t, func(ctx context.Context) (string, error) { + return authcmd.PromptForWorkspaceSelection(ctx, args, oauth2.StaticTokenSource(&oauth2.Token{AccessToken: workspaceTestToken})) + }) + tm.WaitFor(workspacePickerLabel) + tm.Type(tc.input + termtest.KeyEnter) + id, err := tm.Result() + require.NoError(t, err) + assert.Equal(t, tc.want, id) + assert.EqualValues(t, 1, requests.Load()) + assert.NotContains(t, tm.Raw(), workspaceIDPrompt) + assert.NotContains(t, tm.Raw(), "Showing first") + }) + } +} + +func TestPromptForWorkspaceSelectionNoResults(t *testing.T) { + args, requests := newWorkspaceListServer(t, http.StatusOK, []provisioning.Workspace{ + {WorkspaceId: 111, WorkspaceName: "Alpha"}, + {WorkspaceId: 222, WorkspaceName: "Beta"}, + {WorkspaceId: 333}, + }) + tm := termtest.Run(t, func(ctx context.Context) (string, error) { + return authcmd.PromptForWorkspaceSelection(ctx, args, oauth2.StaticTokenSource(&oauth2.Token{AccessToken: workspaceTestToken})) + }) + tm.WaitFor(workspacePickerLabel) + tm.Golden("01-initial") + + tm.Type("999999") + tm.WaitFor("No results") + tm.Type(termtest.KeyEnter) + tm.Type(strings.Repeat(termtest.KeyBackspace, 6)) + require.Eventually(t, func() bool { + return strings.Contains(tm.Snapshot(), workspaceSkipLabel) && !strings.Contains(tm.Snapshot(), "No results") + }, 5*time.Second, 5*time.Millisecond) + tm.Type(termtest.KeyDown + termtest.KeyEnter) + id, err := tm.Result() + require.NoError(t, err) + assert.Equal(t, "222", id) + assert.EqualValues(t, 1, requests.Load()) +} + +func TestPromptForWorkspaceSelectionWithoutPicker(t *testing.T) { + for _, tc := range []struct { + name string + workspaces []provisioning.Workspace + want string + message string + }{ + {name: "no workspaces", workspaces: []provisioning.Workspace{}}, + { + name: "one workspace", + workspaces: []provisioning.Workspace{{WorkspaceId: 123, WorkspaceName: "Only Workspace"}}, + want: "123", + message: `Auto-selected workspace "Only Workspace" (123)`, + }, + } { + t.Run(tc.name, func(t *testing.T) { + args, requests := newWorkspaceListServer(t, http.StatusOK, tc.workspaces) + var output bytes.Buffer + ctx := cmdio.InContext(t.Context(), cmdio.NewTestIO(strings.NewReader(""), io.Discard, &output)) + id, err := authcmd.PromptForWorkspaceSelection(ctx, args, oauth2.StaticTokenSource(&oauth2.Token{AccessToken: workspaceTestToken})) + require.NoError(t, err) + assert.Equal(t, tc.want, id) + assert.EqualValues(t, 1, requests.Load()) + assert.Contains(t, output.String(), tc.message) + assert.NotContains(t, output.String(), workspacePickerLabel) + assert.NotContains(t, output.String(), workspaceIDPrompt) + }) + } +} + +func TestPromptForWorkspaceSelectionNonInteractive(t *testing.T) { + args, requests := newWorkspaceListServer(t, http.StatusOK, nil) + ctx, output := cmdio.NewTestContextWithStderr(t.Context()) + id, err := authcmd.PromptForWorkspaceSelection(ctx, args, oauth2.StaticTokenSource(&oauth2.Token{AccessToken: workspaceTestToken})) + require.NoError(t, err) + assert.Empty(t, id) + assert.Zero(t, requests.Load()) + assert.Contains(t, output.String(), "set workspace_id in your profile") +} + +func TestPromptForWorkspaceSelectionListError(t *testing.T) { + for _, input := range []string{" 123456 ", ""} { + t.Run(input, func(t *testing.T) { + args, requests := newWorkspaceListServer(t, http.StatusForbidden, nil) + tm := termtest.Run(t, func(ctx context.Context) (string, error) { + return authcmd.PromptForWorkspaceSelection(ctx, args, oauth2.StaticTokenSource(&oauth2.Token{AccessToken: workspaceTestToken})) + }) + tm.WaitFor(workspaceIDPrompt) + tm.Type(input + termtest.KeyEnter) + id, err := tm.Result() + require.NoError(t, err) + assert.Equal(t, strings.TrimSpace(input), id) + assert.EqualValues(t, 1, requests.Load()) + assert.NotContains(t, tm.Raw(), workspacePickerLabel) + }) + } +} + +// newWorkspaceListServer exercises the SDK's actual list decoding without a live account. +func newWorkspaceListServer(t *testing.T, status int, workspaces []provisioning.Workspace) (*auth.AuthArguments, *atomic.Int32) { + t.Helper() + // A developer's DATABRICKS_CONFIG_PROFILE would make the SDK client load + // ~/.databrickscfg and fail on a missing profile. Clearing the whole + // environment instead breaks sockets on Windows, which needs SystemRoot. + t.Setenv("DATABRICKS_CONFIG_PROFILE", "") + body, err := json.Marshal(workspaces) + require.NoError(t, err) + if status != http.StatusOK { + body = []byte(`{"error_code":"PERMISSION_DENIED","message":"Account access denied"}`) + } + requests := new(atomic.Int32) + mux := http.NewServeMux() + mux.HandleFunc("GET /api/2.0/accounts/"+workspaceTestAccountID+"/workspaces", func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + assert.Empty(t, r.URL.RawQuery) + assert.Equal(t, "Bearer "+workspaceTestToken, r.Header.Get("Authorization")) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _, err := w.Write(body) + assert.NoError(t, err) + }) + server := httptest.NewServer(mux) + t.Cleanup(server.Close) + return &auth.AuthArguments{Host: server.URL, AccountID: workspaceTestAccountID}, requests +} diff --git a/cmd/auth/testdata/.gitattributes b/cmd/auth/testdata/.gitattributes new file mode 100644 index 00000000000..c111764a37e --- /dev/null +++ b/cmd/auth/testdata/.gitattributes @@ -0,0 +1,3 @@ +# termtest goldens are LF-only snapshots of the VT emulator; keep Windows +# checkouts from converting them to CRLF. +*.golden text eol=lf diff --git a/cmd/auth/testdata/TestPromptForWorkspaceSelectionNoResults/01-initial.golden b/cmd/auth/testdata/TestPromptForWorkspaceSelectionNoResults/01-initial.golden new file mode 100644 index 00000000000..66cd29acfdb --- /dev/null +++ b/cmd/auth/testdata/TestPromptForWorkspaceSelectionNoResults/01-initial.golden @@ -0,0 +1,6 @@ +Search: █ +? Select a workspace (search by name or ID): + ▸ Alpha (111) + Beta (222) + 333 + Skip (account-level access only) diff --git a/libs/cmdio/cmdiotest/termtest/termtest.go b/libs/cmdio/cmdiotest/termtest/termtest.go index 567a44d30d1..6ef0a71df09 100644 --- a/libs/cmdio/cmdiotest/termtest/termtest.go +++ b/libs/cmdio/cmdiotest/termtest/termtest.go @@ -134,7 +134,7 @@ type Term[T any] struct { // it. func NewPrompt(t *testing.T, opts cmdio.PromptOptions) *Term[string] { t.Helper() - return run(t, func(ctx context.Context) (string, error) { + return Run(t, func(ctx context.Context) (string, error) { return cmdio.RunPrompt(ctx, opts) }) } @@ -142,7 +142,7 @@ func NewPrompt(t *testing.T, opts cmdio.PromptOptions) *Term[string] { // NewSecret runs [cmdio.Secret] in a goroutine and returns a harness for it. func NewSecret(t *testing.T, label string) *Term[string] { t.Helper() - return run(t, func(ctx context.Context) (string, error) { + return Run(t, func(ctx context.Context) (string, error) { return cmdio.Secret(ctx, label) }) } @@ -151,7 +151,7 @@ func NewSecret(t *testing.T, label string) *Term[string] { // it. func NewSelect(t *testing.T, opts cmdio.SelectOptions) *Term[int] { t.Helper() - return run(t, func(ctx context.Context) (int, error) { + return Run(t, func(ctx context.Context) (int, error) { return cmdio.RunSelect(ctx, opts) }) } @@ -160,12 +160,14 @@ func NewSelect(t *testing.T, opts cmdio.SelectOptions) *Term[int] { // harness for it. func NewSelectOrdered(t *testing.T, items []cmdio.Tuple, label string) *Term[string] { t.Helper() - return run(t, func(ctx context.Context) (string, error) { + return Run(t, func(ctx context.Context) (string, error) { return cmdio.SelectOrdered(ctx, items, label) }) } -func run[T any](t *testing.T, fn func(ctx context.Context) (T, error)) *Term[T] { +// Run drives an interactive command through the same terminal harness as the prompts. +func Run[T any](t *testing.T, fn func(ctx context.Context) (T, error)) *Term[T] { + t.Helper() inR, inW, err := os.Pipe() require.NoError(t, err, "open input pipe") tt := &Term[T]{