From 393e9a0fe003c994f257d236e307ffc3f78bcc1a Mon Sep 17 00:00:00 2001 From: Akshay Singla Date: Fri, 22 May 2026 04:38:08 +0000 Subject: [PATCH] lakebox: add start command, tab completion, fail-fast on unknown sandbox MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three small additions that close gaps surfaced testing the gateway-host PR end-to-end. 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. Tab completion for sandbox IDs on every command that takes one (ssh / status / stop / start / delete / config / set-default) and key hashes on `ssh-key delete`. Best-effort: completion silently produces no suggestions when the API call fails so the shell stays usable. Standard Cobra ValidArgsFunction pattern, surfaced via `databricks completion bash|zsh|fish` like every other Cobra CLI. 3. `databricks lakebox ssh ` now fails with `no lakebox named "..." — `databricks lakebox list` shows available IDs` instead of falling through to ssh and surfacing `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/completion.go | 63 +++++++++++++++++++++++++++++++++++++++ cmd/lakebox/config.go | 5 ++-- cmd/lakebox/default.go | 5 ++-- cmd/lakebox/delete.go | 5 ++-- cmd/lakebox/lakebox.go | 1 + cmd/lakebox/ssh.go | 26 ++++++++++++---- cmd/lakebox/sshkey.go | 5 ++-- cmd/lakebox/start.go | 60 +++++++++++++++++++++++++++++++++++++ cmd/lakebox/status.go | 5 ++-- cmd/lakebox/stop.go | 5 ++-- 11 files changed, 174 insertions(+), 18 deletions(-) create mode 100644 cmd/lakebox/completion.go create mode 100644 cmd/lakebox/start.go diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go index 840a8316792..1ba1af3c04e 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/completion.go b/cmd/lakebox/completion.go new file mode 100644 index 00000000000..c078832b8b8 --- /dev/null +++ b/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 `, +// 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 +} diff --git a/cmd/lakebox/config.go b/cmd/lakebox/config.go index 55cc47f0f60..e6cda5c7bc4 100644 --- a/cmd/lakebox/config.go +++ b/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) diff --git a/cmd/lakebox/default.go b/cmd/lakebox/default.go index cd96df172d1..e637165e796 100644 --- a/cmd/lakebox/default.go +++ b/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) diff --git a/cmd/lakebox/delete.go b/cmd/lakebox/delete.go index 001a3822524..83d8d751892 100644 --- a/cmd/lakebox/delete.go +++ b/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) diff --git a/cmd/lakebox/lakebox.go b/cmd/lakebox/lakebox.go index b104af080f8..33ce4aa1df0 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 e4d436a0f69..616c9ee0569 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" ) @@ -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)) } } diff --git a/cmd/lakebox/sshkey.go b/cmd/lakebox/sshkey.go index 708d5e4b9d9..005a26d871c 100644 --- a/cmd/lakebox/sshkey.go +++ b/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) diff --git a/cmd/lakebox/start.go b/cmd/lakebox/start.go new file mode 100644 index 00000000000..6b8b0010a3b --- /dev/null +++ b/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 ", + 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 +} diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go index 10aef7daa69..16a60a140a2 100644 --- a/cmd/lakebox/status.go +++ b/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) diff --git a/cmd/lakebox/stop.go b/cmd/lakebox/stop.go index f40d632729a..3c837f9aa33 100644 --- a/cmd/lakebox/stop.go +++ b/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)