diff --git a/cmd/lakebox/completion.go b/cmd/lakebox/completion.go index c078832b8b..f05cc62188 100644 --- a/cmd/lakebox/completion.go +++ b/cmd/lakebox/completion.go @@ -7,38 +7,47 @@ import ( ) // completeSandboxIDs is a Cobra ValidArgsFunction returning the caller's -// sandbox IDs for tab completion. Cobra runs this in a separate process -// from the main command, so we need to bootstrap the workspace client -// ourselves (PreRunE is skipped during completion). +// sandbox IDs (and display names, when distinct from the ID) for tab +// completion. Reads purely from the local cache populated by `lakebox +// list` / `create` / `status` / etc. — no API call on ``, so a flaky +// network or unrefreshed auth token never makes the shell hang. // -// Best-effort: any failure (no auth, no network, lakebox not deployed) -// returns no suggestions instead of an error so the shell stays usable -// and the user can still type the ID by hand. +// Cobra runs ValidArgsFunction in a separate process from the main +// command, so we still need to bootstrap the workspace client just to +// learn which profile we're under. The cache itself is profile-scoped. +// +// Best-effort: any failure (no profile resolvable, empty cache) returns +// no suggestions instead of an error so the shell stays usable and the +// user can still type the ID by hand. func completeSandboxIDs(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) { if len(args) > 0 { return nil, cobra.ShellCompDirectiveNoFileComp } - if err := root.MustWorkspaceClient(cmd, args); err != nil { + profile := completionProfile(cmd, args) + if profile == "" { return nil, cobra.ShellCompDirectiveNoFileComp } - ctx := cmd.Context() - api, err := newLakeboxAPI(cmdctx.WorkspaceClient(ctx)) - if err != nil { + sbs := getSandboxes(cmd.Context(), profile) + if len(sbs) == 0 { return nil, cobra.ShellCompDirectiveNoFileComp } - entries, err := api.list(ctx) - if err != nil { - return nil, cobra.ShellCompDirectiveNoFileComp + // Offer the ID always, plus the display name when the user actually + // set one (server defaults `name` to the sandbox ID, so don't echo + // the same string twice). + suggestions := make([]string, 0, len(sbs)*2) + for _, s := range sbs { + suggestions = append(suggestions, s.ID) + if s.Name != "" && s.Name != s.ID { + suggestions = append(suggestions, s.Name) + } } - ids := make([]string, 0, len(entries)) - for _, e := range entries { - ids = append(ids, e.SandboxID) - } - return ids, cobra.ShellCompDirectiveNoFileComp + return suggestions, cobra.ShellCompDirectiveNoFileComp } // completeSSHKeyHashes is the equivalent for `ssh-key delete `, -// returning the hashes of registered keys. Same best-effort semantics. +// returning the hashes of registered keys. SSH-key hashes aren't cached +// locally (per-user cap is ~100 server-side and listing is cheap), so +// this path still calls the API. func completeSSHKeyHashes(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) { if len(args) > 0 { return nil, cobra.ShellCompDirectiveNoFileComp @@ -61,3 +70,17 @@ func completeSSHKeyHashes(cmd *cobra.Command, args []string, toComplete string) } return hashes, cobra.ShellCompDirectiveNoFileComp } + +// completionProfile returns the profile name the cache is keyed under, +// matching the resolution used by the runtime commands (Profile if set, +// else Host). Returns "" if the workspace client can't be bootstrapped. +func completionProfile(cmd *cobra.Command, args []string) string { + if err := root.MustWorkspaceClient(cmd, args); err != nil { + return "" + } + w := cmdctx.WorkspaceClient(cmd.Context()) + if w.Config.Profile != "" { + return w.Config.Profile + } + return w.Config.Host +} diff --git a/cmd/lakebox/config.go b/cmd/lakebox/config.go index 6b056abdf1..0d63edcf4b 100644 --- a/cmd/lakebox/config.go +++ b/cmd/lakebox/config.go @@ -68,7 +68,15 @@ Examples: } out := cmd.OutOrStdout() - id := args[0] + profile := w.Config.Profile + if profile == "" { + profile = w.Config.Host + } + + id, err := resolveLocalID(ctx, profile, args[0]) + if err != nil { + return err + } // Translate flag presence + value into the proto3 // optional-field semantics the server expects. @@ -101,12 +109,8 @@ Examples: if err != nil { return fmt.Errorf("failed to update lakebox %s: %w", id, err) } - - profile := w.Config.Profile - if profile == "" { - profile = w.Config.Host - } _ = setGatewayHost(ctx, profile, updated.GatewayHost) + _ = upsertSandbox(ctx, profile, updated.SandboxID, updated.Name) blank(out) field(ctx, out, "id", cmdio.Bold(ctx, updated.SandboxID)) diff --git a/cmd/lakebox/create.go b/cmd/lakebox/create.go index 134a816c81..5b6a795b43 100644 --- a/cmd/lakebox/create.go +++ b/cmd/lakebox/create.go @@ -60,6 +60,7 @@ Examples: profile = w.Config.Host } _ = setGatewayHost(ctx, profile, result.GatewayHost) + _ = upsertSandbox(ctx, profile, result.SandboxID, name) currentDefault := getDefault(ctx, profile) shouldSetDefault := currentDefault == "" diff --git a/cmd/lakebox/default.go b/cmd/lakebox/default.go index 8dfec82855..42f58b466a 100644 --- a/cmd/lakebox/default.go +++ b/cmd/lakebox/default.go @@ -35,12 +35,16 @@ Example: profile = w.Config.Host } - lakeboxID := args[0] + lakeboxID, err := resolveLocalID(ctx, profile, args[0]) + if err != nil { + return err + } api, err := newLakeboxAPI(w) if err != nil { return err } - if _, err := api.get(ctx, lakeboxID); err != nil { + entry, err := api.get(ctx, lakeboxID) + if err != nil { if errors.Is(err, apierr.ErrNotFound) { return fmt.Errorf("no lakebox named %q — `databricks lakebox list` shows available IDs", lakeboxID) } @@ -50,6 +54,7 @@ Example: if err := setDefault(ctx, profile, lakeboxID); err != nil { return fmt.Errorf("failed to set default: %w", err) } + _ = upsertSandbox(ctx, profile, entry.SandboxID, entry.Name) fmt.Fprintf(cmd.OutOrStdout(), "Default lakebox set to: %s\n", lakeboxID) return nil }, diff --git a/cmd/lakebox/delete.go b/cmd/lakebox/delete.go index 0e75d3bf6c..809dbda5fb 100644 --- a/cmd/lakebox/delete.go +++ b/cmd/lakebox/delete.go @@ -37,7 +37,15 @@ Examples: return err } - lakeboxID := args[0] + profile := w.Config.Profile + if profile == "" { + profile = w.Config.Host + } + + lakeboxID, err := resolveLocalID(ctx, profile, args[0]) + if err != nil { + return err + } // Validate existence first so `delete ` fails clearly // instead of returning a confident "✓ Removed" on a sandbox @@ -82,10 +90,7 @@ Examples: return fmt.Errorf("failed to delete lakebox %s: %w", lakeboxID, err) } - profile := w.Config.Profile - if profile == "" { - profile = w.Config.Host - } + _ = removeSandbox(ctx, profile, lakeboxID) if getDefault(ctx, profile) == lakeboxID { _ = clearDefault(ctx, profile) s.ok("Removed " + cmdio.Bold(ctx, lakeboxID) + " " + cmdio.Faint(ctx, "(default cleared)")) diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go index ab5217d757..4bd8b22deb 100644 --- a/cmd/lakebox/list.go +++ b/cmd/lakebox/list.go @@ -74,6 +74,15 @@ Example: } } + // Replace the local (id, name) cache from this fresh list, + // so subsequent name-based commands (`lakebox ssh my-project`, + // etc.) resolve locally without another API call. + refs := make([]cachedSandbox, 0, len(entries)) + for _, e := range entries { + refs = append(refs, cachedSandbox{ID: e.SandboxID, Name: e.Name}) + } + _ = setSandboxes(ctx, profile, refs) + if outputJSON { enc := json.NewEncoder(cmd.OutOrStdout()) enc.SetIndent("", " ") diff --git a/cmd/lakebox/resolve.go b/cmd/lakebox/resolve.go new file mode 100644 index 0000000000..d48d97fbd4 --- /dev/null +++ b/cmd/lakebox/resolve.go @@ -0,0 +1,64 @@ +package lakebox + +import ( + "context" + "fmt" + "strings" +) + +// resolveLocalID maps a user-typed argument to a sandbox ID using only +// the local state cache — no API calls. Resolution order: +// +// 1. Exact ID match → return arg unchanged. Anyone scripting against +// IDs (or pasting an ID they already know) gets the fast path. +// 2. Exact name match in the cache: +// - One match → return that sandbox's ID. +// - Several matches → error with the candidate IDs, since the +// server allows multiple sandboxes to share a `--name` and the +// CLI must not silently guess which one was meant. +// 3. No match → pass the arg through. It's either an ID the cache +// hasn't seen yet (fresh sandbox created elsewhere; the user just +// hasn't run `lakebox list` since) or a typo. Either way the next +// API call surfaces a clean 404 instead of us refusing locally. +// +// This is intentionally an offline shortcut: a stale cache only fails +// to find a name; it never causes the CLI to act on the wrong sandbox. +// Run `lakebox list` to refresh. +func resolveLocalID(ctx context.Context, profile, arg string) (string, error) { + if arg == "" { + return "", fmt.Errorf("empty lakebox identifier") + } + + sbs := getSandboxes(ctx, profile) + + // ID-first so that a sandbox whose --name happens to collide with + // another sandbox's ID never gets mistakenly resolved by name. + for _, s := range sbs { + if s.ID == arg { + return arg, nil + } + } + + var matches []cachedSandbox + for _, s := range sbs { + if s.Name != "" && s.Name == arg { + matches = append(matches, s) + } + } + + switch len(matches) { + case 0: + return arg, nil + case 1: + return matches[0].ID, nil + default: + ids := make([]string, len(matches)) + for i, m := range matches { + ids[i] = m.ID + } + return "", fmt.Errorf( + "name %q is ambiguous in local cache; sandboxes with this name: %s — pass the ID directly", + arg, strings.Join(ids, ", "), + ) + } +} diff --git a/cmd/lakebox/resolve_test.go b/cmd/lakebox/resolve_test.go new file mode 100644 index 0000000000..761bf9cdb6 --- /dev/null +++ b/cmd/lakebox/resolve_test.go @@ -0,0 +1,128 @@ +package lakebox + +import ( + "os" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestResolveLocalIDEmptyArgIsError(t *testing.T) { + ctx, _ := stateCtx(t) + _, err := resolveLocalID(ctx, "p", "") + require.Error(t, err) +} + +func TestResolveLocalIDPassThroughWhenCacheEmpty(t *testing.T) { + ctx, _ := stateCtx(t) + id, err := resolveLocalID(ctx, "p", "anything-goes") + require.NoError(t, err) + assert.Equal(t, "anything-goes", id) +} + +func TestResolveLocalIDIDFirstPrecedence(t *testing.T) { + // A sandbox whose --name happens to collide with another sandbox's ID + // must NOT resolve the typed ID via the name path. + ctx, _ := stateCtx(t) + require.NoError(t, setSandboxes(ctx, "p", []cachedSandbox{ + {ID: "happy-panda-1234", Name: "real-panda"}, + {ID: "other-id-5678", Name: "happy-panda-1234"}, + })) + id, err := resolveLocalID(ctx, "p", "happy-panda-1234") + require.NoError(t, err) + assert.Equal(t, "happy-panda-1234", id, "must resolve to the ID, not the entry that has this string as its --name") +} + +func TestResolveLocalIDNameMatch(t *testing.T) { + ctx, _ := stateCtx(t) + require.NoError(t, setSandboxes(ctx, "p", []cachedSandbox{ + {ID: "happy-panda-1234", Name: "my-project"}, + })) + id, err := resolveLocalID(ctx, "p", "my-project") + require.NoError(t, err) + assert.Equal(t, "happy-panda-1234", id) +} + +func TestResolveLocalIDAmbiguousNameErrors(t *testing.T) { + ctx, _ := stateCtx(t) + require.NoError(t, setSandboxes(ctx, "p", []cachedSandbox{ + {ID: "happy-panda-1234", Name: "foo"}, + {ID: "sad-otter-5678", Name: "foo"}, + })) + _, err := resolveLocalID(ctx, "p", "foo") + require.Error(t, err) + assert.Contains(t, err.Error(), "ambiguous") + assert.Contains(t, err.Error(), "happy-panda-1234") + assert.Contains(t, err.Error(), "sad-otter-5678") +} + +func TestResolveLocalIDEmptyNamesIgnored(t *testing.T) { + // Sandboxes without a --name should not match an empty-string arg or + // any other arg by virtue of their name being empty. + ctx, _ := stateCtx(t) + require.NoError(t, setSandboxes(ctx, "p", []cachedSandbox{ + {ID: "happy-panda-1234", Name: ""}, + })) + _, err := resolveLocalID(ctx, "p", "") + require.Error(t, err, "empty arg always errors") + + id, err := resolveLocalID(ctx, "p", "happy-panda-1234") + require.NoError(t, err) + assert.Equal(t, "happy-panda-1234", id) +} + +func TestSandboxesPerProfileIsolated(t *testing.T) { + ctx, _ := stateCtx(t) + require.NoError(t, setSandboxes(ctx, "p-a", []cachedSandbox{{ID: "a-1", Name: "shared-name"}})) + require.NoError(t, setSandboxes(ctx, "p-b", []cachedSandbox{{ID: "b-1", Name: "shared-name"}})) + + id, err := resolveLocalID(ctx, "p-a", "shared-name") + require.NoError(t, err) + assert.Equal(t, "a-1", id) + + id, err = resolveLocalID(ctx, "p-b", "shared-name") + require.NoError(t, err) + assert.Equal(t, "b-1", id) +} + +func TestUpsertSandboxAddsNewEntry(t *testing.T) { + ctx, _ := stateCtx(t) + require.NoError(t, upsertSandbox(ctx, "p", "id-1", "name-1")) + assert.Equal(t, []cachedSandbox{{ID: "id-1", Name: "name-1"}}, getSandboxes(ctx, "p")) +} + +func TestUpsertSandboxUpdatesExistingEntry(t *testing.T) { + ctx, _ := stateCtx(t) + require.NoError(t, upsertSandbox(ctx, "p", "id-1", "old")) + require.NoError(t, upsertSandbox(ctx, "p", "id-1", "new")) + assert.Equal(t, []cachedSandbox{{ID: "id-1", Name: "new"}}, getSandboxes(ctx, "p")) +} + +func TestUpsertSandboxSameValueIsNoop(t *testing.T) { + ctx, path := stateCtx(t) + require.NoError(t, upsertSandbox(ctx, "p", "id-1", "name-1")) + before, err := os.Stat(path) + require.NoError(t, err) + require.NoError(t, upsertSandbox(ctx, "p", "id-1", "name-1")) + after, err := os.Stat(path) + require.NoError(t, err) + assert.Equal(t, before.ModTime(), after.ModTime(), "no-op upsert must not rewrite the file") +} + +func TestRemoveSandboxDropsEntry(t *testing.T) { + ctx, _ := stateCtx(t) + require.NoError(t, setSandboxes(ctx, "p", []cachedSandbox{ + {ID: "keep-me", Name: ""}, + {ID: "drop-me", Name: "doomed"}, + })) + require.NoError(t, removeSandbox(ctx, "p", "drop-me")) + assert.Equal(t, []cachedSandbox{{ID: "keep-me", Name: ""}}, getSandboxes(ctx, "p")) +} + +func TestRemoveSandboxMissingIsNoop(t *testing.T) { + ctx, _ := stateCtx(t) + require.NoError(t, setSandboxes(ctx, "p", []cachedSandbox{{ID: "keep-me"}})) + require.NoError(t, removeSandbox(ctx, "p", "nope")) + assert.Equal(t, []cachedSandbox{{ID: "keep-me"}}, getSandboxes(ctx, "p")) +} diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go index e54201ab52..465d59d451 100644 --- a/cmd/lakebox/ssh.go +++ b/cmd/lakebox/ssh.go @@ -89,6 +89,16 @@ Examples: extraArgs = args[dashAt:] } + // Resolve a user-typed name to its ID using the local cache. + // No-op when arg is already an ID or not in cache. + if lakeboxID != "" { + resolved, err := resolveLocalID(ctx, profile, lakeboxID) + if err != nil { + return err + } + lakeboxID = resolved + } + // sandboxGatewayHost captures the gateway hostname from any // Sandbox response we touch in this command, so the resolution // below can prefer it over the cached value. Stays "" when we diff --git a/cmd/lakebox/start.go b/cmd/lakebox/start.go index 6b8b0010a3..5b97145b09 100644 --- a/cmd/lakebox/start.go +++ b/cmd/lakebox/start.go @@ -35,7 +35,16 @@ Example: return err } - lakeboxID := args[0] + profile := w.Config.Profile + if profile == "" { + profile = w.Config.Host + } + + lakeboxID, err := resolveLocalID(ctx, profile, args[0]) + if err != nil { + return err + } + s := spin(ctx, "Starting "+lakeboxID+"…") defer s.Close() @@ -45,11 +54,8 @@ Example: return fmt.Errorf("failed to start lakebox %s: %w", lakeboxID, err) } - profile := w.Config.Profile - if profile == "" { - profile = w.Config.Host - } _ = setGatewayHost(ctx, profile, updated.GatewayHost) + _ = upsertSandbox(ctx, profile, updated.SandboxID, updated.Name) s.ok("Started " + cmdio.Bold(ctx, updated.SandboxID)) return nil diff --git a/cmd/lakebox/state.go b/cmd/lakebox/state.go index 17d098bcc7..7e38868437 100644 --- a/cmd/lakebox/state.go +++ b/cmd/lakebox/state.go @@ -22,6 +22,21 @@ type stateFile struct { // the sandbox just to learn where to connect. Empty until the first // command that reads a sandbox response populates it. GatewayHosts map[string]string `json:"gatewayHosts,omitempty"` + // Profile name → cached (id, name) pairs, used purely for local + // name-based resolution (`lakebox ssh my-project` finds the matching + // sandbox ID without an extra API call). Refreshed in full by + // `lakebox list`; mutated in-place by `create`, `config --name`, + // `delete`, and `status`. Cache misses fall through to treating the + // argument as an ID, so this never *blocks* an operation — it only + // adds the name-to-ID shortcut. + Sandboxes map[string][]cachedSandbox `json:"sandboxes,omitempty"` +} + +// cachedSandbox is the minimal (id, name) pair we need to resolve a +// user-typed name to a wire ID without calling the server. +type cachedSandbox struct { + ID string `json:"id"` + Name string `json:"name,omitempty"` } func stateFilePath(ctx context.Context) (string, error) { @@ -135,3 +150,100 @@ func setGatewayHost(ctx context.Context, profile, host string) error { state.GatewayHosts[profile] = host return saveState(ctx, state) } + +// getSandboxes returns the locally-cached (id, name) pairs for `profile`, +// or nil if nothing has been cached yet. +func getSandboxes(ctx context.Context, profile string) []cachedSandbox { + state, err := loadState(ctx) + if err != nil { + return nil + } + return state.Sandboxes[profile] +} + +// setSandboxes replaces the cached sandbox list for `profile`. Used by +// `lakebox list` to keep the cache in sync with the server's view. +// No-op when the new list is byte-for-byte identical to the cached one, +// to avoid rewriting the state file on every `list` call. +func setSandboxes(ctx context.Context, profile string, sbs []cachedSandbox) error { + state, err := loadState(ctx) + if err != nil { + return err + } + if sandboxesEqual(state.Sandboxes[profile], sbs) { + return nil + } + if state.Sandboxes == nil { + state.Sandboxes = make(map[string][]cachedSandbox) + } + if sbs == nil { + sbs = []cachedSandbox{} + } + state.Sandboxes[profile] = sbs + return saveState(ctx, state) +} + +// upsertSandbox adds or updates a single cached (id, name) entry for +// `profile`. Use this when a single command observes one sandbox's +// state (create, status, config), so the cache stays warm without +// requiring a full `list`. +func upsertSandbox(ctx context.Context, profile, id, name string) error { + if id == "" { + return nil + } + state, err := loadState(ctx) + if err != nil { + return err + } + existing := state.Sandboxes[profile] + for i, s := range existing { + if s.ID == id { + if s.Name == name { + return nil // no change + } + existing[i].Name = name + if state.Sandboxes == nil { + state.Sandboxes = make(map[string][]cachedSandbox) + } + state.Sandboxes[profile] = existing + return saveState(ctx, state) + } + } + if state.Sandboxes == nil { + state.Sandboxes = make(map[string][]cachedSandbox) + } + state.Sandboxes[profile] = append(existing, cachedSandbox{ID: id, Name: name}) + return saveState(ctx, state) +} + +// removeSandbox drops a single cached entry for `profile`. Called from +// `delete` so the cache doesn't keep referencing sandboxes that no +// longer exist server-side. +func removeSandbox(ctx context.Context, profile, id string) error { + state, err := loadState(ctx) + if err != nil { + return err + } + existing := state.Sandboxes[profile] + for i, s := range existing { + if s.ID == id { + state.Sandboxes[profile] = append(existing[:i], existing[i+1:]...) + return saveState(ctx, state) + } + } + return nil +} + +// sandboxesEqual reports whether two slices contain the same entries in +// the same order. Used by setSandboxes to short-circuit no-op writes. +func sandboxesEqual(a, b []cachedSandbox) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go index 362f4fdce5..924a69f6b3 100644 --- a/cmd/lakebox/status.go +++ b/cmd/lakebox/status.go @@ -32,18 +32,23 @@ Example: return err } - lakeboxID := args[0] + profile := w.Config.Profile + if profile == "" { + profile = w.Config.Host + } + + lakeboxID, err := resolveLocalID(ctx, profile, args[0]) + if err != nil { + return err + } entry, err := api.get(ctx, lakeboxID) if err != nil { return fmt.Errorf("failed to get lakebox %s: %w", lakeboxID, err) } - profile := w.Config.Profile - if profile == "" { - profile = w.Config.Host - } _ = setGatewayHost(ctx, profile, entry.GatewayHost) + _ = upsertSandbox(ctx, profile, entry.SandboxID, entry.Name) if outputJSON { enc := json.NewEncoder(cmd.OutOrStdout()) diff --git a/cmd/lakebox/stop.go b/cmd/lakebox/stop.go index 3c837f9aa3..c83212166b 100644 --- a/cmd/lakebox/stop.go +++ b/cmd/lakebox/stop.go @@ -34,7 +34,16 @@ Example: return err } - lakeboxID := args[0] + profile := w.Config.Profile + if profile == "" { + profile = w.Config.Host + } + + lakeboxID, err := resolveLocalID(ctx, profile, args[0]) + if err != nil { + return err + } + s := spin(ctx, "Stopping "+lakeboxID+"…") defer s.Close() @@ -44,11 +53,8 @@ Example: return fmt.Errorf("failed to stop lakebox %s: %w", lakeboxID, err) } - profile := w.Config.Profile - if profile == "" { - profile = w.Config.Host - } _ = setGatewayHost(ctx, profile, updated.GatewayHost) + _ = upsertSandbox(ctx, profile, updated.SandboxID, updated.Name) s.ok("Stopped " + cmdio.Bold(ctx, updated.SandboxID)) return nil