From 45bb44d7a42a8c0a06bc6f0a1ced3218db4153f0 Mon Sep 17 00:00:00 2001 From: Akshay Singla Date: Fri, 22 May 2026 04:56:00 +0000 Subject: [PATCH] lakebox: add start command, fail-fast on unknown sandbox in ssh MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two small additions surfaced during end-to-end testing of #5292. 1. `databricks lakebox start ` — wraps the existing StartSandbox RPC. Symmetric to `stop`. Useful for pre-warming a sandbox before connecting (the gateway also auto-starts on ssh, so this is mostly for scripting). 2. `databricks lakebox ssh ` now fails with `no lakebox named "..." — `databricks lakebox list` shows available IDs` instead of falling through to ssh and surfacing the generic `Permission denied (publickey)` from the gateway. Implemented by making the existing pre-ssh `api.get` (added in #5292 for gateway discovery) treat `apierr.ErrNotFound` as fatal. Non-NotFound errors still fall through so transient API hiccups don't block a connection the gateway can still route. Co-authored-by: Isaac --- cmd/lakebox/api.go | 12 +++++++++ cmd/lakebox/lakebox.go | 1 + cmd/lakebox/ssh.go | 22 +++++++++++----- cmd/lakebox/start.go | 59 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 88 insertions(+), 6 deletions(-) create mode 100644 cmd/lakebox/start.go diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go index 840a831679..1ba1af3c04 100644 --- a/cmd/lakebox/api.go +++ b/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. diff --git a/cmd/lakebox/lakebox.go b/cmd/lakebox/lakebox.go index b104af080f..33ce4aa1df 100644 --- a/cmd/lakebox/lakebox.go +++ b/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()) diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go index e4d436a0f6..d8d567271c 100644 --- a/cmd/lakebox/ssh.go +++ b/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" ) @@ -131,13 +132,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)) } } diff --git a/cmd/lakebox/start.go b/cmd/lakebox/start.go new file mode 100644 index 0000000000..9005670d34 --- /dev/null +++ b/cmd/lakebox/start.go @@ -0,0 +1,59 @@ +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 ", + 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, + 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 +}