diff --git a/cmd/cmd.go b/cmd/cmd.go index 014471f763..ddbb70f451 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -2,117 +2,93 @@ package cmd import ( "context" - "strings" + "fmt" - "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/databricks/cli/libs/cmdctx" "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) - - // 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.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. + +Getting started: + lakebox auth login --host https://... # authenticate to Databricks workspace and lakebox service + lakebox ssh # SSH to your default lakebox + +Common workflows: + 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 + + 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(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) } + cli.AddCommand(authCmd) - // 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..94877b4a42 --- /dev/null +++ b/cmd/lakebox/api.go @@ -0,0 +1,199 @@ +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"` +} + +// 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)) +} + +// 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 { + 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..4afa321241 --- /dev/null +++ b/cmd/lakebox/lakebox.go @@ -0,0 +1,43 @@ +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. + +Getting started: + lakebox auth login --host https://... # authenticate to Databricks workspace and lakebox service + lakebox ssh # SSH to your default lakebox + +Common workflows: + 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' +`, + } + + cmd.AddCommand(newRegisterCommand()) + 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..90139d6be8 --- /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 %s\n", "ID", "STATUS", "DEFAULT") + for _, e := range entries { + id := extractLakeboxID(e.Name) + def := "" + if id == defaultID { + def = "*" + } + fmt.Fprintf(cmd.OutOrStdout(), " %-30s %-12s %s\n", id, e.Status, def) + } + return nil + }, + } + + cmd.Flags().BoolVar(&outputJSON, "json", false, "Output as JSON") + + return cmd +} diff --git a/cmd/lakebox/register.go b/cmd/lakebox/register.go new file mode 100644 index 0000000000..a1da60422b --- /dev/null +++ b/cmd/lakebox/register.go @@ -0,0 +1,133 @@ +package lakebox + +import ( + "context" + "fmt" + "os" + "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" +) + +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 +} + +// 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) +} diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go new file mode 100644 index 0000000000..8868f38e81 --- /dev/null +++ b/cmd/lakebox/ssh.go @@ -0,0 +1,134 @@ +package lakebox + +import ( + "fmt" + "os" + "os/exec" + "runtime" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/spf13/cobra" +) + +const ( + defaultGatewayHost = "uw2.dbrx.dev" + defaultGatewayPort = "2222" +) + +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 + } + + // Use the dedicated lakebox SSH key. + keyPath, err := lakeboxKeyPath() + if err != nil { + 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) + + // 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) + } + } + + fmt.Fprintf(cmd.ErrOrStderr(), "Connecting to %s@%s:%s...\n", + lakeboxID, gatewayHost, gatewayPort) + return execSSHDirect(lakeboxID, gatewayHost, gatewayPort, keyPath) + }, + } + + cmd.Flags().StringVar(&gatewayHost, "gateway", defaultGatewayHost, "Lakebox gateway hostname") + cmd.Flags().StringVar(&gatewayPort, "port", defaultGatewayPort, "Lakebox gateway SSH port") + + return cmd +} + +// 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", + "-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:]...) + 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..4bb130496d --- /dev/null +++ b/cmd/lakebox/status.go @@ -0,0 +1,55 @@ +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) + } + return nil + }, + } + + cmd.Flags().BoolVar(&outputJSON, "json", false, "Output as JSON") + + return cmd +}