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/completion.go b/cmd/lakebox/completion.go new file mode 100644 index 0000000000..c078832b8b --- /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 55cc47f0f6..e6cda5c7bc 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 cd96df172d..e637165e79 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 001a382252..83d8d75189 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 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..616c9ee056 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 708d5e4b9d..005a26d871 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 0000000000..6b8b0010a3 --- /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 10aef7daa6..16a60a140a 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 f40d632729..3c837f9aa3 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)