lakebox: tab completion for sandbox IDs and ssh-key hashes

cmd/lakebox/completion.go

@@ -0,0 +1,63 @@
+package lakebox
+
+import (
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/spf13/cobra"
+)
+
+// 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).
+//
+// 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.
+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 {
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+	ctx := cmd.Context()
+	api, err := newLakeboxAPI(cmdctx.WorkspaceClient(ctx))
+	if err != nil {
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+	entries, err := api.list(ctx)
+	if err != nil {
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+	ids := make([]string, 0, len(entries))
+	for _, e := range entries {
+		ids = append(ids, e.SandboxID)
+	}
+	return ids, cobra.ShellCompDirectiveNoFileComp
+}
+
+// completeSSHKeyHashes is the equivalent for `ssh-key delete <hash>`,
+// returning the hashes of registered keys. Same best-effort semantics.
+func completeSSHKeyHashes(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 {
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+	ctx := cmd.Context()
+	api, err := newLakeboxAPI(cmdctx.WorkspaceClient(ctx))
+	if err != nil {
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+	keys, err := api.listKeys(ctx)
+	if err != nil {
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+	hashes := make([]string, 0, len(keys))
+	for _, k := range keys {
+		hashes = append(hashes, k.KeyHash)
+	}
+	return hashes, cobra.ShellCompDirectiveNoFileComp
+}

cmd/lakebox/config.go

@@ -56,8 +56,9 @@ Examples:
   databricks lakebox config happy-panda-1234 --no-autostop                  # never auto-stop
   databricks lakebox config happy-panda-1234 --no-autostop=false            # back to timeout path
   databricks lakebox config happy-panda-1234 --idle-timeout 30m --no-autostop=false`,
-		Args:    cobra.ExactArgs(1),
-		PreRunE: root.MustWorkspaceClient,
+		Args:              cobra.ExactArgs(1),
+		PreRunE:           root.MustWorkspaceClient,
+		ValidArgsFunction: completeSandboxIDs,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)

cmd/lakebox/default.go

@@ -18,8 +18,9 @@ The default is stored locally in ~/.databricks/lakebox.json per profile.
 
 Example:
   databricks lakebox set-default happy-panda-1234`,
-		Args:    cobra.ExactArgs(1),
-		PreRunE: root.MustWorkspaceClient,
+		Args:              cobra.ExactArgs(1),
+		PreRunE:           root.MustWorkspaceClient,
+		ValidArgsFunction: completeSandboxIDs,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)

cmd/lakebox/delete.go

@@ -19,8 +19,9 @@ Permanently terminates and removes the specified lakebox.
 
 Example:
   databricks lakebox delete happy-panda-1234`,
-		Args:    cobra.ExactArgs(1),
-		PreRunE: root.MustWorkspaceClient,
+		Args:              cobra.ExactArgs(1),
+		PreRunE:           root.MustWorkspaceClient,
+		ValidArgsFunction: completeSandboxIDs,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)

cmd/lakebox/ssh.go

@@ -51,8 +51,9 @@ Examples:
   databricks lakebox ssh -- ls -la /home                  # run command on default lakebox
   databricks lakebox ssh happy-panda-1234 -- cat /etc/os-release  # run command on specific lakebox
   databricks lakebox ssh -- -L 8080:localhost:8080        # port forwarding on default lakebox`,
-		Args:    cobra.ArbitraryArgs,
-		PreRunE: root.MustWorkspaceClient,
+		Args:              cobra.ArbitraryArgs,
+		PreRunE:           root.MustWorkspaceClient,
+		ValidArgsFunction: completeSandboxIDs,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)

cmd/lakebox/sshkey.go

@@ -35,8 +35,9 @@ private key will fail until the key is re-registered.
 
 Example:
   databricks lakebox ssh-key delete a1b2c3d4e5f6...`,
-		Args:    cobra.ExactArgs(1),
-		PreRunE: root.MustWorkspaceClient,
+		Args:              cobra.ExactArgs(1),
+		PreRunE:           root.MustWorkspaceClient,
+		ValidArgsFunction: completeSSHKeyHashes,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)

cmd/lakebox/start.go

@@ -24,8 +24,9 @@ Starting an already-running sandbox is a no-op.
 
 Example:
   databricks lakebox start happy-panda-1234`,
-		Args:    cobra.ExactArgs(1),
-		PreRunE: root.MustWorkspaceClient,
+		Args:              cobra.ExactArgs(1),
+		PreRunE:           root.MustWorkspaceClient,
+		ValidArgsFunction: completeSandboxIDs,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)

cmd/lakebox/status.go

@@ -21,8 +21,9 @@ func newStatusCommand() *cobra.Command {
 Example:
   lakebox status happy-panda-1234
   lakebox status happy-panda-1234 --json`,
-		Args:    cobra.ExactArgs(1),
-		PreRunE: root.MustWorkspaceClient,
+		Args:              cobra.ExactArgs(1),
+		PreRunE:           root.MustWorkspaceClient,
+		ValidArgsFunction: completeSandboxIDs,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)

cmd/lakebox/stop.go

@@ -23,8 +23,9 @@ Stopping an already-stopped sandbox is a no-op.
 
 Example:
   databricks lakebox stop happy-panda-1234`,
-		Args:    cobra.ExactArgs(1),
-		PreRunE: root.MustWorkspaceClient,
+		Args:              cobra.ExactArgs(1),
+		PreRunE:           root.MustWorkspaceClient,
+		ValidArgsFunction: completeSandboxIDs,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)