From 756dfa2ebf729269206662d60e64370463bc1945 Mon Sep 17 00:00:00 2001 From: shuochen0311 Date: Fri, 10 Apr 2026 18:20:26 +0000 Subject: [PATCH 1/7] Add Lakebox CLI for managing Databricks sandbox environments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Lakebox provides SSH-accessible development environments backed by microVM isolation. This adds CLI commands for lifecycle management: - `lakebox auth login` — authenticate to a Databricks workspace - `lakebox create` — create a new lakebox (with optional SSH public key) - `lakebox list` — list your lakeboxes (shows status, key hash, default) - `lakebox ssh` — SSH to your default lakebox (or create one on first use) - `lakebox status ` — show lakebox details - `lakebox delete ` — delete a lakebox - `lakebox set-default ` — change the default lakebox Features: - Default lakebox management stored at ~/.databricks/lakebox.json per profile - Automatic SSH config management (~/.ssh/config) - Public key auth only (password/keyboard-interactive disabled in SSH config) - Creates and sets default on first `lakebox ssh` if none exists --- cmd/cmd.go | 126 +++++---------------- cmd/lakebox/api.go | 175 +++++++++++++++++++++++++++++ cmd/lakebox/create.go | 83 ++++++++++++++ cmd/lakebox/default.go | 39 +++++++ cmd/lakebox/delete.go | 51 +++++++++ cmd/lakebox/exec_unix.go | 13 +++ cmd/lakebox/lakebox.go | 40 +++++++ cmd/lakebox/list.go | 70 ++++++++++++ cmd/lakebox/ssh.go | 235 +++++++++++++++++++++++++++++++++++++++ cmd/lakebox/state.go | 90 +++++++++++++++ cmd/lakebox/status.go | 58 ++++++++++ 11 files changed, 880 insertions(+), 100 deletions(-) create mode 100644 cmd/lakebox/api.go create mode 100644 cmd/lakebox/create.go create mode 100644 cmd/lakebox/default.go create mode 100644 cmd/lakebox/delete.go create mode 100644 cmd/lakebox/exec_unix.go create mode 100644 cmd/lakebox/lakebox.go create mode 100644 cmd/lakebox/list.go create mode 100644 cmd/lakebox/ssh.go create mode 100644 cmd/lakebox/state.go create mode 100644 cmd/lakebox/status.go diff --git a/cmd/cmd.go b/cmd/cmd.go index 014471f763..fe81149c08 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -2,117 +2,43 @@ package cmd import ( "context" - "strings" - "github.com/databricks/cli/cmd/psql" - ssh "github.com/databricks/cli/experimental/ssh/cmd" - - "github.com/databricks/cli/cmd/account" - "github.com/databricks/cli/cmd/api" "github.com/databricks/cli/cmd/auth" - "github.com/databricks/cli/cmd/bundle" - "github.com/databricks/cli/cmd/cache" - "github.com/databricks/cli/cmd/completion" - "github.com/databricks/cli/cmd/configure" - "github.com/databricks/cli/cmd/experimental" - "github.com/databricks/cli/cmd/fs" - "github.com/databricks/cli/cmd/labs" - "github.com/databricks/cli/cmd/pipelines" + "github.com/databricks/cli/cmd/lakebox" "github.com/databricks/cli/cmd/root" - "github.com/databricks/cli/cmd/selftest" - "github.com/databricks/cli/cmd/sync" - "github.com/databricks/cli/cmd/version" - "github.com/databricks/cli/cmd/workspace" - "github.com/databricks/cli/libs/cmdgroup" "github.com/spf13/cobra" ) -const ( - mainGroup = "main" - permissionsGroup = "permissions" -) - -// configureGroups adds groups to the command, only if a group -// has at least one available command. -func configureGroups(cmd *cobra.Command, groups []cobra.Group) { - filteredGroups := cmdgroup.FilterGroups(groups, cmd.Commands()) - for i := range filteredGroups { - cmd.AddGroup(&filteredGroups[i]) - } -} - -func accountCommand() *cobra.Command { - cmd := account.New() - configureGroups(cmd, account.Groups()) - return cmd -} - func New(ctx context.Context) *cobra.Command { cli := root.New(ctx) + cli.Use = "lakebox" + cli.Short = "Lakebox CLI — manage Databricks sandbox environments" + cli.Long = `Lakebox CLI — manage Databricks sandbox environments. + +Lakebox provides SSH-accessible development environments backed by +microVM isolation. Each lakebox is a personal sandbox with pre-installed +tooling (Python, Node.js, Rust, Databricks CLI) and persistent storage. + +Common workflows: + lakebox auth login # authenticate to Databricks + lakebox ssh # SSH to your default lakebox + lakebox ssh my-project # SSH to a named lakebox + lakebox list # list your lakeboxes + lakebox create # create a new lakebox + lakebox delete my-project # delete a lakebox + lakebox status my-project # show lakebox status + +The CLI manages your ~/.ssh/config so you can also connect directly: + ssh my-project # after 'lakebox ssh' +` + cli.CompletionOptions.DisableDefaultCmd = true - // Add account subcommand. - cli.AddCommand(accountCommand()) - - // Add workspace subcommands. - workspaceCommands := workspace.All() - for _, cmd := range workspaceCommands { - // Order the permissions subcommands after the main commands. - for _, sub := range cmd.Commands() { - // some commands override groups in overrides.go, leave them as-is - if sub.GroupID != "" { - continue - } - - switch { - case strings.HasSuffix(sub.Name(), "-permissions"), strings.HasSuffix(sub.Name(), "-permission-levels"): - sub.GroupID = permissionsGroup - default: - sub.GroupID = mainGroup - } - } - - cli.AddCommand(cmd) - - // Built-in groups for the workspace commands. - groups := []cobra.Group{ - { - ID: mainGroup, - Title: "Available Commands", - }, - { - ID: pipelines.ManagementGroupID, - Title: "Management Commands", - }, - { - ID: permissionsGroup, - Title: "Permission Commands", - }, - } - - configureGroups(cmd, groups) - } - - // Add other subcommands. - cli.AddCommand(api.New()) cli.AddCommand(auth.New()) - cli.AddCommand(completion.New()) - cli.AddCommand(bundle.New()) - cli.AddCommand(cache.New()) - cli.AddCommand(experimental.New()) - cli.AddCommand(psql.New()) - cli.AddCommand(configure.New()) - cli.AddCommand(fs.New()) - cli.AddCommand(labs.New(ctx)) - cli.AddCommand(sync.New()) - cli.AddCommand(version.New()) - cli.AddCommand(selftest.New()) - cli.AddCommand(ssh.New()) - // Add workspace command groups, filtering out empty groups or groups with only hidden commands. - configureGroups(cli, append(workspace.Groups(), cobra.Group{ - ID: "development", - Title: "Developer Tools", - })) + // Register lakebox subcommands directly at root level. + for _, sub := range lakebox.New().Commands() { + cli.AddCommand(sub) + } return cli } diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go new file mode 100644 index 0000000000..ff8f7d30b1 --- /dev/null +++ b/cmd/lakebox/api.go @@ -0,0 +1,175 @@ +package lakebox + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + + "github.com/databricks/databricks-sdk-go" +) + +const lakeboxAPIPath = "/api/2.0/lakebox" + +// lakeboxAPI wraps raw HTTP calls to the lakebox REST API. +type lakeboxAPI struct { + w *databricks.WorkspaceClient +} + +// createRequest is the JSON body for POST /api/2.0/lakebox. +type createRequest struct { + PublicKey string `json:"public_key,omitempty"` +} + +// createResponse is the JSON body returned by POST /api/2.0/lakebox. +type createResponse struct { + LakeboxID string `json:"lakebox_id"` + Status string `json:"status"` +} + +// lakeboxEntry is a single item in the list response. +type lakeboxEntry struct { + Name string `json:"name"` + Status string `json:"status"` + FQDN string `json:"fqdn"` + PubkeyHashPrefix string `json:"pubkey_hash_prefix,omitempty"` +} + +// listResponse is the JSON body returned by GET /api/2.0/lakebox. +type listResponse struct { + Lakeboxes []lakeboxEntry `json:"lakeboxes"` +} + +// apiError is the error body returned by the lakebox API. +type apiError struct { + ErrorCode string `json:"error_code"` + Message string `json:"message"` +} + +func (e *apiError) Error() string { + return fmt.Sprintf("%s: %s", e.ErrorCode, e.Message) +} + +func newLakeboxAPI(w *databricks.WorkspaceClient) *lakeboxAPI { + return &lakeboxAPI{w: w} +} + +// create calls POST /api/2.0/lakebox with an optional public key. +func (a *lakeboxAPI) create(ctx context.Context, publicKey string) (*createResponse, error) { + body := createRequest{PublicKey: publicKey} + jsonBody, err := json.Marshal(body) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + resp, err := a.doRequest(ctx, "POST", lakeboxAPIPath, bytes.NewReader(jsonBody)) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusCreated { + return nil, parseAPIError(resp) + } + + var result createResponse + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, fmt.Errorf("failed to decode response: %w", err) + } + return &result, nil +} + +// list calls GET /api/2.0/lakebox. +func (a *lakeboxAPI) list(ctx context.Context) ([]lakeboxEntry, error) { + resp, err := a.doRequest(ctx, "GET", lakeboxAPIPath, nil) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, parseAPIError(resp) + } + + var result listResponse + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, fmt.Errorf("failed to decode response: %w", err) + } + return result.Lakeboxes, nil +} + +// get calls GET /api/2.0/lakebox/{id}. +func (a *lakeboxAPI) get(ctx context.Context, id string) (*lakeboxEntry, error) { + resp, err := a.doRequest(ctx, "GET", lakeboxAPIPath+"/"+id, nil) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, parseAPIError(resp) + } + + var result lakeboxEntry + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, fmt.Errorf("failed to decode response: %w", err) + } + return &result, nil +} + +// delete calls DELETE /api/2.0/lakebox/{id}. +func (a *lakeboxAPI) delete(ctx context.Context, id string) error { + resp, err := a.doRequest(ctx, "DELETE", lakeboxAPIPath+"/"+id, nil) + if err != nil { + return err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusNoContent { + return parseAPIError(resp) + } + return nil +} + +// doRequest makes an authenticated HTTP request to the workspace. +func (a *lakeboxAPI) doRequest(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) { + host := strings.TrimRight(a.w.Config.Host, "/") + url := host + path + + req, err := http.NewRequestWithContext(ctx, method, url, body) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + if err := a.w.Config.Authenticate(req); err != nil { + return nil, fmt.Errorf("failed to authenticate: %w", err) + } + + if body != nil { + req.Header.Set("Content-Type", "application/json") + } + + return http.DefaultClient.Do(req) +} + +func parseAPIError(resp *http.Response) error { + body, _ := io.ReadAll(resp.Body) + var apiErr apiError + if json.Unmarshal(body, &apiErr) == nil && apiErr.Message != "" { + return &apiErr + } + return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body)) +} + +// extractLakeboxID extracts the short ID from a full resource name. +// e.g. "apps/lakebox/instances/happy-panda-1234" -> "happy-panda-1234" +func extractLakeboxID(name string) string { + parts := strings.Split(name, "/") + if len(parts) > 0 { + return parts[len(parts)-1] + } + return name +} diff --git a/cmd/lakebox/create.go b/cmd/lakebox/create.go new file mode 100644 index 0000000000..872776cc8d --- /dev/null +++ b/cmd/lakebox/create.go @@ -0,0 +1,83 @@ +package lakebox + +import ( + "fmt" + "os" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/spf13/cobra" +) + +func newCreateCommand() *cobra.Command { + var publicKeyFile string + + cmd := &cobra.Command{ + Use: "create", + Short: "Create a new Lakebox environment", + Long: `Create a new Lakebox environment. + +Creates a new personal development environment backed by a microVM. +Blocks until the lakebox is running and prints the lakebox ID. + +If --public-key-file is provided, the key is installed in the lakebox's +authorized_keys so you can SSH directly. Otherwise the gateway key is used. + +Example: + databricks lakebox create + databricks lakebox create --public-key-file ~/.ssh/id_ed25519.pub`, + PreRunE: root.MustWorkspaceClient, + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + api := newLakeboxAPI(w) + + var publicKey string + if publicKeyFile != "" { + data, err := os.ReadFile(publicKeyFile) + if err != nil { + return fmt.Errorf("failed to read public key file %s: %w", publicKeyFile, err) + } + publicKey = string(data) + } + + fmt.Fprintf(cmd.ErrOrStderr(), "Creating lakebox...\n") + + result, err := api.create(ctx, publicKey) + if err != nil { + return fmt.Errorf("failed to create lakebox: %w", err) + } + + profile := w.Config.Profile + if profile == "" { + profile = w.Config.Host + } + + // Set as default if no default exists, or the current default + // has been deleted (no longer in the list). + currentDefault := getDefault(profile) + shouldSetDefault := currentDefault == "" + if !shouldSetDefault && currentDefault != "" { + // Check if the current default still exists. + if _, err := api.get(ctx, currentDefault); err != nil { + shouldSetDefault = true + } + } + if shouldSetDefault { + if err := setDefault(profile, result.LakeboxID); err != nil { + fmt.Fprintf(cmd.ErrOrStderr(), "Warning: failed to save default: %v\n", err) + } else { + fmt.Fprintf(cmd.ErrOrStderr(), "Set as default lakebox.\n") + } + } + + fmt.Fprintf(cmd.ErrOrStderr(), "Lakebox created (status: %s)\n", result.Status) + fmt.Fprintln(cmd.OutOrStdout(), result.LakeboxID) + return nil + }, + } + + cmd.Flags().StringVar(&publicKeyFile, "public-key-file", "", "Path to SSH public key file to install in the lakebox") + + return cmd +} diff --git a/cmd/lakebox/default.go b/cmd/lakebox/default.go new file mode 100644 index 0000000000..9d5a366c9c --- /dev/null +++ b/cmd/lakebox/default.go @@ -0,0 +1,39 @@ +package lakebox + +import ( + "fmt" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/spf13/cobra" +) + +func newSetDefaultCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "set-default ", + Short: "Set the default Lakebox for SSH", + Long: `Set the default Lakebox that 'databricks lakebox ssh' connects to. + +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, + RunE: func(cmd *cobra.Command, args []string) error { + w := cmdctx.WorkspaceClient(cmd.Context()) + profile := w.Config.Profile + if profile == "" { + profile = w.Config.Host + } + + lakeboxID := args[0] + if err := setDefault(profile, lakeboxID); err != nil { + return fmt.Errorf("failed to set default: %w", err) + } + fmt.Fprintf(cmd.OutOrStdout(), "Default lakebox set to: %s\n", lakeboxID) + return nil + }, + } + return cmd +} diff --git a/cmd/lakebox/delete.go b/cmd/lakebox/delete.go new file mode 100644 index 0000000000..a814083ed3 --- /dev/null +++ b/cmd/lakebox/delete.go @@ -0,0 +1,51 @@ +package lakebox + +import ( + "fmt" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/spf13/cobra" +) + +func newDeleteCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "delete ", + Short: "Delete a Lakebox environment", + Long: `Delete a Lakebox environment. + +Permanently terminates and removes the specified lakebox. Only the +creator (same auth token) can delete a lakebox. + +Example: + databricks lakebox delete 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 := newLakeboxAPI(w) + + lakeboxID := args[0] + + if err := api.delete(ctx, lakeboxID); err != nil { + return fmt.Errorf("failed to delete lakebox %s: %w", lakeboxID, err) + } + + // Clear default if we just deleted it. + profile := w.Config.Profile + if profile == "" { + profile = w.Config.Host + } + if getDefault(profile) == lakeboxID { + _ = clearDefault(profile) + fmt.Fprintf(cmd.ErrOrStderr(), "Cleared default lakebox.\n") + } + + fmt.Fprintf(cmd.ErrOrStderr(), "Deleted lakebox %s\n", lakeboxID) + return nil + }, + } + + return cmd +} diff --git a/cmd/lakebox/exec_unix.go b/cmd/lakebox/exec_unix.go new file mode 100644 index 0000000000..d47f629572 --- /dev/null +++ b/cmd/lakebox/exec_unix.go @@ -0,0 +1,13 @@ +//go:build !windows + +package lakebox + +import ( + "os" + "syscall" +) + +// execSyscall replaces the current process with the given command (Unix only). +func execSyscall(path string, args []string) error { + return syscall.Exec(path, args, os.Environ()) +} diff --git a/cmd/lakebox/lakebox.go b/cmd/lakebox/lakebox.go new file mode 100644 index 0000000000..6523debef9 --- /dev/null +++ b/cmd/lakebox/lakebox.go @@ -0,0 +1,40 @@ +package lakebox + +import ( + "github.com/spf13/cobra" +) + +func New() *cobra.Command { + cmd := &cobra.Command{ + Use: "lakebox", + Short: "Manage Databricks Lakebox environments", + Long: `Manage Databricks Lakebox environments. + +Lakebox provides SSH-accessible development environments backed by +microVM isolation. Each lakebox is a personal sandbox with pre-installed +tooling (Python, Node.js, Rust, Databricks CLI) and persistent storage. + +Common workflows: + databricks lakebox login # authenticate to Databricks + databricks lakebox ssh # SSH to your default lakebox + databricks lakebox ssh my-project # SSH to a named lakebox + databricks lakebox list # list your lakeboxes + databricks lakebox create --name my-project # create a new lakebox + databricks lakebox delete my-project # delete a lakebox + databricks lakebox status # show current lakebox status + +The CLI manages your ~/.ssh/config so you can also connect directly: + ssh my-project # after 'lakebox ssh --setup' +`, + } + + cmd.AddCommand(newLoginCommand()) + cmd.AddCommand(newSSHCommand()) + cmd.AddCommand(newListCommand()) + cmd.AddCommand(newCreateCommand()) + cmd.AddCommand(newDeleteCommand()) + cmd.AddCommand(newStatusCommand()) + cmd.AddCommand(newSetDefaultCommand()) + + return cmd +} diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go new file mode 100644 index 0000000000..bf80a9919e --- /dev/null +++ b/cmd/lakebox/list.go @@ -0,0 +1,70 @@ +package lakebox + +import ( + "encoding/json" + "fmt" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/spf13/cobra" +) + +func newListCommand() *cobra.Command { + var outputJSON bool + + cmd := &cobra.Command{ + Use: "list", + Short: "List your Lakebox environments", + Long: `List your Lakebox environments. + +Shows all lakeboxes associated with your account, including their +current status and ID. + +Example: + databricks lakebox list + databricks lakebox list --json`, + PreRunE: root.MustWorkspaceClient, + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + api := newLakeboxAPI(w) + + entries, err := api.list(ctx) + if err != nil { + return fmt.Errorf("failed to list lakeboxes: %w", err) + } + + if outputJSON { + enc := json.NewEncoder(cmd.OutOrStdout()) + enc.SetIndent("", " ") + return enc.Encode(entries) + } + + if len(entries) == 0 { + fmt.Fprintln(cmd.ErrOrStderr(), "No lakeboxes found.") + return nil + } + + profile := w.Config.Profile + if profile == "" { + profile = w.Config.Host + } + defaultID := getDefault(profile) + + fmt.Fprintf(cmd.OutOrStdout(), " %-30s %-12s %-10s %s\n", "ID", "STATUS", "KEY", "DEFAULT") + for _, e := range entries { + id := extractLakeboxID(e.Name) + def := "" + if id == defaultID { + def = "*" + } + fmt.Fprintf(cmd.OutOrStdout(), " %-30s %-12s %-10s %s\n", id, e.Status, e.PubkeyHashPrefix, def) + } + return nil + }, + } + + cmd.Flags().BoolVar(&outputJSON, "json", false, "Output as JSON") + + return cmd +} diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go new file mode 100644 index 0000000000..1978dec684 --- /dev/null +++ b/cmd/lakebox/ssh.go @@ -0,0 +1,235 @@ +package lakebox + +import ( + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/spf13/cobra" +) + +const ( + defaultGatewayHost = "uw2.dbrx.dev" + defaultGatewayPort = "2222" + + // SSH config block markers for idempotent updates. + sshConfigMarkerStart = "# --- Lakebox managed start ---" + sshConfigMarkerEnd = "# --- Lakebox managed end ---" +) + +func newSSHCommand() *cobra.Command { + var gatewayHost string + var gatewayPort string + + cmd := &cobra.Command{ + Use: "ssh [lakebox-id]", + Short: "SSH into a Lakebox environment", + Long: `SSH into a Lakebox environment. + +This command: +1. Authenticates to the Databricks workspace +2. Ensures you have a local SSH key (~/.ssh/id_ed25519) +3. Creates a lakebox if one doesn't exist (installs your public key) +4. Updates ~/.ssh/config with a Host entry for the lakebox +5. Connects via SSH using the lakebox ID as the SSH username + +Without arguments, creates a new lakebox. With a lakebox ID argument, +connects to the specified lakebox. + +Example: + databricks lakebox ssh # create and connect to a new lakebox + databricks lakebox ssh happy-panda-1234 # connect to existing lakebox`, + Args: cobra.MaximumNArgs(1), + PreRunE: func(cmd *cobra.Command, args []string) error { + return root.MustWorkspaceClient(cmd, args) + }, + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + + profile := w.Config.Profile + if profile == "" { + profile = w.Config.Host + } + + // Ensure SSH key exists. + keyPath, err := ensureSSHKey() + if err != nil { + return fmt.Errorf("failed to ensure SSH key: %w", err) + } + fmt.Fprintf(cmd.ErrOrStderr(), "Using SSH key: %s\n", keyPath) + + // Determine lakebox ID: + // 1. Explicit arg → use it + // 2. Local default exists → use it + // 3. Neither → create a new one and set as default + var lakeboxID string + if len(args) > 0 { + lakeboxID = args[0] + } else if def := getDefault(profile); def != "" { + lakeboxID = def + fmt.Fprintf(cmd.ErrOrStderr(), "Using default lakebox: %s\n", lakeboxID) + } else { + api := newLakeboxAPI(w) + pubKeyData, err := os.ReadFile(keyPath + ".pub") + if err != nil { + return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err) + } + + fmt.Fprintf(cmd.ErrOrStderr(), "Creating lakebox...\n") + result, err := api.create(ctx, string(pubKeyData)) + if err != nil { + return fmt.Errorf("failed to create lakebox: %w", err) + } + lakeboxID = result.LakeboxID + fmt.Fprintf(cmd.ErrOrStderr(), "Lakebox %s created (status: %s)\n", lakeboxID, result.Status) + + if err := setDefault(profile, lakeboxID); err != nil { + fmt.Fprintf(cmd.ErrOrStderr(), "Warning: failed to save default: %v\n", err) + } + } + + // Write SSH config entry for this lakebox. + sshConfigPath, err := sshConfigFilePath() + if err != nil { + return err + } + entry := buildSSHConfigEntry(lakeboxID, gatewayHost, gatewayPort, keyPath) + if err := writeSSHConfigEntry(sshConfigPath, lakeboxID, entry); err != nil { + return fmt.Errorf("failed to update SSH config: %w", err) + } + + fmt.Fprintf(cmd.ErrOrStderr(), "Connecting to %s@%s:%s...\n", + lakeboxID, gatewayHost, gatewayPort) + return execSSH(lakeboxID) + }, + } + + cmd.Flags().StringVar(&gatewayHost, "gateway", defaultGatewayHost, "Lakebox gateway hostname") + cmd.Flags().StringVar(&gatewayPort, "port", defaultGatewayPort, "Lakebox gateway SSH port") + + return cmd +} + +// ensureSSHKey checks for an existing SSH key and generates one if missing. +func ensureSSHKey() (string, error) { + homeDir, err := os.UserHomeDir() + if err != nil { + return "", err + } + + candidates := []string{ + filepath.Join(homeDir, ".ssh", "id_ed25519"), + filepath.Join(homeDir, ".ssh", "id_rsa"), + } + for _, p := range candidates { + if _, err := os.Stat(p); err == nil { + return p, nil + } + } + + // Generate ed25519 key. + keyPath := candidates[0] + sshDir := filepath.Dir(keyPath) + if err := os.MkdirAll(sshDir, 0700); err != nil { + return "", fmt.Errorf("failed to create %s: %w", sshDir, err) + } + + cmd := exec.Command("ssh-keygen", "-t", "ed25519", "-f", keyPath, "-N", "", "-q") + cmd.Stdin = os.Stdin + cmd.Stdout = os.Stderr + cmd.Stderr = os.Stderr + if err := cmd.Run(); err != nil { + return "", fmt.Errorf("ssh-keygen failed: %w", err) + } + + return keyPath, nil +} + +func sshConfigFilePath() (string, error) { + homeDir, err := os.UserHomeDir() + if err != nil { + return "", err + } + return filepath.Join(homeDir, ".ssh", "config"), nil +} + +// buildSSHConfigEntry creates the SSH config block for a lakebox. +// The lakebox ID is used as both the Host alias and the SSH User. +func buildSSHConfigEntry(lakeboxID, host, port, keyPath string) string { + return fmt.Sprintf(`Host %s + HostName %s + Port %s + User %s + IdentityFile %s + IdentitiesOnly yes + PreferredAuthentications publickey + PasswordAuthentication no + KbdInteractiveAuthentication no + StrictHostKeyChecking no + UserKnownHostsFile /dev/null + LogLevel INFO +`, lakeboxID, host, port, lakeboxID, keyPath) +} + +// writeSSHConfigEntry idempotently writes a single lakebox entry to ~/.ssh/config. +// Replaces any existing lakebox block in-place. +func writeSSHConfigEntry(configPath, lakeboxID, entry string) error { + sshDir := filepath.Dir(configPath) + if err := os.MkdirAll(sshDir, 0700); err != nil { + return err + } + + existing, err := os.ReadFile(configPath) + if err != nil && !os.IsNotExist(err) { + return err + } + + wrappedEntry := fmt.Sprintf("%s\n%s%s\n", sshConfigMarkerStart, entry, sshConfigMarkerEnd) + content := string(existing) + + // Remove existing lakebox block if present. + startIdx := strings.Index(content, sshConfigMarkerStart) + if startIdx >= 0 { + endIdx := strings.Index(content[startIdx:], sshConfigMarkerEnd) + if endIdx >= 0 { + endIdx += startIdx + len(sshConfigMarkerEnd) + if endIdx < len(content) && content[endIdx] == '\n' { + endIdx++ + } + content = content[:startIdx] + content[endIdx:] + } + } + + if !strings.HasSuffix(content, "\n") && len(content) > 0 { + content += "\n" + } + content += wrappedEntry + + return os.WriteFile(configPath, []byte(content), 0600) +} + +// execSSH execs into ssh using the lakebox ID as the Host alias. +func execSSH(lakeboxID string) error { + sshPath, err := exec.LookPath("ssh") + if err != nil { + return fmt.Errorf("ssh not found in PATH: %w", err) + } + + args := []string{"ssh", lakeboxID} + + if runtime.GOOS == "windows" { + cmd := exec.Command(sshPath, args[1:]...) + cmd.Stdin = os.Stdin + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + return cmd.Run() + } + + return execSyscall(sshPath, args) +} diff --git a/cmd/lakebox/state.go b/cmd/lakebox/state.go new file mode 100644 index 0000000000..c0c8ad2d84 --- /dev/null +++ b/cmd/lakebox/state.go @@ -0,0 +1,90 @@ +package lakebox + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" +) + +// stateFile stores per-profile lakebox defaults on the local filesystem. +// Located at ~/.databricks/lakebox.json. +type stateFile struct { + // Profile name → default lakebox ID. + Defaults map[string]string `json:"defaults"` +} + +func stateFilePath() (string, error) { + home, err := os.UserHomeDir() + if err != nil { + return "", err + } + return filepath.Join(home, ".databricks", "lakebox.json"), nil +} + +func loadState() (*stateFile, error) { + path, err := stateFilePath() + if err != nil { + return nil, err + } + + data, err := os.ReadFile(path) + if os.IsNotExist(err) { + return &stateFile{Defaults: make(map[string]string)}, nil + } + if err != nil { + return nil, fmt.Errorf("failed to read %s: %w", path, err) + } + + var state stateFile + if err := json.Unmarshal(data, &state); err != nil { + return &stateFile{Defaults: make(map[string]string)}, nil + } + if state.Defaults == nil { + state.Defaults = make(map[string]string) + } + return &state, nil +} + +func saveState(state *stateFile) error { + path, err := stateFilePath() + if err != nil { + return err + } + + if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil { + return err + } + + data, err := json.MarshalIndent(state, "", " ") + if err != nil { + return err + } + return os.WriteFile(path, data, 0600) +} + +func getDefault(profile string) string { + state, err := loadState() + if err != nil { + return "" + } + return state.Defaults[profile] +} + +func setDefault(profile, lakeboxID string) error { + state, err := loadState() + if err != nil { + return err + } + state.Defaults[profile] = lakeboxID + return saveState(state) +} + +func clearDefault(profile string) error { + state, err := loadState() + if err != nil { + return err + } + delete(state.Defaults, profile) + return saveState(state) +} diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go new file mode 100644 index 0000000000..1afd968211 --- /dev/null +++ b/cmd/lakebox/status.go @@ -0,0 +1,58 @@ +package lakebox + +import ( + "encoding/json" + "fmt" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/spf13/cobra" +) + +func newStatusCommand() *cobra.Command { + var outputJSON bool + + cmd := &cobra.Command{ + Use: "status ", + Short: "Show Lakebox environment status", + Long: `Show detailed status of a Lakebox environment. + +Example: + databricks lakebox status happy-panda-1234 + databricks lakebox status happy-panda-1234 --json`, + Args: cobra.ExactArgs(1), + PreRunE: root.MustWorkspaceClient, + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + api := newLakeboxAPI(w) + + lakeboxID := args[0] + + entry, err := api.get(ctx, lakeboxID) + if err != nil { + return fmt.Errorf("failed to get lakebox %s: %w", lakeboxID, err) + } + + if outputJSON { + enc := json.NewEncoder(cmd.OutOrStdout()) + enc.SetIndent("", " ") + return enc.Encode(entry) + } + + fmt.Fprintf(cmd.OutOrStdout(), "ID: %s\n", extractLakeboxID(entry.Name)) + fmt.Fprintf(cmd.OutOrStdout(), "Status: %s\n", entry.Status) + if entry.FQDN != "" { + fmt.Fprintf(cmd.OutOrStdout(), "FQDN: %s\n", entry.FQDN) + } + if entry.PubkeyHashPrefix != "" { + fmt.Fprintf(cmd.OutOrStdout(), "Key: %s\n", entry.PubkeyHashPrefix) + } + return nil + }, + } + + cmd.Flags().BoolVar(&outputJSON, "json", false, "Output as JSON") + + return cmd +} From c20c6dfaa65e5db081292d051fd5f45517ff6c1a Mon Sep 17 00:00:00 2001 From: shuochen0311 Date: Mon, 13 Apr 2026 20:29:55 +0000 Subject: [PATCH 2/7] Remove KEY column from list, add register-key command - Remove PubkeyHashPrefix field from lakeboxEntry (no longer returned by API) - Remove KEY column from list output - Remove Key line from status output - Add register-key subcommand for SSH public key registration Co-authored-by: Isaac --- cmd/lakebox/api.go | 32 ++++++++++++++++++--- cmd/lakebox/lakebox.go | 19 +++++++------ cmd/lakebox/list.go | 4 +-- cmd/lakebox/register_key.go | 55 +++++++++++++++++++++++++++++++++++++ cmd/lakebox/status.go | 3 -- 5 files changed, 95 insertions(+), 18 deletions(-) create mode 100644 cmd/lakebox/register_key.go diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go index ff8f7d30b1..94877b4a42 100644 --- a/cmd/lakebox/api.go +++ b/cmd/lakebox/api.go @@ -32,10 +32,9 @@ type createResponse struct { // lakeboxEntry is a single item in the list response. type lakeboxEntry struct { - Name string `json:"name"` - Status string `json:"status"` - FQDN string `json:"fqdn"` - PubkeyHashPrefix string `json:"pubkey_hash_prefix,omitempty"` + Name string `json:"name"` + Status string `json:"status"` + FQDN string `json:"fqdn"` } // listResponse is the JSON body returned by GET /api/2.0/lakebox. @@ -164,6 +163,31 @@ func parseAPIError(resp *http.Response) error { return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body)) } +// registerKeyRequest is the JSON body for POST /api/2.0/lakebox/register-key. +type registerKeyRequest struct { + PublicKey string `json:"public_key"` +} + +// registerKey calls POST /api/2.0/lakebox/register-key. +func (a *lakeboxAPI) registerKey(ctx context.Context, publicKey string) error { + body := registerKeyRequest{PublicKey: publicKey} + jsonBody, err := json.Marshal(body) + if err != nil { + return fmt.Errorf("failed to marshal request: %w", err) + } + + resp, err := a.doRequest(ctx, "POST", lakeboxAPIPath+"/register-key", bytes.NewReader(jsonBody)) + if err != nil { + return err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return parseAPIError(resp) + } + return nil +} + // extractLakeboxID extracts the short ID from a full resource name. // e.g. "apps/lakebox/instances/happy-panda-1234" -> "happy-panda-1234" func extractLakeboxID(name string) string { diff --git a/cmd/lakebox/lakebox.go b/cmd/lakebox/lakebox.go index 6523debef9..aa9463bca8 100644 --- a/cmd/lakebox/lakebox.go +++ b/cmd/lakebox/lakebox.go @@ -15,26 +15,27 @@ microVM isolation. Each lakebox is a personal sandbox with pre-installed tooling (Python, Node.js, Rust, Databricks CLI) and persistent storage. Common workflows: - databricks lakebox login # authenticate to Databricks - databricks lakebox ssh # SSH to your default lakebox - databricks lakebox ssh my-project # SSH to a named lakebox - databricks lakebox list # list your lakeboxes - databricks lakebox create --name my-project # create a new lakebox - databricks lakebox delete my-project # delete a lakebox - databricks lakebox status # show current lakebox status + lakebox auth login # authenticate to Databricks + lakebox ssh # SSH to your default lakebox + lakebox ssh my-project # SSH to a named lakebox + lakebox list # list your lakeboxes + lakebox create # create a new lakebox + lakebox delete my-project # delete a lakebox + lakebox status my-project # show lakebox status + lakebox register-key --public-key-file ~/.ssh/id_rsa.pub # register SSH key The CLI manages your ~/.ssh/config so you can also connect directly: - ssh my-project # after 'lakebox ssh --setup' + ssh my-project # after 'lakebox ssh' `, } - cmd.AddCommand(newLoginCommand()) cmd.AddCommand(newSSHCommand()) cmd.AddCommand(newListCommand()) cmd.AddCommand(newCreateCommand()) cmd.AddCommand(newDeleteCommand()) cmd.AddCommand(newStatusCommand()) cmd.AddCommand(newSetDefaultCommand()) + cmd.AddCommand(newRegisterKeyCommand()) return cmd } diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go index bf80a9919e..90139d6be8 100644 --- a/cmd/lakebox/list.go +++ b/cmd/lakebox/list.go @@ -51,14 +51,14 @@ Example: } defaultID := getDefault(profile) - fmt.Fprintf(cmd.OutOrStdout(), " %-30s %-12s %-10s %s\n", "ID", "STATUS", "KEY", "DEFAULT") + fmt.Fprintf(cmd.OutOrStdout(), " %-30s %-12s %s\n", "ID", "STATUS", "DEFAULT") for _, e := range entries { id := extractLakeboxID(e.Name) def := "" if id == defaultID { def = "*" } - fmt.Fprintf(cmd.OutOrStdout(), " %-30s %-12s %-10s %s\n", id, e.Status, e.PubkeyHashPrefix, def) + fmt.Fprintf(cmd.OutOrStdout(), " %-30s %-12s %s\n", id, e.Status, def) } return nil }, diff --git a/cmd/lakebox/register_key.go b/cmd/lakebox/register_key.go new file mode 100644 index 0000000000..5a19cc4f57 --- /dev/null +++ b/cmd/lakebox/register_key.go @@ -0,0 +1,55 @@ +package lakebox + +import ( + "fmt" + "os" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/spf13/cobra" +) + +func newRegisterKeyCommand() *cobra.Command { + var publicKeyFile string + + cmd := &cobra.Command{ + Use: "register-key", + Short: "Register an SSH public key for lakebox access", + Long: `Register an SSH public key with the lakebox service. + +Once registered, the key can be used to SSH into any of your lakeboxes. +A user can have multiple registered keys; any of them grants access to +all lakeboxes owned by that user. + +Example: + databricks lakebox register-key --public-key-file ~/.ssh/id_ed25519.pub`, + PreRunE: root.MustWorkspaceClient, + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + api := newLakeboxAPI(w) + + if publicKeyFile == "" { + return fmt.Errorf("--public-key-file is required") + } + + data, err := os.ReadFile(publicKeyFile) + if err != nil { + return fmt.Errorf("failed to read public key file %s: %w", publicKeyFile, err) + } + + publicKey := string(data) + if err := api.registerKey(ctx, publicKey); err != nil { + return fmt.Errorf("failed to register key: %w", err) + } + + fmt.Fprintln(cmd.ErrOrStderr(), "SSH public key registered.") + return nil + }, + } + + cmd.Flags().StringVar(&publicKeyFile, "public-key-file", "", "Path to SSH public key file to register") + _ = cmd.MarkFlagRequired("public-key-file") + + return cmd +} diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go index 1afd968211..4bb130496d 100644 --- a/cmd/lakebox/status.go +++ b/cmd/lakebox/status.go @@ -45,9 +45,6 @@ Example: if entry.FQDN != "" { fmt.Fprintf(cmd.OutOrStdout(), "FQDN: %s\n", entry.FQDN) } - if entry.PubkeyHashPrefix != "" { - fmt.Fprintf(cmd.OutOrStdout(), "Key: %s\n", entry.PubkeyHashPrefix) - } return nil }, } From f8f8cc1aa04add672a448c1b399589ecb1a49435 Mon Sep 17 00:00:00 2001 From: shuochen0311 Date: Tue, 14 Apr 2026 03:11:22 +0000 Subject: [PATCH 3/7] Simplify SSH flow: register command, direct SSH args, remove config writes - Add 'register' command: generates ~/.ssh/lakebox_rsa and registers with API - Remove 'register-key' command (replaced by 'register') - Remove 'login' command (use 'auth login' + 'register' separately) - SSH command passes options directly as args instead of writing ~/.ssh/config - Check for ssh-keygen availability with helpful install instructions Co-authored-by: Isaac --- cmd/cmd.go | 6 +- cmd/lakebox/lakebox.go | 23 +++--- cmd/lakebox/register.go | 110 ++++++++++++++++++++++++++++ cmd/lakebox/register_key.go | 55 -------------- cmd/lakebox/ssh.go | 141 +++++------------------------------- 5 files changed, 148 insertions(+), 187 deletions(-) create mode 100644 cmd/lakebox/register.go delete mode 100644 cmd/lakebox/register_key.go diff --git a/cmd/cmd.go b/cmd/cmd.go index fe81149c08..c120f25aa7 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -19,8 +19,12 @@ Lakebox provides SSH-accessible development environments backed by microVM isolation. Each lakebox is a personal sandbox with pre-installed tooling (Python, Node.js, Rust, Databricks CLI) and persistent storage. +Getting started: + lakebox auth login --host https://... # authenticate to Databricks + lakebox register # generate SSH key and register + lakebox ssh # SSH to your default lakebox + Common workflows: - lakebox auth login # authenticate to Databricks lakebox ssh # SSH to your default lakebox lakebox ssh my-project # SSH to a named lakebox lakebox list # list your lakeboxes diff --git a/cmd/lakebox/lakebox.go b/cmd/lakebox/lakebox.go index aa9463bca8..127b5d93bf 100644 --- a/cmd/lakebox/lakebox.go +++ b/cmd/lakebox/lakebox.go @@ -14,28 +14,31 @@ Lakebox provides SSH-accessible development environments backed by microVM isolation. Each lakebox is a personal sandbox with pre-installed tooling (Python, Node.js, Rust, Databricks CLI) and persistent storage. +Getting started: + lakebox auth login --host https://... # authenticate to Databricks + lakebox register # generate SSH key and register + lakebox ssh # SSH to your default lakebox + Common workflows: - lakebox auth login # authenticate to Databricks - lakebox ssh # SSH to your default lakebox - lakebox ssh my-project # SSH to a named lakebox - lakebox list # list your lakeboxes - lakebox create # create a new lakebox - lakebox delete my-project # delete a lakebox - lakebox status my-project # show lakebox status - lakebox register-key --public-key-file ~/.ssh/id_rsa.pub # register SSH key + lakebox ssh # SSH to your default lakebox + lakebox ssh my-project # SSH to a named lakebox + lakebox list # list your lakeboxes + lakebox create # create a new lakebox + lakebox delete my-project # delete a lakebox + lakebox status my-project # show lakebox status The CLI manages your ~/.ssh/config so you can also connect directly: - ssh my-project # after 'lakebox ssh' + ssh my-project # after 'lakebox ssh' `, } + cmd.AddCommand(newRegisterCommand()) cmd.AddCommand(newSSHCommand()) cmd.AddCommand(newListCommand()) cmd.AddCommand(newCreateCommand()) cmd.AddCommand(newDeleteCommand()) cmd.AddCommand(newStatusCommand()) cmd.AddCommand(newSetDefaultCommand()) - cmd.AddCommand(newRegisterKeyCommand()) return cmd } diff --git a/cmd/lakebox/register.go b/cmd/lakebox/register.go new file mode 100644 index 0000000000..7286a14bf5 --- /dev/null +++ b/cmd/lakebox/register.go @@ -0,0 +1,110 @@ +package lakebox + +import ( + "fmt" + "os" + "os/exec" + "path/filepath" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/spf13/cobra" +) + +const lakeboxKeyName = "lakebox_rsa" + +func newRegisterCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "register", + Short: "Register this machine for lakebox SSH access", + Long: `Generate a dedicated SSH key for lakebox and register it with the service. + +This command: +1. Generates an RSA SSH key at ~/.ssh/lakebox_rsa (if it doesn't exist) +2. Registers the public key with the lakebox service + +After registration, 'lakebox ssh' will use this key automatically. +Run this once per machine. + +Example: + lakebox register`, + PreRunE: root.MustWorkspaceClient, + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + api := newLakeboxAPI(w) + + keyPath, generated, err := ensureLakeboxKey() + if err != nil { + return fmt.Errorf("failed to ensure lakebox SSH key: %w", err) + } + + if generated { + fmt.Fprintf(cmd.ErrOrStderr(), "Generated SSH key: %s\n", keyPath) + } else { + fmt.Fprintf(cmd.ErrOrStderr(), "Using existing SSH key: %s\n", keyPath) + } + + pubKeyData, err := os.ReadFile(keyPath + ".pub") + if err != nil { + return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err) + } + + if err := api.registerKey(ctx, string(pubKeyData)); err != nil { + return fmt.Errorf("failed to register key: %w", err) + } + + fmt.Fprintln(cmd.ErrOrStderr(), "Registered. You can now use 'lakebox ssh' to connect.") + return nil + }, + } + + return cmd +} + +// lakeboxKeyPath returns the path to the dedicated lakebox SSH key. +func lakeboxKeyPath() (string, error) { + homeDir, err := os.UserHomeDir() + if err != nil { + return "", err + } + return filepath.Join(homeDir, ".ssh", lakeboxKeyName), nil +} + +// ensureLakeboxKey returns the path to the lakebox SSH key, generating it if +// it doesn't exist. Returns (path, wasGenerated, error). +func ensureLakeboxKey() (string, bool, error) { + keyPath, err := lakeboxKeyPath() + if err != nil { + return "", false, err + } + + if _, err := os.Stat(keyPath); err == nil { + return keyPath, false, nil + } + + // Check that ssh-keygen is available before trying to generate. + if _, err := exec.LookPath("ssh-keygen"); err != nil { + return "", false, fmt.Errorf( + "ssh-keygen not found in PATH.\n" + + "Please install OpenSSH and run 'lakebox register' again.\n" + + " macOS: brew install openssh\n" + + " Ubuntu: sudo apt install openssh-client\n" + + " Windows: install Git for Windows (includes ssh-keygen)") + } + + sshDir := filepath.Dir(keyPath) + if err := os.MkdirAll(sshDir, 0700); err != nil { + return "", false, fmt.Errorf("failed to create %s: %w", sshDir, err) + } + + genCmd := exec.Command("ssh-keygen", "-t", "rsa", "-b", "4096", "-f", keyPath, "-N", "", "-q", "-C", "lakebox") + genCmd.Stdin = os.Stdin + genCmd.Stdout = os.Stderr + genCmd.Stderr = os.Stderr + if err := genCmd.Run(); err != nil { + return "", false, fmt.Errorf("ssh-keygen failed: %w", err) + } + + return keyPath, true, nil +} diff --git a/cmd/lakebox/register_key.go b/cmd/lakebox/register_key.go deleted file mode 100644 index 5a19cc4f57..0000000000 --- a/cmd/lakebox/register_key.go +++ /dev/null @@ -1,55 +0,0 @@ -package lakebox - -import ( - "fmt" - "os" - - "github.com/databricks/cli/cmd/root" - "github.com/databricks/cli/libs/cmdctx" - "github.com/spf13/cobra" -) - -func newRegisterKeyCommand() *cobra.Command { - var publicKeyFile string - - cmd := &cobra.Command{ - Use: "register-key", - Short: "Register an SSH public key for lakebox access", - Long: `Register an SSH public key with the lakebox service. - -Once registered, the key can be used to SSH into any of your lakeboxes. -A user can have multiple registered keys; any of them grants access to -all lakeboxes owned by that user. - -Example: - databricks lakebox register-key --public-key-file ~/.ssh/id_ed25519.pub`, - PreRunE: root.MustWorkspaceClient, - RunE: func(cmd *cobra.Command, args []string) error { - ctx := cmd.Context() - w := cmdctx.WorkspaceClient(ctx) - api := newLakeboxAPI(w) - - if publicKeyFile == "" { - return fmt.Errorf("--public-key-file is required") - } - - data, err := os.ReadFile(publicKeyFile) - if err != nil { - return fmt.Errorf("failed to read public key file %s: %w", publicKeyFile, err) - } - - publicKey := string(data) - if err := api.registerKey(ctx, publicKey); err != nil { - return fmt.Errorf("failed to register key: %w", err) - } - - fmt.Fprintln(cmd.ErrOrStderr(), "SSH public key registered.") - return nil - }, - } - - cmd.Flags().StringVar(&publicKeyFile, "public-key-file", "", "Path to SSH public key file to register") - _ = cmd.MarkFlagRequired("public-key-file") - - return cmd -} diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go index 1978dec684..8868f38e81 100644 --- a/cmd/lakebox/ssh.go +++ b/cmd/lakebox/ssh.go @@ -4,9 +4,7 @@ import ( "fmt" "os" "os/exec" - "path/filepath" "runtime" - "strings" "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" @@ -16,10 +14,6 @@ import ( const ( defaultGatewayHost = "uw2.dbrx.dev" defaultGatewayPort = "2222" - - // SSH config block markers for idempotent updates. - sshConfigMarkerStart = "# --- Lakebox managed start ---" - sshConfigMarkerEnd = "# --- Lakebox managed end ---" ) func newSSHCommand() *cobra.Command { @@ -57,10 +51,13 @@ Example: profile = w.Config.Host } - // Ensure SSH key exists. - keyPath, err := ensureSSHKey() + // Use the dedicated lakebox SSH key. + keyPath, err := lakeboxKeyPath() if err != nil { - return fmt.Errorf("failed to ensure SSH key: %w", err) + return fmt.Errorf("failed to determine lakebox key path: %w", err) + } + if _, err := os.Stat(keyPath); os.IsNotExist(err) { + return fmt.Errorf("lakebox SSH key not found at %s — run 'lakebox register' first", keyPath) } fmt.Fprintf(cmd.ErrOrStderr(), "Using SSH key: %s\n", keyPath) @@ -94,19 +91,9 @@ Example: } } - // Write SSH config entry for this lakebox. - sshConfigPath, err := sshConfigFilePath() - if err != nil { - return err - } - entry := buildSSHConfigEntry(lakeboxID, gatewayHost, gatewayPort, keyPath) - if err := writeSSHConfigEntry(sshConfigPath, lakeboxID, entry); err != nil { - return fmt.Errorf("failed to update SSH config: %w", err) - } - fmt.Fprintf(cmd.ErrOrStderr(), "Connecting to %s@%s:%s...\n", lakeboxID, gatewayHost, gatewayPort) - return execSSH(lakeboxID) + return execSSHDirect(lakeboxID, gatewayHost, gatewayPort, keyPath) }, } @@ -116,112 +103,24 @@ Example: return cmd } -// ensureSSHKey checks for an existing SSH key and generates one if missing. -func ensureSSHKey() (string, error) { - homeDir, err := os.UserHomeDir() - if err != nil { - return "", err - } - - candidates := []string{ - filepath.Join(homeDir, ".ssh", "id_ed25519"), - filepath.Join(homeDir, ".ssh", "id_rsa"), - } - for _, p := range candidates { - if _, err := os.Stat(p); err == nil { - return p, nil - } - } - - // Generate ed25519 key. - keyPath := candidates[0] - sshDir := filepath.Dir(keyPath) - if err := os.MkdirAll(sshDir, 0700); err != nil { - return "", fmt.Errorf("failed to create %s: %w", sshDir, err) - } - - cmd := exec.Command("ssh-keygen", "-t", "ed25519", "-f", keyPath, "-N", "", "-q") - cmd.Stdin = os.Stdin - cmd.Stdout = os.Stderr - cmd.Stderr = os.Stderr - if err := cmd.Run(); err != nil { - return "", fmt.Errorf("ssh-keygen failed: %w", err) - } - - return keyPath, nil -} - -func sshConfigFilePath() (string, error) { - homeDir, err := os.UserHomeDir() - if err != nil { - return "", err - } - return filepath.Join(homeDir, ".ssh", "config"), nil -} - -// buildSSHConfigEntry creates the SSH config block for a lakebox. -// The lakebox ID is used as both the Host alias and the SSH User. -func buildSSHConfigEntry(lakeboxID, host, port, keyPath string) string { - return fmt.Sprintf(`Host %s - HostName %s - Port %s - User %s - IdentityFile %s - IdentitiesOnly yes - PreferredAuthentications publickey - PasswordAuthentication no - KbdInteractiveAuthentication no - StrictHostKeyChecking no - UserKnownHostsFile /dev/null - LogLevel INFO -`, lakeboxID, host, port, lakeboxID, keyPath) -} - -// writeSSHConfigEntry idempotently writes a single lakebox entry to ~/.ssh/config. -// Replaces any existing lakebox block in-place. -func writeSSHConfigEntry(configPath, lakeboxID, entry string) error { - sshDir := filepath.Dir(configPath) - if err := os.MkdirAll(sshDir, 0700); err != nil { - return err - } - - existing, err := os.ReadFile(configPath) - if err != nil && !os.IsNotExist(err) { - return err - } - - wrappedEntry := fmt.Sprintf("%s\n%s%s\n", sshConfigMarkerStart, entry, sshConfigMarkerEnd) - content := string(existing) - - // Remove existing lakebox block if present. - startIdx := strings.Index(content, sshConfigMarkerStart) - if startIdx >= 0 { - endIdx := strings.Index(content[startIdx:], sshConfigMarkerEnd) - if endIdx >= 0 { - endIdx += startIdx + len(sshConfigMarkerEnd) - if endIdx < len(content) && content[endIdx] == '\n' { - endIdx++ - } - content = content[:startIdx] + content[endIdx:] - } - } - - if !strings.HasSuffix(content, "\n") && len(content) > 0 { - content += "\n" - } - content += wrappedEntry - - return os.WriteFile(configPath, []byte(content), 0600) -} - -// execSSH execs into ssh using the lakebox ID as the Host alias. -func execSSH(lakeboxID string) error { +// execSSHDirect execs into ssh with all options passed as args (no ~/.ssh/config needed). +func execSSHDirect(lakeboxID, host, port, keyPath string) error { sshPath, err := exec.LookPath("ssh") if err != nil { return fmt.Errorf("ssh not found in PATH: %w", err) } - args := []string{"ssh", lakeboxID} + args := []string{ + "ssh", + "-i", keyPath, + "-p", port, + "-o", "IdentitiesOnly=yes", + "-o", "PreferredAuthentications=publickey", + "-o", "StrictHostKeyChecking=no", + "-o", "UserKnownHostsFile=/dev/null", + "-o", "LogLevel=ERROR", + fmt.Sprintf("%s@%s", lakeboxID, host), + } if runtime.GOOS == "windows" { cmd := exec.Command(sshPath, args[1:]...) From 4b4186113ebf7cc8790ec9a0766ba2102d9f48cb Mon Sep 17 00:00:00 2001 From: shuochen0311 Date: Tue, 14 Apr 2026 03:22:20 +0000 Subject: [PATCH 4/7] Auto-register SSH key after auth login, fix login hook matching - Hook into auth login PostRun to auto-generate ~/.ssh/lakebox_rsa and register it after OAuth completes - Fix hook: match on sub.Name() not sub.Use (Use includes args) - Export EnsureAndReadKey and RegisterKey for use by auth hook - Update help text Co-authored-by: Isaac --- cmd/cmd.go | 52 ++++++++++++++++++++++++++++++++++++++--- cmd/lakebox/lakebox.go | 3 +-- cmd/lakebox/register.go | 23 ++++++++++++++++++ 3 files changed, 73 insertions(+), 5 deletions(-) diff --git a/cmd/cmd.go b/cmd/cmd.go index c120f25aa7..ddbb70f451 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -2,10 +2,12 @@ package cmd import ( "context" + "fmt" "github.com/databricks/cli/cmd/auth" "github.com/databricks/cli/cmd/lakebox" "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" "github.com/spf13/cobra" ) @@ -20,8 +22,7 @@ microVM isolation. Each lakebox is a personal sandbox with pre-installed tooling (Python, Node.js, Rust, Databricks CLI) and persistent storage. Getting started: - lakebox auth login --host https://... # authenticate to Databricks - lakebox register # generate SSH key and register + lakebox auth login --host https://... # authenticate to Databricks workspace and lakebox service lakebox ssh # SSH to your default lakebox Common workflows: @@ -37,7 +38,52 @@ The CLI manages your ~/.ssh/config so you can also connect directly: ` cli.CompletionOptions.DisableDefaultCmd = true - cli.AddCommand(auth.New()) + authCmd := auth.New() + // Hook into 'auth login' to auto-register SSH key after OAuth completes. + for _, sub := range authCmd.Commands() { + if sub.Name() == "login" { + origRunE := sub.RunE + sub.RunE = func(cmd *cobra.Command, args []string) error { + // Run the original auth login. + if err := origRunE(cmd, args); err != nil { + return err + } + + // Auto-register: generate lakebox SSH key and register it. + fmt.Fprintln(cmd.ErrOrStderr(), "") + fmt.Fprintln(cmd.ErrOrStderr(), "Setting up SSH access...") + + keyPath, pubKey, err := lakebox.EnsureAndReadKey() + if err != nil { + fmt.Fprintf(cmd.ErrOrStderr(), + "SSH key setup failed: %v\n"+ + "You can set it up later with: lakebox register\n", err) + return nil + } + fmt.Fprintf(cmd.ErrOrStderr(), "Using SSH key: %s\n", keyPath) + + if err := root.MustWorkspaceClient(cmd, args); err != nil { + fmt.Fprintf(cmd.ErrOrStderr(), + "Could not initialize workspace client for key registration.\n"+ + "Run 'lakebox register' to complete setup.\n") + return nil + } + + w := cmdctx.WorkspaceClient(cmd.Context()) + if err := lakebox.RegisterKey(cmd.Context(), w, pubKey); err != nil { + fmt.Fprintf(cmd.ErrOrStderr(), + "Key registration failed: %v\n"+ + "Run 'lakebox register' to retry.\n", err) + return nil + } + + fmt.Fprintln(cmd.ErrOrStderr(), "SSH key registered. You're ready to use 'lakebox ssh'.") + return nil + } + break + } + } + cli.AddCommand(authCmd) // Register lakebox subcommands directly at root level. for _, sub := range lakebox.New().Commands() { diff --git a/cmd/lakebox/lakebox.go b/cmd/lakebox/lakebox.go index 127b5d93bf..4afa321241 100644 --- a/cmd/lakebox/lakebox.go +++ b/cmd/lakebox/lakebox.go @@ -15,8 +15,7 @@ microVM isolation. Each lakebox is a personal sandbox with pre-installed tooling (Python, Node.js, Rust, Databricks CLI) and persistent storage. Getting started: - lakebox auth login --host https://... # authenticate to Databricks - lakebox register # generate SSH key and register + lakebox auth login --host https://... # authenticate to Databricks workspace and lakebox service lakebox ssh # SSH to your default lakebox Common workflows: diff --git a/cmd/lakebox/register.go b/cmd/lakebox/register.go index 7286a14bf5..a1da60422b 100644 --- a/cmd/lakebox/register.go +++ b/cmd/lakebox/register.go @@ -1,6 +1,7 @@ package lakebox import ( + "context" "fmt" "os" "os/exec" @@ -8,6 +9,7 @@ import ( "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/databricks-sdk-go" "github.com/spf13/cobra" ) @@ -108,3 +110,24 @@ func ensureLakeboxKey() (string, bool, error) { return keyPath, true, nil } + +// EnsureAndReadKey generates the lakebox SSH key if needed and returns +// (keyPath, publicKeyContent, error). Exported for use by the auth login hook. +func EnsureAndReadKey() (string, string, error) { + keyPath, _, err := ensureLakeboxKey() + if err != nil { + return "", "", err + } + pubKeyData, err := os.ReadFile(keyPath + ".pub") + if err != nil { + return "", "", fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err) + } + return keyPath, string(pubKeyData), nil +} + +// RegisterKey registers a public key with the lakebox API. Exported for use +// by the auth login hook. +func RegisterKey(ctx context.Context, w *databricks.WorkspaceClient, pubKey string) error { + api := newLakeboxAPI(w) + return api.registerKey(ctx, pubKey) +} From df599e9273beebe13ec81c7d5341b4b77001e2a8 Mon Sep 17 00:00:00 2001 From: shuochen0311 Date: Tue, 14 Apr 2026 21:42:07 +0000 Subject: [PATCH 5/7] Support passthrough args and remote commands in lakebox ssh Everything after -- is passed directly to the ssh process, enabling: lakebox ssh -- echo hello # run command and return lakebox ssh -- cat /etc/os-release lakebox ssh -- -L 8080:localhost:8080 # port forwarding Co-authored-by: Isaac --- cmd/lakebox/ssh.go | 90 +++++++++++++++++++++++++--------------------- 1 file changed, 50 insertions(+), 40 deletions(-) diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go index 8868f38e81..7559893bfb 100644 --- a/cmd/lakebox/ssh.go +++ b/cmd/lakebox/ssh.go @@ -21,24 +21,24 @@ func newSSHCommand() *cobra.Command { var gatewayPort string cmd := &cobra.Command{ - Use: "ssh [lakebox-id]", + Use: "ssh [lakebox-id] [-- ...]", Short: "SSH into a Lakebox environment", Long: `SSH into a Lakebox environment. -This command: -1. Authenticates to the Databricks workspace -2. Ensures you have a local SSH key (~/.ssh/id_ed25519) -3. Creates a lakebox if one doesn't exist (installs your public key) -4. Updates ~/.ssh/config with a Host entry for the lakebox -5. Connects via SSH using the lakebox ID as the SSH username - -Without arguments, creates a new lakebox. With a lakebox ID argument, -connects to the specified lakebox. - -Example: - databricks lakebox ssh # create and connect to a new lakebox - databricks lakebox ssh happy-panda-1234 # connect to existing lakebox`, - Args: cobra.MaximumNArgs(1), +Connect to your default or a named lakebox via SSH. Extra arguments +after -- are passed directly to the ssh process. This lets you run +remote commands, set up port forwarding, or pass any other ssh flags. + +Examples: + lakebox ssh # interactive shell on default lakebox + lakebox ssh happy-panda-1234 # interactive shell on specific lakebox + lakebox ssh -- ls -la /home # run command on default lakebox + lakebox ssh happy-panda-1234 -- cat /etc/os-release # run command on specific lakebox + lakebox ssh -- -L 8080:localhost:8080 # port forwarding on default lakebox`, + // Disable flag parsing after -- so extra args are passed through. + DisableFlagParsing: false, + // Accept any number of args: [lakebox-id] [-- extra...] + Args: cobra.ArbitraryArgs, PreRunE: func(cmd *cobra.Command, args []string) error { return root.MustWorkspaceClient(cmd, args) }, @@ -61,39 +61,47 @@ Example: } fmt.Fprintf(cmd.ErrOrStderr(), "Using SSH key: %s\n", keyPath) - // Determine lakebox ID: - // 1. Explicit arg → use it - // 2. Local default exists → use it - // 3. Neither → create a new one and set as default + // Parse args: first arg (if not starting with -) is lakebox ID, + // everything else is passed through to ssh. var lakeboxID string - if len(args) > 0 { + var extraArgs []string + + if len(args) > 0 && args[0] != "--" && args[0][0] != '-' { lakeboxID = args[0] - } else if def := getDefault(profile); def != "" { - lakeboxID = def - fmt.Fprintf(cmd.ErrOrStderr(), "Using default lakebox: %s\n", lakeboxID) + extraArgs = args[1:] } else { - api := newLakeboxAPI(w) - pubKeyData, err := os.ReadFile(keyPath + ".pub") - if err != nil { - return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err) - } - - fmt.Fprintf(cmd.ErrOrStderr(), "Creating lakebox...\n") - result, err := api.create(ctx, string(pubKeyData)) - if err != nil { - return fmt.Errorf("failed to create lakebox: %w", err) - } - lakeboxID = result.LakeboxID - fmt.Fprintf(cmd.ErrOrStderr(), "Lakebox %s created (status: %s)\n", lakeboxID, result.Status) + extraArgs = args + } - if err := setDefault(profile, lakeboxID); err != nil { - fmt.Fprintf(cmd.ErrOrStderr(), "Warning: failed to save default: %v\n", err) + // Determine lakebox ID if not explicit. + if lakeboxID == "" { + if def := getDefault(profile); def != "" { + lakeboxID = def + fmt.Fprintf(cmd.ErrOrStderr(), "Using default lakebox: %s\n", lakeboxID) + } else { + api := newLakeboxAPI(w) + pubKeyData, err := os.ReadFile(keyPath + ".pub") + if err != nil { + return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err) + } + + fmt.Fprintf(cmd.ErrOrStderr(), "Creating lakebox...\n") + result, err := api.create(ctx, string(pubKeyData)) + if err != nil { + return fmt.Errorf("failed to create lakebox: %w", err) + } + lakeboxID = result.LakeboxID + fmt.Fprintf(cmd.ErrOrStderr(), "Lakebox %s created (status: %s)\n", lakeboxID, result.Status) + + if err := setDefault(profile, lakeboxID); err != nil { + fmt.Fprintf(cmd.ErrOrStderr(), "Warning: failed to save default: %v\n", err) + } } } fmt.Fprintf(cmd.ErrOrStderr(), "Connecting to %s@%s:%s...\n", lakeboxID, gatewayHost, gatewayPort) - return execSSHDirect(lakeboxID, gatewayHost, gatewayPort, keyPath) + return execSSHDirect(lakeboxID, gatewayHost, gatewayPort, keyPath, extraArgs) }, } @@ -104,7 +112,8 @@ Example: } // execSSHDirect execs into ssh with all options passed as args (no ~/.ssh/config needed). -func execSSHDirect(lakeboxID, host, port, keyPath string) error { +// Extra args are appended after the destination (for remote commands or ssh flags). +func execSSHDirect(lakeboxID, host, port, keyPath string, extraArgs []string) error { sshPath, err := exec.LookPath("ssh") if err != nil { return fmt.Errorf("ssh not found in PATH: %w", err) @@ -121,6 +130,7 @@ func execSSHDirect(lakeboxID, host, port, keyPath string) error { "-o", "LogLevel=ERROR", fmt.Sprintf("%s@%s", lakeboxID, host), } + args = append(args, extraArgs...) if runtime.GOOS == "windows" { cmd := exec.Command(sshPath, args[1:]...) From cd2579760e15edbc4af014da6bc9963a4669110e Mon Sep 17 00:00:00 2001 From: Stas Kelvich Date: Tue, 14 Apr 2026 15:30:12 -0700 Subject: [PATCH 6/7] Fix workspace client init after login, persist last profile After 'lakebox auth login --host ', the post-login hook now constructs the workspace client directly from the --host/--profile flags instead of using MustWorkspaceClient (which started with an empty config and fell back to the DEFAULT profile). All lakebox commands now use a mustWorkspaceClient wrapper that reads the last-login profile from ~/.databricks/lakebox.json, so 'lakebox ssh' uses the correct profile without requiring --profile on every invocation. Also adds install.sh and upload.sh scripts. --- cmd/cmd.go | 30 +++++++++++++--- cmd/lakebox/create.go | 3 +- cmd/lakebox/default.go | 3 +- cmd/lakebox/delete.go | 3 +- cmd/lakebox/lakebox.go | 15 +++++++- cmd/lakebox/list.go | 3 +- cmd/lakebox/register.go | 3 +- cmd/lakebox/ssh.go | 5 +-- cmd/lakebox/state.go | 21 +++++++++++ cmd/lakebox/status.go | 3 +- install.sh | 80 +++++++++++++++++++++++++++++++++++++++++ upload.sh | 13 +++++++ 12 files changed, 160 insertions(+), 22 deletions(-) create mode 100755 install.sh create mode 100755 upload.sh diff --git a/cmd/cmd.go b/cmd/cmd.go index ddbb70f451..8a70375514 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -3,11 +3,12 @@ package cmd import ( "context" "fmt" + "strings" "github.com/databricks/cli/cmd/auth" "github.com/databricks/cli/cmd/lakebox" "github.com/databricks/cli/cmd/root" - "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/databricks-sdk-go" "github.com/spf13/cobra" ) @@ -62,14 +63,33 @@ The CLI manages your ~/.ssh/config so you can also connect directly: } fmt.Fprintf(cmd.ErrOrStderr(), "Using SSH key: %s\n", keyPath) - if err := root.MustWorkspaceClient(cmd, args); err != nil { + host := cmd.Flag("host").Value.String() + if host == "" && len(args) > 0 { + host = args[0] + } + profile := cmd.Flag("profile").Value.String() + if profile == "" && host != "" { + // Derive profile name the same way auth login does. + h := strings.TrimPrefix(host, "https://") + h = strings.TrimPrefix(h, "http://") + profile = strings.SplitN(h, ".", 2)[0] + } + if profile != "" { + if err := lakebox.SetLastProfile(profile); err != nil { + fmt.Fprintf(cmd.ErrOrStderr(), "Warning: failed to save last profile: %v\n", err) + } + } + w, err := databricks.NewWorkspaceClient(&databricks.Config{ + Host: host, + Profile: profile, + }) + if err != nil { fmt.Fprintf(cmd.ErrOrStderr(), - "Could not initialize workspace client for key registration.\n"+ - "Run 'lakebox register' to complete setup.\n") + "Could not initialize workspace client for key registration: %v\n"+ + "Run 'lakebox register' to complete setup.\n", err) return nil } - w := cmdctx.WorkspaceClient(cmd.Context()) if err := lakebox.RegisterKey(cmd.Context(), w, pubKey); err != nil { fmt.Fprintf(cmd.ErrOrStderr(), "Key registration failed: %v\n"+ diff --git a/cmd/lakebox/create.go b/cmd/lakebox/create.go index 872776cc8d..db1a22ebb7 100644 --- a/cmd/lakebox/create.go +++ b/cmd/lakebox/create.go @@ -4,7 +4,6 @@ import ( "fmt" "os" - "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" "github.com/spf13/cobra" ) @@ -26,7 +25,7 @@ authorized_keys so you can SSH directly. Otherwise the gateway key is used. Example: databricks lakebox create databricks lakebox create --public-key-file ~/.ssh/id_ed25519.pub`, - PreRunE: root.MustWorkspaceClient, + PreRunE: mustWorkspaceClient, 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 9d5a366c9c..b632c5984a 100644 --- a/cmd/lakebox/default.go +++ b/cmd/lakebox/default.go @@ -3,7 +3,6 @@ package lakebox import ( "fmt" - "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" "github.com/spf13/cobra" ) @@ -19,7 +18,7 @@ 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, + PreRunE: mustWorkspaceClient, RunE: func(cmd *cobra.Command, args []string) error { w := cmdctx.WorkspaceClient(cmd.Context()) profile := w.Config.Profile diff --git a/cmd/lakebox/delete.go b/cmd/lakebox/delete.go index a814083ed3..9c8ce93963 100644 --- a/cmd/lakebox/delete.go +++ b/cmd/lakebox/delete.go @@ -3,7 +3,6 @@ package lakebox import ( "fmt" - "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" "github.com/spf13/cobra" ) @@ -20,7 +19,7 @@ creator (same auth token) can delete a lakebox. Example: databricks lakebox delete happy-panda-1234`, Args: cobra.ExactArgs(1), - PreRunE: root.MustWorkspaceClient, + PreRunE: mustWorkspaceClient, 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 4afa321241..6a968df87a 100644 --- a/cmd/lakebox/lakebox.go +++ b/cmd/lakebox/lakebox.go @@ -1,6 +1,7 @@ package lakebox import ( + "github.com/databricks/cli/cmd/root" "github.com/spf13/cobra" ) @@ -32,12 +33,24 @@ The CLI manages your ~/.ssh/config so you can also connect directly: } cmd.AddCommand(newRegisterCommand()) + cmd.AddCommand(newSetDefaultCommand()) cmd.AddCommand(newSSHCommand()) cmd.AddCommand(newListCommand()) cmd.AddCommand(newCreateCommand()) cmd.AddCommand(newDeleteCommand()) cmd.AddCommand(newStatusCommand()) - cmd.AddCommand(newSetDefaultCommand()) return cmd } + +// mustWorkspaceClient applies the saved last-login profile when the user +// hasn't explicitly set --profile, then delegates to root.MustWorkspaceClient. +func mustWorkspaceClient(cmd *cobra.Command, args []string) error { + profileFlag := cmd.Flag("profile") + if profileFlag != nil && !profileFlag.Changed { + if last := GetLastProfile(); last != "" { + _ = profileFlag.Value.Set(last) + } + } + return root.MustWorkspaceClient(cmd, args) +} diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go index 90139d6be8..3222d1c10c 100644 --- a/cmd/lakebox/list.go +++ b/cmd/lakebox/list.go @@ -4,7 +4,6 @@ import ( "encoding/json" "fmt" - "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" "github.com/spf13/cobra" ) @@ -23,7 +22,7 @@ current status and ID. Example: databricks lakebox list databricks lakebox list --json`, - PreRunE: root.MustWorkspaceClient, + PreRunE: mustWorkspaceClient, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) diff --git a/cmd/lakebox/register.go b/cmd/lakebox/register.go index a1da60422b..27d6cc59a1 100644 --- a/cmd/lakebox/register.go +++ b/cmd/lakebox/register.go @@ -7,7 +7,6 @@ import ( "os/exec" "path/filepath" - "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/databricks-sdk-go" "github.com/spf13/cobra" @@ -30,7 +29,7 @@ Run this once per machine. Example: lakebox register`, - PreRunE: root.MustWorkspaceClient, + PreRunE: mustWorkspaceClient, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go index 8868f38e81..86098baf5a 100644 --- a/cmd/lakebox/ssh.go +++ b/cmd/lakebox/ssh.go @@ -6,7 +6,6 @@ import ( "os/exec" "runtime" - "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" "github.com/spf13/cobra" ) @@ -39,9 +38,7 @@ Example: databricks lakebox ssh # create and connect to a new lakebox databricks lakebox ssh happy-panda-1234 # connect to existing lakebox`, Args: cobra.MaximumNArgs(1), - PreRunE: func(cmd *cobra.Command, args []string) error { - return root.MustWorkspaceClient(cmd, args) - }, + PreRunE: mustWorkspaceClient, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) diff --git a/cmd/lakebox/state.go b/cmd/lakebox/state.go index c0c8ad2d84..b84b5b16e1 100644 --- a/cmd/lakebox/state.go +++ b/cmd/lakebox/state.go @@ -12,6 +12,8 @@ import ( type stateFile struct { // Profile name → default lakebox ID. Defaults map[string]string `json:"defaults"` + // Last profile used with 'lakebox auth login'. + LastProfile string `json:"last_profile,omitempty"` } func stateFilePath() (string, error) { @@ -80,6 +82,25 @@ func setDefault(profile, lakeboxID string) error { return saveState(state) } +// GetLastProfile returns the profile saved by the most recent 'lakebox auth login'. +func GetLastProfile() string { + state, err := loadState() + if err != nil { + return "" + } + return state.LastProfile +} + +// SetLastProfile persists the profile used during 'lakebox auth login'. +func SetLastProfile(profile string) error { + state, err := loadState() + if err != nil { + return err + } + state.LastProfile = profile + return saveState(state) +} + func clearDefault(profile string) error { state, err := loadState() if err != nil { diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go index 4bb130496d..eaeeb8d7cc 100644 --- a/cmd/lakebox/status.go +++ b/cmd/lakebox/status.go @@ -4,7 +4,6 @@ import ( "encoding/json" "fmt" - "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" "github.com/spf13/cobra" ) @@ -21,7 +20,7 @@ Example: databricks lakebox status happy-panda-1234 databricks lakebox status happy-panda-1234 --json`, Args: cobra.ExactArgs(1), - PreRunE: root.MustWorkspaceClient, + PreRunE: mustWorkspaceClient, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) diff --git a/install.sh b/install.sh new file mode 100755 index 0000000000..acdf259b4c --- /dev/null +++ b/install.sh @@ -0,0 +1,80 @@ +#!/bin/sh +# Lakebox CLI installer — . <(curl -s devbox.dbrx.dev) + +_lakebox_install() { + INSTALL_DIR="$HOME/.lakebox/bin" + REMOTE_NAME="databricks" + LOCAL_NAME="lakebox" + BASE_URL="https://devbox.dbrx.dev" + + case "$(uname -s)" in + Linux*) OS="linux" ;; + Darwin*) OS="darwin" ;; + *) printf "error: unsupported OS: %s\n" "$(uname -s)" >&2; return 1 ;; + esac + + case "$(uname -m)" in + x86_64|amd64) ARCH="amd64" ;; + aarch64|arm64) ARCH="arm64" ;; + *) printf "error: unsupported arch: %s\n" "$(uname -m)" >&2; return 1 ;; + esac + + url="${BASE_URL}/${REMOTE_NAME}-${OS}-${ARCH}" + + printf "📦 Installing Lakebox CLI (%s/%s)...\n" "$OS" "$ARCH" + + mkdir -p "$INSTALL_DIR" || { printf "error: could not create %s\n" "$INSTALL_DIR" >&2; return 1; } + + if command -v curl >/dev/null 2>&1; then + curl -fSL --progress-bar "$url" -o "$INSTALL_DIR/$LOCAL_NAME" || { printf "error: download failed\n" >&2; return 1; } + elif command -v wget >/dev/null 2>&1; then + wget -q --show-progress "$url" -O "$INSTALL_DIR/$LOCAL_NAME" || { printf "error: download failed\n" >&2; return 1; } + else + printf "error: curl or wget is required\n" >&2; return 1 + fi + + chmod +x "$INSTALL_DIR/$LOCAL_NAME" + + PATH_LINE="export PATH=\"\$HOME/.lakebox/bin:\$PATH\"" + case ":$PATH:" in + *":$INSTALL_DIR:"*) ;; + *) + added=0 + for rc in "$HOME/.zshrc" "$HOME/.bashrc"; do + [ -f "$rc" ] || continue + if ! grep -qF '.lakebox/bin' "$rc" 2>/dev/null; then + printf '\n# Lakebox CLI\n%s\n' "$PATH_LINE" >> "$rc" + printf "📝 Updated %s\n" "$rc" + added=1 + fi + done + if [ "$added" = 0 ]; then + if [ "$OS" = "darwin" ]; then + rc="$HOME/.zshrc" + else + rc="$HOME/.bashrc" + fi + printf '\n# Lakebox CLI\n%s\n' "$PATH_LINE" >> "$rc" + printf "📝 Updated %s\n" "$rc" + fi + export PATH="$INSTALL_DIR:$PATH" + ;; + esac + + printf "\n✅ Lakebox CLI installed to %s\n" "$INSTALL_DIR/$LOCAL_NAME" + + LAKEBOX_HOST="https://dbsql-dev-testing-default.dev.databricks.com" + LAKEBOX_PROFILE="dbsql-dev-testing-default" + if ! grep -qF "$LAKEBOX_PROFILE" "$HOME/.databrickscfg" 2>/dev/null; then + printf "\n🔑 Logging in...\n" + lakebox auth login --host "$LAKEBOX_HOST" --profile "$LAKEBOX_PROFILE" + fi + + printf "\nCommon workflows:\n" + printf " lakebox ssh # SSH to your default lakebox\n" + printf " lakebox ssh my-project # SSH to a named lakebox\n" + printf " lakebox list # list your lakeboxes\n" +} + +_lakebox_install +unset -f _lakebox_install \ No newline at end of file diff --git a/upload.sh b/upload.sh new file mode 100755 index 0000000000..c55c0aa182 --- /dev/null +++ b/upload.sh @@ -0,0 +1,13 @@ +#!/bin/sh +set -eu + +HOST="arca.ssh" +FILES="install.sh databricks-darwin-amd64 databricks-darwin-arm64 databricks-linux-amd64 databricks-linux-arm64" + +for f in $FILES; do + printf "Uploading %s...\n" "$f" + scp "$f" "$HOST:~/" + ssh "$HOST" "~/unp-upload.sh ~/$f" +done + +printf "\nDone.\n" From c1168a414f429f47517557ab8b9ce3bc7084fe28 Mon Sep 17 00:00:00 2001 From: shuochen0311 Date: Thu, 16 Apr 2026 06:48:54 +0000 Subject: [PATCH 7/7] Add consistent terminal UI: spinners, colors, aligned output Single cyan accent color throughout. Bold for IDs, dim for metadata. Braille spinner with elapsed time during async operations. - create: animated spinner during provisioning - list: aligned columns with colored status, cyan bold for running - status: clean field layout - delete: spinner during removal - ssh: spinner during connection - register: spinner during key registration - Shared ui.go with all primitives Co-authored-by: Isaac --- cmd/lakebox/create.go | 21 +++--- cmd/lakebox/delete.go | 15 +++-- cmd/lakebox/list.go | 46 +++++++++++-- cmd/lakebox/register.go | 11 +++- cmd/lakebox/ssh.go | 17 ++--- cmd/lakebox/status.go | 13 ++-- cmd/lakebox/ui.go | 141 ++++++++++++++++++++++++++++++++++++++++ 7 files changed, 221 insertions(+), 43 deletions(-) create mode 100644 cmd/lakebox/ui.go diff --git a/cmd/lakebox/create.go b/cmd/lakebox/create.go index db1a22ebb7..c4ce3a439e 100644 --- a/cmd/lakebox/create.go +++ b/cmd/lakebox/create.go @@ -19,17 +19,14 @@ func newCreateCommand() *cobra.Command { Creates a new personal development environment backed by a microVM. Blocks until the lakebox is running and prints the lakebox ID. -If --public-key-file is provided, the key is installed in the lakebox's -authorized_keys so you can SSH directly. Otherwise the gateway key is used. - Example: - databricks lakebox create - databricks lakebox create --public-key-file ~/.ssh/id_ed25519.pub`, + lakebox create`, PreRunE: mustWorkspaceClient, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) api := newLakeboxAPI(w) + stderr := cmd.ErrOrStderr() var publicKey string if publicKeyFile != "" { @@ -40,37 +37,37 @@ Example: publicKey = string(data) } - fmt.Fprintf(cmd.ErrOrStderr(), "Creating lakebox...\n") + s := spin(stderr, "Provisioning your lakebox…") result, err := api.create(ctx, publicKey) if err != nil { + s.fail("Failed to create lakebox") return fmt.Errorf("failed to create lakebox: %w", err) } + s.ok(fmt.Sprintf("Lakebox %s is %s", bold(result.LakeboxID), status(result.Status))) + profile := w.Config.Profile if profile == "" { profile = w.Config.Host } - // Set as default if no default exists, or the current default - // has been deleted (no longer in the list). currentDefault := getDefault(profile) shouldSetDefault := currentDefault == "" if !shouldSetDefault && currentDefault != "" { - // Check if the current default still exists. if _, err := api.get(ctx, currentDefault); err != nil { shouldSetDefault = true } } if shouldSetDefault { if err := setDefault(profile, result.LakeboxID); err != nil { - fmt.Fprintf(cmd.ErrOrStderr(), "Warning: failed to save default: %v\n", err) + warn(stderr, fmt.Sprintf("Could not save default: %v", err)) } else { - fmt.Fprintf(cmd.ErrOrStderr(), "Set as default lakebox.\n") + field(stderr, "default", result.LakeboxID) } } - fmt.Fprintf(cmd.ErrOrStderr(), "Lakebox created (status: %s)\n", result.Status) + blank(stderr) fmt.Fprintln(cmd.OutOrStdout(), result.LakeboxID) return nil }, diff --git a/cmd/lakebox/delete.go b/cmd/lakebox/delete.go index 9c8ce93963..ba56e2a508 100644 --- a/cmd/lakebox/delete.go +++ b/cmd/lakebox/delete.go @@ -13,35 +13,36 @@ func newDeleteCommand() *cobra.Command { Short: "Delete a Lakebox environment", Long: `Delete a Lakebox environment. -Permanently terminates and removes the specified lakebox. Only the -creator (same auth token) can delete a lakebox. +Permanently terminates and removes the specified lakebox. Example: - databricks lakebox delete happy-panda-1234`, + lakebox delete happy-panda-1234`, Args: cobra.ExactArgs(1), PreRunE: mustWorkspaceClient, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) api := newLakeboxAPI(w) + stderr := cmd.ErrOrStderr() lakeboxID := args[0] + s := spin(stderr, fmt.Sprintf("Removing %s…", lakeboxID)) if err := api.delete(ctx, lakeboxID); err != nil { + s.fail(fmt.Sprintf("Failed to delete %s", lakeboxID)) return fmt.Errorf("failed to delete lakebox %s: %w", lakeboxID, err) } - // Clear default if we just deleted it. profile := w.Config.Profile if profile == "" { profile = w.Config.Host } if getDefault(profile) == lakeboxID { _ = clearDefault(profile) - fmt.Fprintf(cmd.ErrOrStderr(), "Cleared default lakebox.\n") + s.ok(fmt.Sprintf("Removed %s %s", bold(lakeboxID), dim("(default cleared)"))) + } else { + s.ok(fmt.Sprintf("Removed %s", bold(lakeboxID))) } - - fmt.Fprintf(cmd.ErrOrStderr(), "Deleted lakebox %s\n", lakeboxID) return nil }, } diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go index 3222d1c10c..2ed3149658 100644 --- a/cmd/lakebox/list.go +++ b/cmd/lakebox/list.go @@ -3,6 +3,7 @@ package lakebox import ( "encoding/json" "fmt" + "strings" "github.com/databricks/cli/libs/cmdctx" "github.com/spf13/cobra" @@ -20,8 +21,8 @@ Shows all lakeboxes associated with your account, including their current status and ID. Example: - databricks lakebox list - databricks lakebox list --json`, + lakebox list + lakebox list --json`, PreRunE: mustWorkspaceClient, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() @@ -40,7 +41,7 @@ Example: } if len(entries) == 0 { - fmt.Fprintln(cmd.ErrOrStderr(), "No lakeboxes found.") + fmt.Fprintf(cmd.ErrOrStderr(), " %sNo lakeboxes found.%s\n", dm, rs) return nil } @@ -50,15 +51,48 @@ Example: } defaultID := getDefault(profile) - fmt.Fprintf(cmd.OutOrStdout(), " %-30s %-12s %s\n", "ID", "STATUS", "DEFAULT") + out := cmd.OutOrStdout() + + // Compute column width. + col := 10 + for _, e := range entries { + if l := len(extractLakeboxID(e.Name)); l > col { + col = l + } + } + col += 2 + + blank(out) + fmt.Fprintf(out, " %s%-*s %-10s %s%s\n", dm, col, "ID", "STATUS", "DEFAULT", rs) + fmt.Fprintf(out, " %s%s%s\n", dm, strings.Repeat("─", col+22), rs) + for _, e := range entries { id := extractLakeboxID(e.Name) def := "" if id == defaultID { - def = "*" + def = accent("*") + } + // Pad ID manually to avoid ANSI codes breaking alignment. + idPad := col - len(id) + if idPad < 0 { + idPad = 0 + } + st := status(e.Status) + // Pad status to 10 visible chars. + stPad := 10 - len(e.Status) + if stPad < 0 { + stPad = 0 + } + idStr := bold(id) + if strings.EqualFold(e.Status, "running") { + idStr = cyan + bo + id + rs } - fmt.Fprintf(cmd.OutOrStdout(), " %-30s %-12s %s\n", id, e.Status, def) + fmt.Fprintf(out, " %s%s %s%s %s\n", + idStr, strings.Repeat(" ", idPad), + st, strings.Repeat(" ", stPad), + def) } + blank(out) return nil }, } diff --git a/cmd/lakebox/register.go b/cmd/lakebox/register.go index 27d6cc59a1..f3550d8e5d 100644 --- a/cmd/lakebox/register.go +++ b/cmd/lakebox/register.go @@ -40,10 +40,11 @@ Example: return fmt.Errorf("failed to ensure lakebox SSH key: %w", err) } + stderr := cmd.ErrOrStderr() if generated { - fmt.Fprintf(cmd.ErrOrStderr(), "Generated SSH key: %s\n", keyPath) + ok(stderr, fmt.Sprintf("Generated SSH key at %s", dim(keyPath))) } else { - fmt.Fprintf(cmd.ErrOrStderr(), "Using existing SSH key: %s\n", keyPath) + field(stderr, "key", keyPath) } pubKeyData, err := os.ReadFile(keyPath + ".pub") @@ -51,11 +52,15 @@ Example: return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err) } + s := spin(stderr, "Registering key…") if err := api.registerKey(ctx, string(pubKeyData)); err != nil { + s.fail("Failed to register key") return fmt.Errorf("failed to register key: %w", err) } + s.ok("SSH key registered") - fmt.Fprintln(cmd.ErrOrStderr(), "Registered. You can now use 'lakebox ssh' to connect.") + blank(stderr) + fmt.Fprintf(stderr, " Run %s to connect.\n\n", bold("lakebox ssh")) return nil }, } diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go index 04a999bd40..483dbd38a8 100644 --- a/cmd/lakebox/ssh.go +++ b/cmd/lakebox/ssh.go @@ -53,7 +53,7 @@ Examples: if _, err := os.Stat(keyPath); os.IsNotExist(err) { return fmt.Errorf("lakebox SSH key not found at %s — run 'lakebox register' first", keyPath) } - fmt.Fprintf(cmd.ErrOrStderr(), "Using SSH key: %s\n", keyPath) + stderr := cmd.ErrOrStderr() // Parse args: everything before -- is the optional lakebox ID, // everything after -- is passed through to ssh. @@ -62,15 +62,12 @@ Examples: dashAt := cmd.ArgsLenAtDash() if dashAt == -1 { - // No -- found: first arg (if any) is lakebox ID. if len(args) > 0 { lakeboxID = args[0] } } else if dashAt == 0 { - // -- is first: no lakebox ID, rest is extra args. extraArgs = args[dashAt:] } else { - // lakebox ID before --, extra args after. lakeboxID = args[0] extraArgs = args[dashAt:] } @@ -79,7 +76,6 @@ Examples: if lakeboxID == "" { if def := getDefault(profile); def != "" { lakeboxID = def - fmt.Fprintf(cmd.ErrOrStderr(), "Using default lakebox: %s\n", lakeboxID) } else { api := newLakeboxAPI(w) pubKeyData, err := os.ReadFile(keyPath + ".pub") @@ -87,22 +83,23 @@ Examples: return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err) } - fmt.Fprintf(cmd.ErrOrStderr(), "Creating lakebox...\n") + s := spin(stderr, "Provisioning your lakebox…") result, err := api.create(ctx, string(pubKeyData)) if err != nil { + s.fail("Failed to create lakebox") return fmt.Errorf("failed to create lakebox: %w", err) } lakeboxID = result.LakeboxID - fmt.Fprintf(cmd.ErrOrStderr(), "Lakebox %s created (status: %s)\n", lakeboxID, result.Status) + s.ok(fmt.Sprintf("Lakebox %s ready", bold(lakeboxID))) if err := setDefault(profile, lakeboxID); err != nil { - fmt.Fprintf(cmd.ErrOrStderr(), "Warning: failed to save default: %v\n", err) + warn(stderr, fmt.Sprintf("Could not save default: %v", err)) } } } - fmt.Fprintf(cmd.ErrOrStderr(), "Connecting to %s@%s:%s...\n", - lakeboxID, gatewayHost, gatewayPort) + s := spin(stderr, fmt.Sprintf("Connecting to %s…", bold(lakeboxID))) + s.ok(fmt.Sprintf("Connected to %s", bold(lakeboxID))) return execSSHDirect(lakeboxID, gatewayHost, gatewayPort, keyPath, extraArgs) }, } diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go index eaeeb8d7cc..bf2efbcaba 100644 --- a/cmd/lakebox/status.go +++ b/cmd/lakebox/status.go @@ -17,8 +17,8 @@ func newStatusCommand() *cobra.Command { Long: `Show detailed status of a Lakebox environment. Example: - databricks lakebox status happy-panda-1234 - databricks lakebox status happy-panda-1234 --json`, + lakebox status happy-panda-1234 + lakebox status happy-panda-1234 --json`, Args: cobra.ExactArgs(1), PreRunE: mustWorkspaceClient, RunE: func(cmd *cobra.Command, args []string) error { @@ -39,11 +39,14 @@ Example: return enc.Encode(entry) } - fmt.Fprintf(cmd.OutOrStdout(), "ID: %s\n", extractLakeboxID(entry.Name)) - fmt.Fprintf(cmd.OutOrStdout(), "Status: %s\n", entry.Status) + out := cmd.OutOrStdout() + blank(out) + field(out, "id", bold(extractLakeboxID(entry.Name))) + field(out, "status", status(entry.Status)) if entry.FQDN != "" { - fmt.Fprintf(cmd.OutOrStdout(), "FQDN: %s\n", entry.FQDN) + field(out, "fqdn", dim(entry.FQDN)) } + blank(out) return nil }, } diff --git a/cmd/lakebox/ui.go b/cmd/lakebox/ui.go new file mode 100644 index 0000000000..2eab33310c --- /dev/null +++ b/cmd/lakebox/ui.go @@ -0,0 +1,141 @@ +package lakebox + +import ( + "fmt" + "io" + "os" + "strings" + "sync" + "time" +) + +// Single accent color throughout. Bold for emphasis. Dim for metadata. +const ( + rs = "\033[0m" // reset + bo = "\033[1m" // bold + dm = "\033[2m" // dim + cyan = "\033[36m" // accent +) + +func isTTY(w io.Writer) bool { + if f, ok := w.(*os.File); ok { + fi, err := f.Stat() + if err != nil { + return false + } + return fi.Mode()&os.ModeCharDevice != 0 + } + return false +} + +// spinner shows a braille spinner like Claude Code. +type spinner struct { + w io.Writer + msg string + done chan struct{} + once sync.Once + started time.Time +} + +func spin(w io.Writer, msg string) *spinner { + s := &spinner{w: w, msg: msg, done: make(chan struct{}), started: time.Now()} + if isTTY(w) { + go s.run() + } else { + fmt.Fprintf(w, "* %s\n", msg) + } + return s +} + +func (s *spinner) run() { + frames := []string{"⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷"} + i := 0 + ticker := time.NewTicker(80 * time.Millisecond) + defer ticker.Stop() + for { + select { + case <-s.done: + return + case <-ticker.C: + elapsed := time.Since(s.started).Truncate(time.Second) + fmt.Fprintf(s.w, "\r %s%s%s %s%s%s %s(%s)%s ", + cyan, frames[i%len(frames)], rs, + bo, s.msg, rs, + dm, elapsed, rs) + i++ + } + } +} + +func (s *spinner) ok(msg string) { + s.once.Do(func() { + close(s.done) + if isTTY(s.w) { + fmt.Fprintf(s.w, "\r\033[K %s✓%s %s\n", cyan, rs, msg) + } else { + fmt.Fprintf(s.w, "✓ %s\n", msg) + } + }) +} + +func (s *spinner) fail(msg string) { + s.once.Do(func() { + close(s.done) + if isTTY(s.w) { + fmt.Fprintf(s.w, "\r\033[K %s✗%s %s\n", cyan, rs, msg) + } else { + fmt.Fprintf(s.w, "✗ %s\n", msg) + } + }) +} + +// --- Consistent output primitives --- + +// status formats a status string with the accent color. +func status(s string) string { + switch strings.ToLower(s) { + case "running": + return cyan + "running" + rs + case "stopped": + return dm + "stopped" + rs + case "creating": + return cyan + bo + "creating…" + rs + default: + return dm + strings.ToLower(s) + rs + } +} + +// field prints " label value" +func field(w io.Writer, label, value string) { + fmt.Fprintf(w, " %s%-10s%s %s\n", dm, label, rs, value) +} + +// ok prints " ✓ message" +func ok(w io.Writer, msg string) { + fmt.Fprintf(w, " %s✓%s %s\n", cyan, rs, msg) +} + +// warn prints " ! message" +func warn(w io.Writer, msg string) { + fmt.Fprintf(w, " %s!%s %s\n", cyan, rs, msg) +} + +// blank prints an empty line. +func blank(w io.Writer) { + fmt.Fprintln(w) +} + +// accent wraps text in the accent color. +func accent(s string) string { + return cyan + s + rs +} + +// bold wraps text in bold. +func bold(s string) string { + return bo + s + rs +} + +// dim wraps text in dim. +func dim(s string) string { + return dm + s + rs +}