diff --git a/cmd/cmd.go b/cmd/cmd.go index 263c1da0c1..e76da540be 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -18,6 +18,7 @@ import ( "github.com/databricks/cli/cmd/experimental" "github.com/databricks/cli/cmd/fs" "github.com/databricks/cli/cmd/labs" + "github.com/databricks/cli/cmd/lakebox" "github.com/databricks/cli/cmd/pipelines" "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/cmd/selftest" @@ -105,6 +106,7 @@ func New(ctx context.Context) *cobra.Command { cli.AddCommand(configure.New()) cli.AddCommand(fs.New()) cli.AddCommand(labs.New(ctx)) + cli.AddCommand(lakebox.New()) cli.AddCommand(sync.New()) cli.AddCommand(version.New()) cli.AddCommand(selftest.New()) diff --git a/cmd/fuzz_panic_test.go b/cmd/fuzz_panic_test.go index 4fb5d5b9d3..e4037b4ef8 100644 --- a/cmd/fuzz_panic_test.go +++ b/cmd/fuzz_panic_test.go @@ -208,6 +208,7 @@ func isAutoGenerated(leaf leafCommand) bool { "configure": true, "experimental": true, "labs": true, + "lakebox": true, "pipelines": true, "psql": true, "selftest": true, diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go new file mode 100644 index 0000000000..219a7db142 --- /dev/null +++ b/cmd/lakebox/api.go @@ -0,0 +1,342 @@ +package lakebox + +import ( + "context" + "fmt" + "net/http" + "strings" + "time" + + "github.com/databricks/cli/libs/auth" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" +) + +// Sandboxes live under the `/sandboxes` sub-collection of the lakebox service +// namespace (see `lakebox.proto` `LakeboxService.CreateSandbox`). +const lakeboxAPIPath = "/api/2.0/lakebox/sandboxes" + +// SSH keys are nested under the lakebox service namespace alongside +// `sandboxes/` (see `LakeboxService.CreateSshKey`). +const lakeboxKeysAPIPath = "/api/2.0/lakebox/ssh-keys" + +// orgIDHeader is sent by multi-workspace gateways (e.g. dogfood staging) so +// the gateway can scope the credential to a specific workspace. Without it, +// requests fail with "Credential was not sent or was of an unsupported type +// for this API." +const orgIDHeader = "X-Databricks-Org-Id" + +// lakeboxAPI wraps the SDK ApiClient with workspace-id-aware request headers. +type lakeboxAPI struct { + c *client.DatabricksClient +} + +// sandboxCreateBody is the inner `Sandbox` message in the create payload. +// Only `name` is caller-settable today; all other fields are server-chosen. +type sandboxCreateBody struct { + Name string `json:"name,omitempty"` +} + +// createRequest is the JSON body for POST /api/2.0/lakebox/sandboxes. +// `CreateSandboxRequest { Sandbox sandbox = 1 }` has `body: "*"`, so the +// wire body is the full request with a `sandbox` wrapper. +type createRequest struct { + Sandbox sandboxCreateBody `json:"sandbox"` +} + +// createResponse is the JSON body returned by POST /api/2.0/lakebox/sandboxes. +// Mirrors the `Sandbox` proto message after JSON transcoding. +// +// `FQDN` is the manager's internal routing hostname — not user-actionable, +// SSH always goes through the gateway. Tagged `omitempty` so the day the +// manager stops returning it, both this struct and downstream `--json` +// output drop the field cleanly instead of leaking a ghost empty string. +type createResponse struct { + SandboxID string `json:"sandboxId"` + Status string `json:"status"` + FQDN string `json:"fqdn,omitempty"` +} + +// sandboxEntry is a single item in the list response. +// Mirrors the `Sandbox` proto message after JSON transcoding. +// +// IdleTimeout and NoAutostop correspond to the proto's `optional` fields; +// they're pointers so we can tell "field absent on the wire" (server has +// the global default) from "explicitly set to 0 / false." +// +// `IdleTimeout` is a `google.protobuf.Duration`. Proto3 JSON canonical +// form serializes Duration as a string with an `s` suffix (e.g. +// `"900s"`), so the Go field is `*string` and we parse on read. +type sandboxEntry struct { + SandboxID string `json:"sandboxId"` + Status string `json:"status"` + FQDN string `json:"fqdn,omitempty"` + Name string `json:"name,omitempty"` + CreateTime string `json:"createTime,omitempty"` + LastStartTime string `json:"lastStartTime,omitempty"` + IdleTimeout *string `json:"idleTimeout,omitempty"` + NoAutostop *bool `json:"noAutostop,omitempty"` +} + +// idleTimeoutSecs parses the proto3-canonical Duration string off +// `IdleTimeout` (e.g. `"900s"` → `900`). Returns 0 when unset or when +// the string is not a recognizable Duration. Sub-second precision is +// dropped — the watchdog only acts on whole seconds. +func (e *sandboxEntry) idleTimeoutSecs() int64 { + if e.IdleTimeout == nil { + return 0 + } + s := *e.IdleTimeout + if !strings.HasSuffix(s, "s") { + return 0 + } + d, err := time.ParseDuration(s) + if err != nil { + return 0 + } + return int64(d.Seconds()) +} + +// defaultAutoStopSecs mirrors the manager's `watchdog_idle_grace_secs` +// fallback (10 minutes) used when a sandbox has no per-record override. +// The value is also documented in `lakebox/CLAUDE.md` ("Sandbox +// Watchdog" section). Hardcoded here so list/status can render the +// effective timeout without an extra round-trip to fetch manager config. +const defaultAutoStopSecs int64 = 600 + +// autoStopLabel renders the auto-stop policy advertised by the manager +// for one sandbox into a short human-readable string. Mirrors the wire +// semantics from `lakebox/proto/lakebox.proto`: +// - `no_autostop == true` → never auto-stops +// - `idle_timeout` set and positive → that many seconds +// - otherwise → manager's global default (`defaultAutoStopSecs`) +func (e *sandboxEntry) autoStopLabel() string { + if e.NoAutostop != nil && *e.NoAutostop { + return "never" + } + if secs := e.idleTimeoutSecs(); secs > 0 { + return formatDurationSecs(secs) + } + return formatDurationSecs(defaultAutoStopSecs) +} + +// formatDurationSecs prints `secs` as a compact duration (e.g. `90s`, +// `15m`, `2h`, `1h30m`). Falls back to seconds if it's not a clean +// minute/hour multiple. Avoids pulling in a dependency just for this. +func formatDurationSecs(secs int64) string { + if secs < 60 { + return fmt.Sprintf("%ds", secs) + } + if secs%3600 == 0 { + return fmt.Sprintf("%dh", secs/3600) + } + if secs >= 3600 { + return fmt.Sprintf("%dh%dm", secs/3600, (secs%3600)/60) + } + if secs%60 == 0 { + return fmt.Sprintf("%dm", secs/60) + } + return fmt.Sprintf("%ds", secs) +} + +// listResponse is the JSON body returned by GET /api/2.0/lakebox/sandboxes. +// `nextPageToken` is empty on the final page (or when the result fits in one). +type listResponse struct { + Sandboxes []sandboxEntry `json:"sandboxes"` + NextPageToken string `json:"nextPageToken,omitempty"` +} + +// listPageSize matches the manager-side default. Typical user fleets are +// well under this, so one round-trip covers them; the pagination loop in +// `list` handles the rare larger fleet. +const listPageSize = 100 + +// updateBody is the PATCH request body. The proto declares +// `UpdateSandboxRequest { Sandbox sandbox = 1 }` with `body: "sandbox"` +// in the (google.api.http) annotation, so the HTTP body is the inner +// `Sandbox` message directly — there is no `{"sandbox": {...}}` +// wrapping on the wire. +// +// Pointer fields encode the proto3 `optional` semantics — only the +// fields we explicitly set are emitted, leaving everything else +// server-untouched. `IdleTimeout` is a proto3-canonical Duration +// string (e.g. `"900s"`); the server-side wire type is +// `google.protobuf.Duration`. +type updateBody struct { + SandboxID string `json:"sandbox_id"` + Name *string `json:"name,omitempty"` + IdleTimeout *string `json:"idle_timeout,omitempty"` + NoAutostop *bool `json:"no_autostop,omitempty"` +} + +// registerKeyRequest is the JSON body for POST /api/2.0/lakebox/ssh-keys. +type registerKeyRequest struct { + PublicKey string `json:"public_key"` + Name string `json:"name,omitempty"` +} + +func newLakeboxAPI(w *databricks.WorkspaceClient) (*lakeboxAPI, error) { + c, err := client.New(w.Config) + if err != nil { + return nil, fmt.Errorf("failed to create lakebox API client: %w", err) + } + return &lakeboxAPI{c: c}, nil +} + +// headers attaches the workspace routing identifier so multi-workspace +// gateways (e.g. SPOG hosts) can scope the credential. Mirrors the pattern +// in libs/telemetry, libs/filer, and SDK-generated workspace services. The +// auth.WorkspaceIDNone sentinel ("none") is treated as unset so the literal +// string never goes on the wire. +func (a *lakeboxAPI) headers() map[string]string { + wsID := a.c.Config.WorkspaceID + if wsID == "" || wsID == auth.WorkspaceIDNone { + return nil + } + return map[string]string{orgIDHeader: wsID} +} + +// create calls POST /api/2.0/lakebox/sandboxes. An empty `name` is omitted +// from the wire payload so the server treats it as "unset" rather than +// "explicit empty string." +func (a *lakeboxAPI) create(ctx context.Context, name string) (*createResponse, error) { + body := createRequest{Sandbox: sandboxCreateBody{Name: name}} + var resp createResponse + err := a.c.Do(ctx, http.MethodPost, lakeboxAPIPath, a.headers(), nil, body, &resp) + if err != nil { + return nil, err + } + return &resp, nil +} + +// list calls GET /api/2.0/lakebox/sandboxes, following pagination until the +// server stops sending `next_page_token`. Returns the full set in one slice. +func (a *lakeboxAPI) list(ctx context.Context) ([]sandboxEntry, error) { + var all []sandboxEntry + pageToken := "" + for { + page, err := a.listPage(ctx, pageToken) + if err != nil { + return nil, err + } + all = append(all, page.Sandboxes...) + if page.NextPageToken == "" { + return all, nil + } + pageToken = page.NextPageToken + } +} + +// listPage fetches a single page of sandboxes. An empty `pageToken` requests +// the first page; the server enforces ordering across pages. +// +// `query` is passed in slot 6 (`request`), not slot 5 (`queryParams`). On +// GET, the SDK's makeRequestBody serializes `request` into the URL query +// string and sends an empty body. Routing through `queryParams` instead +// makes it write a literal `null` body, which the lakebox manager rejects +// with `INVALID_PARAMETER_VALUE: Request body must be a JSON object`. See +// databricks-sdk-go/httpclient/request.go:makeRequestBody. +func (a *lakeboxAPI) listPage(ctx context.Context, pageToken string) (*listResponse, error) { + query := map[string]any{"page_size": listPageSize} + if pageToken != "" { + query["page_token"] = pageToken + } + var resp listResponse + err := a.c.Do(ctx, http.MethodGet, lakeboxAPIPath, a.headers(), nil, query, &resp) + if err != nil { + return nil, err + } + return &resp, nil +} + +// get calls GET /api/2.0/lakebox/sandboxes/{id}. +func (a *lakeboxAPI) get(ctx context.Context, id string) (*sandboxEntry, error) { + var resp sandboxEntry + err := a.c.Do(ctx, http.MethodGet, lakeboxAPIPath+"/"+id, a.headers(), nil, nil, &resp) + if err != nil { + return nil, err + } + return &resp, nil +} + +// update calls PATCH /api/2.0/lakebox/sandboxes/{id} with whichever of +// `idle_timeout` / `no_autostop` the caller chose to set. Fields left +// nil are omitted from the wire payload, so the server preserves their +// current values. Returns the refreshed `sandboxEntry`. +func (a *lakeboxAPI) update(ctx context.Context, id string, name *string, idleTimeoutSecs *int64, noAutostop *bool) (*sandboxEntry, error) { + var idleTimeout *string + if idleTimeoutSecs != nil { + s := fmt.Sprintf("%ds", *idleTimeoutSecs) + idleTimeout = &s + } + body := updateBody{ + SandboxID: id, + Name: name, + IdleTimeout: idleTimeout, + NoAutostop: noAutostop, + } + var resp sandboxEntry + err := a.c.Do(ctx, http.MethodPatch, lakeboxAPIPath+"/"+id, a.headers(), nil, body, &resp) + if err != nil { + return nil, err + } + return &resp, nil +} + +// delete calls DELETE /api/2.0/lakebox/sandboxes/{id}. +func (a *lakeboxAPI) delete(ctx context.Context, id string) error { + return a.c.Do(ctx, http.MethodDelete, lakeboxAPIPath+"/"+id, a.headers(), nil, nil, nil) +} + +// stop calls POST /api/2.0/lakebox/sandboxes/{id}/stop and returns the +// refreshed sandbox. The proto's `StopSandboxRequest` carries `sandbox_id` +// (redundant with the URL path) under `body: "*"`, so we mirror it +// explicitly even though the transcoder fills the field from the path. +func (a *lakeboxAPI) stop(ctx context.Context, id string) (*sandboxEntry, error) { + body := map[string]string{"sandbox_id": id} + var resp sandboxEntry + err := a.c.Do(ctx, http.MethodPost, lakeboxAPIPath+"/"+id+"/stop", a.headers(), nil, body, &resp) + if err != nil { + return nil, err + } + return &resp, nil +} + +// registerKey calls POST /api/2.0/lakebox/ssh-keys. An empty `name` is +// omitted from the wire payload so the server records "unset" rather than +// an explicit empty string. +func (a *lakeboxAPI) registerKey(ctx context.Context, publicKey, name string) error { + return a.c.Do(ctx, http.MethodPost, lakeboxKeysAPIPath, a.headers(), nil, registerKeyRequest{PublicKey: publicKey, Name: name}, nil) +} + +// sshKeyEntry is a single item in the ssh-key list response. Mirrors the +// `SshKey` proto message after JSON transcoding (`key_hash` → `keyHash`, +// timestamps as RFC 3339 strings). +type sshKeyEntry struct { + KeyHash string `json:"keyHash"` + Name string `json:"name,omitempty"` + CreateTime string `json:"createTime,omitempty"` + LastUseTime string `json:"lastUseTime,omitempty"` +} + +// listKeysResponse is the JSON body returned by GET /api/2.0/lakebox/ssh-keys. +// Per-user keys are hard-capped at 100 server-side, so the full set fits in +// one response — no pagination. +type listKeysResponse struct { + SshKeys []sshKeyEntry `json:"sshKeys"` +} + +// listKeys calls GET /api/2.0/lakebox/ssh-keys. +func (a *lakeboxAPI) listKeys(ctx context.Context) ([]sshKeyEntry, error) { + var resp listKeysResponse + err := a.c.Do(ctx, http.MethodGet, lakeboxKeysAPIPath, a.headers(), nil, nil, &resp) + if err != nil { + return nil, err + } + return resp.SshKeys, nil +} + +// deleteKey calls DELETE /api/2.0/lakebox/ssh-keys/{key_hash}. +func (a *lakeboxAPI) deleteKey(ctx context.Context, keyHash string) error { + return a.c.Do(ctx, http.MethodDelete, lakeboxKeysAPIPath+"/"+keyHash, a.headers(), nil, nil, nil) +} diff --git a/cmd/lakebox/config.go b/cmd/lakebox/config.go new file mode 100644 index 0000000000..9f1ef6429a --- /dev/null +++ b/cmd/lakebox/config.go @@ -0,0 +1,152 @@ +package lakebox + +import ( + "errors" + "fmt" + "time" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/spf13/cobra" +) + +// MIN_IDLE_TIMEOUT_SECS / MAX_IDLE_TIMEOUT_SECS mirror the manager-side +// constants in lakebox/src/api/handlers/sandbox.rs. Pre-flighting client-side +// gives a clearer error than waiting for the server's INVALID_ARGUMENT. +const ( + minIdleTimeoutSecs = 60 + maxIdleTimeoutSecs = 86_400 +) + +func newConfigCommand() *cobra.Command { + var idleTimeoutFlag string + var noAutostopFlag bool + var nameFlag string + + cmd := &cobra.Command{ + Use: "config ", + Short: "Configure a Lakebox's name and auto-stop policy", + Long: `Configure a Lakebox's name and auto-stop policy. + +Three knobs are independent — pass any combination: + + --name