Add Lakebox CLI for managing Databricks sandbox environments

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
 }

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
+}