lakebox: resolve names → IDs from local cache

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 `<TAB>`, 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 <hash>`,
-// 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
+}

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

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 == ""

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

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 <typo>` 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)"))

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("", "  ")

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, ", "),
+		)
+	}
+}

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