Remove auth workspace picker limit and support ID search

.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))

cmd/auth/export_test.go

@@ -0,0 +1,6 @@
+package auth
+
+var (
+	NewWorkspaceSearcher        = newWorkspaceSearcher
+	PromptForWorkspaceSelection = promptForWorkspaceSelection
+)

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.

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
+}

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

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)

libs/cmdio/cmdiotest/termtest/termtest.go

@@ -134,15 +134,15 @@ 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)
 	})
 }
 
 // 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]{