lakebox: add start command, tab completion, fail-fast on unknown sandbox

cmd/lakebox/api.go

@@ -305,6 +305,18 @@ func (a *lakeboxAPI) stop(ctx context.Context, id string) (*sandboxEntry, error)
 	return &resp, nil
 }
 
+// start calls POST /api/2.0/lakebox/sandboxes/{id}/start and returns the
+// refreshed sandbox. Mirror of `stop`; same body shape per `body: "*"`.
+func (a *lakeboxAPI) start(ctx context.Context, id string) (*sandboxEntry, error) {
+	body := map[string]string{"sandbox_id": id}
+	var resp sandboxEntry
+	err := a.c.Do(ctx, http.MethodPost, lakeboxAPIPath+"/"+id+"/start", a.headers(), nil, body, &resp)
+	if err != nil {
+		return nil, err
+	}
+	return &resp, nil
+}
+
 // registerKey calls POST /api/2.0/lakebox/ssh-keys. An empty `name` is
 // omitted from the wire payload so the server records "unset" rather than
 // an explicit empty string.

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/lakebox.go

@@ -39,6 +39,7 @@ Common workflows:
 	cmd.AddCommand(newCreateCommand())
 	cmd.AddCommand(newDeleteCommand())
 	cmd.AddCommand(newStopCommand())
+	cmd.AddCommand(newStartCommand())
 	cmd.AddCommand(newStatusCommand())
 	cmd.AddCommand(newConfigCommand())
 

cmd/lakebox/ssh.go

@@ -11,6 +11,7 @@ import (
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/execv"
+	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/spf13/cobra"
 )
 
@@ -52,6 +53,10 @@ Examples:
   databricks lakebox ssh -- -L 8080:localhost:8080        # port forwarding on default lakebox`,
 		Args:    cobra.ArbitraryArgs,
 		PreRunE: root.MustWorkspaceClient,
+		// Tab-complete the optional first positional only. Cobra strips
+		// anything after `--` before reaching us, so `len(args) > 0`
+		// suffices to detect "user is past the first positional."
+		ValidArgsFunction: completeSandboxIDs,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
@@ -131,13 +136,22 @@ Examples:
 						warn(ctx, fmt.Sprintf("Could not save default: %v", err))
 					}
 				}
-			} else if getGatewayHost(ctx, profile) == "" {
-				// Explicit-id ssh on a profile we have no cached gateway for:
-				// one-time `get` to learn it. Subsequent invocations hit the
-				// cache and skip the round-trip. Failure here is non-fatal —
-				// we fall through to the workspace-host heuristic.
-				if sb, err := api.get(ctx, lakeboxID); err == nil {
+			} else {
+				// Validate the explicit ID against the server. Two reasons:
+				//   1. Surface `lakebox ssh fake-id` as a clear 404 instead of
+				//      letting the user wade through `Permission denied` from
+				//      ssh when the gateway can't route an unknown sandbox.
+				//   2. Capture `gateway_host` to drive the resolution below.
+				// Non-NotFound errors fall through so transient API hiccups
+				// don't block a connection the gateway can still route.
+				sb, err := api.get(ctx, lakeboxID)
+				switch {
+				case err == nil:
 					sandboxGatewayHost = sb.GatewayHost
+				case errors.Is(err, apierr.ErrNotFound):
+					return fmt.Errorf("no lakebox named %q — `databricks lakebox list` shows available IDs", lakeboxID)
+				default:
+					warn(ctx, fmt.Sprintf("could not validate lakebox %s: %v", lakeboxID, err))
 				}
 			}
 

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

@@ -0,0 +1,60 @@
+package lakebox
+
+import (
+	"fmt"
+
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/cli/libs/cmdio"
+	"github.com/spf13/cobra"
+)
+
+func newStartCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "start <lakebox-id>",
+		Short: "Start a stopped Lakebox environment",
+		Long: `Start a stopped Lakebox environment.
+
+Boots the backing microVM and brings the sandbox to Running.
+'databricks lakebox ssh' already auto-starts a stopped sandbox on
+connection, so this command is mostly useful for pre-warming an
+environment without immediately connecting.
+
+Starting an already-running sandbox is a no-op.
+
+Example:
+  databricks lakebox start happy-panda-1234`,
+		Args:              cobra.ExactArgs(1),
+		PreRunE:           root.MustWorkspaceClient,
+		ValidArgsFunction: completeSandboxIDs,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+			w := cmdctx.WorkspaceClient(ctx)
+			api, err := newLakeboxAPI(w)
+			if err != nil {
+				return err
+			}
+
+			lakeboxID := args[0]
+			s := spin(ctx, "Starting "+lakeboxID+"…")
+			defer s.Close()
+
+			updated, err := api.start(ctx, lakeboxID)
+			if err != nil {
+				s.fail("Failed to start " + lakeboxID)
+				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)
+
+			s.ok("Started " + cmdio.Bold(ctx, updated.SandboxID))
+			return nil
+		},
+	}
+
+	return cmd
+}

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)