Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
146 changes: 50 additions & 96 deletions cmd/lakebox/api.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,39 +15,31 @@ import (

// sandboxPath returns the URL path for a single sandbox resource. The ID is
// path-escaped so a value like `foo;rm -rf /` lands on
// `/sandboxes/foo%3Brm%20-rf%20%2F` and gets a clean 400 from
// validate_sandbox_id on the server, rather than its unescaped `/` re-routing
// the request to the list endpoint (which silently returns an empty result the
// CLI then renders as an all-zero sandbox record).
// `/sandboxes/foo%3Brm%20-rf%20%2F` and gets a clean 400 from the server,
// rather than its unescaped `/` re-routing the request to the list endpoint
// (which silently returns an empty result the CLI then renders as an
// all-zero sandbox record).
func sandboxPath(id string) string {
return lakeboxAPIPath + "/" + url.PathEscape(id)
}

// 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"
// Sub-collections under the lakebox service namespace.
const (
lakeboxAPIPath = "/api/2.0/lakebox/sandboxes"
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."
// orgIDHeader scopes the credential to a workspace on multi-workspace
// gateways. Without it, requests fail with "Credential was not sent or was
// of an unsupported type for this API."
const orgIDHeader = "X-Databricks-Org-Id"

// maxNameBytes mirrors the manager-side `Sandbox.name` length cap. The
// server measures bytes, not characters, so a name made of emoji or other
// multi-byte UTF-8 hits the limit in fewer visible characters than the user
// expects. Mirroring the constant client-side lets us fail fast with the
// observed byte count instead of paying a round-trip for a 400.
// maxNameBytes mirrors the server-side `Sandbox.name` cap. The server
// measures bytes (not runes), so emoji hit the limit faster than expected;
// mirroring it client-side lets us fail fast with the observed byte count.
const maxNameBytes = 256

// validateName returns an error when `name` exceeds the wire limit. The
// error names the observed byte count so the user can recover in one shot
// (the original server message just said "exceeds 256 bytes" without saying
// by how much, which is unhelpful when emoji are in play).
// validateName rejects names that exceed the wire limit (counted in bytes).
func validateName(name string) error {
if n := len(name); n > maxNameBytes {
return fmt.Errorf("--name is %d bytes; limit is %d (emoji and most non-ASCII characters count as 2-4 bytes each)", n, maxNameBytes)
Expand All @@ -61,43 +53,32 @@ type lakeboxAPI struct {
}

// sandboxCreateBody is the inner `Sandbox` message in the create payload.
// Only `name` is caller-settable today; all other fields are server-chosen.
// Only `name` is caller-settable; the rest 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.
// createRequest is the wrapped POST body for sandbox creation.
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.
// `GatewayHost` is the public SSH gateway hostname for the workspace,
// stamped by the manager (universe#1966484) so the CLI no longer needs to
// hardcode regional defaults. Both are `omitempty` so old/new wire shapes
// round-trip cleanly.
// createResponse mirrors the Sandbox proto after JSON transcoding. FQDN is
// the manager's internal routing host (not user-actionable); GatewayHost is
// the public SSH gateway. Both are `omitempty` so old and new server
// versions round-trip cleanly.
type createResponse struct {
SandboxID string `json:"sandboxId"`
Status string `json:"status"`
FQDN string `json:"fqdn,omitempty"`
GatewayHost string `json:"gatewayHost,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.
// sandboxEntry mirrors the Sandbox proto after JSON transcoding.
// IdleTimeout and NoAutostop are pointer-typed so we can distinguish
// "field absent on the wire" (server uses its default) from "explicitly
// set to 0 / false". IdleTimeout is a proto3-canonical Duration string
// (see idleTimeoutSecs).
type sandboxEntry struct {
SandboxID string `json:"sandboxId"`
Status string `json:"status"`
Expand Down Expand Up @@ -129,21 +110,12 @@ func (e *sandboxEntry) idleTimeoutSecs() int64 {
return int64(d.Seconds())
}

// 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`:
// autoStopLabel renders the auto-stop policy for one sandbox:
// - `no_autostop == true` → never auto-stops
// - `idle_timeout` set and positive → that many seconds
// - otherwise → no enforcement today; render as "never"
//
// The "otherwise" branch used to render a hardcoded `10m` claiming to
// mirror a manager-side `watchdog_idle_grace_secs` fallback. That
// fallback does not exist in the current tree (only a stale comment in
// `lakebox/proto/lakebox.proto`); the ESM-side `LakeboxChecker` is also
// gated off via the `lakeboxCheckerEnabled` SAFE flag, so unset
// `idle_timeout` is functionally "never auto-stops" today. Once the
// manager enforces a real default, swap this branch back to a duration
// label.
// If the manager later enforces an idle-grace default, render it here.
func (e *sandboxEntry) autoStopLabel() string {
if e.NoAutostop != nil && *e.NoAutostop {
return "never"
Expand All @@ -156,7 +128,7 @@ func (e *sandboxEntry) autoStopLabel() string {

// 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.
// minute/hour multiple.
func formatDurationSecs(secs int64) string {
if secs < 60 {
return fmt.Sprintf("%ds", secs)
Expand All @@ -174,28 +146,17 @@ func formatDurationSecs(secs int64) string {
}

// 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.
// listPageSize matches the manager-side default.
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`.
// updateBody is the PATCH body; the server takes the inner `Sandbox`
// message directly with no `{"sandbox": ...}` wrapping. Pointer fields
// encode proto3 optional semantics (see sandboxEntry).
type updateBody struct {
SandboxID string `json:"sandbox_id"`
Name *string `json:"name,omitempty"`
Expand All @@ -209,6 +170,7 @@ type registerKeyRequest struct {
Name string `json:"name,omitempty"`
}

// newLakeboxAPI returns a lakeboxAPI bound to the workspace client's config.
func newLakeboxAPI(w *databricks.WorkspaceClient) (*lakeboxAPI, error) {
c, err := client.New(w.Config)
if err != nil {
Expand All @@ -218,10 +180,9 @@ func newLakeboxAPI(w *databricks.WorkspaceClient) (*lakeboxAPI, error) {
}

// 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.
// gateways (e.g. SPOG hosts) can scope the credential. 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 {
Expand All @@ -231,8 +192,7 @@ func (a *lakeboxAPI) headers() map[string]string {
}

// 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."
// 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
Expand All @@ -244,7 +204,7 @@ func (a *lakeboxAPI) create(ctx context.Context, name string) (*createResponse,
}

// 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.
// server stops sending `next_page_token`.
func (a *lakeboxAPI) list(ctx context.Context) ([]sandboxEntry, error) {
var all []sandboxEntry
pageToken := ""
Expand All @@ -261,8 +221,7 @@ func (a *lakeboxAPI) list(ctx context.Context) ([]sandboxEntry, error) {
}
}

// listPage fetches a single page of sandboxes. An empty `pageToken` requests
// the first page; the server enforces ordering across pages.
// listPage fetches a single page of sandboxes.
//
// `query` is passed in slot 6 (`request`), not slot 5 (`queryParams`). On
// GET, the SDK's makeRequestBody serializes `request` into the URL query
Expand Down Expand Up @@ -294,9 +253,9 @@ func (a *lakeboxAPI) get(ctx context.Context, id string) (*sandboxEntry, error)
}

// 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`.
// `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 {
Expand All @@ -323,9 +282,7 @@ func (a *lakeboxAPI) delete(ctx context.Context, id string) error {
}

// 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.
// refreshed sandbox.
func (a *lakeboxAPI) stop(ctx context.Context, id string) (*sandboxEntry, error) {
body := map[string]string{"sandbox_id": id}
var resp sandboxEntry
Expand All @@ -337,7 +294,7 @@ func (a *lakeboxAPI) stop(ctx context.Context, id string) (*sandboxEntry, error)
}

// start calls POST /api/2.0/lakebox/sandboxes/{id}/start and returns the
// refreshed sandbox. Mirror of `stop`; same body shape per `body: "*"`.
// refreshed sandbox.
func (a *lakeboxAPI) start(ctx context.Context, id string) (*sandboxEntry, error) {
body := map[string]string{"sandbox_id": id}
var resp sandboxEntry
Expand All @@ -349,15 +306,12 @@ func (a *lakeboxAPI) start(ctx context.Context, id string) (*sandboxEntry, error
}

// 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.
// omitted 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).
// sshKeyEntry is a single item in the ssh-key list response.
type sshKeyEntry struct {
KeyHash string `json:"keyHash"`
Name string `json:"name,omitempty"`
Expand All @@ -366,8 +320,8 @@ type sshKeyEntry struct {
}

// 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.
// Per-user keys are hard-capped server-side, so the full set fits in one
// response — no pagination.
type listKeysResponse struct {
SshKeys []sshKeyEntry `json:"sshKeys"`
}
Expand Down
26 changes: 6 additions & 20 deletions cmd/lakebox/completion.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,19 +6,9 @@ import (
"github.com/spf13/cobra"
)

// completeSandboxIDs is a Cobra ValidArgsFunction returning the caller's
// sandbox IDs (and display names, when distinct from the ID) for tab
// completion. Reads purely from the local cache populated by `lakebox
// list` / `create` / `status` / etc. — no API call on `<TAB>`, so a flaky
// network or unrefreshed auth token never makes the shell hang.
//
// Cobra runs ValidArgsFunction in a separate process from the main
// command, so we still need to bootstrap the workspace client just to
// learn which profile we're under. The cache itself is profile-scoped.
//
// Best-effort: any failure (no profile resolvable, empty cache) returns
// no suggestions instead of an error so the shell stays usable and the
// user can still type the ID by hand.
// completeSandboxIDs returns sandbox IDs and (distinct) display names
// from the local cache for tab completion. Cache-only so an unrefreshed
// token never hangs the shell; any failure yields no suggestions.
func completeSandboxIDs(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
if len(args) > 0 {
return nil, cobra.ShellCompDirectiveNoFileComp
Expand All @@ -31,9 +21,7 @@ func completeSandboxIDs(cmd *cobra.Command, args []string, toComplete string) ([
if len(sbs) == 0 {
return nil, cobra.ShellCompDirectiveNoFileComp
}
// Offer the ID always, plus the display name when the user actually
// set one (server defaults `name` to the sandbox ID, so don't echo
// the same string twice).
// Server defaults `name` to the ID, so only emit the name when it's distinct.
suggestions := make([]string, 0, len(sbs)*2)
for _, s := range sbs {
suggestions = append(suggestions, s.ID)
Expand All @@ -44,10 +32,8 @@ func completeSandboxIDs(cmd *cobra.Command, args []string, toComplete string) ([
return suggestions, cobra.ShellCompDirectiveNoFileComp
}

// completeSSHKeyHashes is the equivalent for `ssh-key delete <hash>`,
// returning the hashes of registered keys. SSH-key hashes aren't cached
// locally (per-user cap is ~100 server-side and listing is cheap), so
// this path still calls the API.
// completeSSHKeyHashes returns registered key hashes for `ssh-key delete`.
// Hashes aren't cached locally, so this path calls the API.
func completeSSHKeyHashes(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
if len(args) > 0 {
return nil, cobra.ShellCompDirectiveNoFileComp
Expand Down
15 changes: 6 additions & 9 deletions cmd/lakebox/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,9 @@ import (
"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.
// minIdleTimeoutSecs / maxIdleTimeoutSecs mirror the server-side bounds
// on `idle_timeout`. Pre-flighting client-side gives a clearer error
// than waiting for the server's INVALID_ARGUMENT.
const (
minIdleTimeoutSecs = 60
maxIdleTimeoutSecs = 86_400
Expand Down Expand Up @@ -78,8 +78,6 @@ Examples:
return err
}

// Translate flag presence + value into the proto3
// optional-field semantics the server expects.
var idleSecs *int64
if cmd.Flags().Changed("idle-timeout") {
secs, err := parseIdleTimeoutFlag(idleTimeoutFlag)
Expand Down Expand Up @@ -157,10 +155,9 @@ func checkIdleSecs(secs int64) (int64, error) {
return 0, nil // clear / revert to global default
}
if secs < minIdleTimeoutSecs || secs > maxIdleTimeoutSecs {
// Format both the bounds and the offending value as Go-style
// durations to match the input form the user typed and the
// flag's --help text (Anwell flagged the prior `86400s` /
// `90000s` echoes as confusing — same unit as input now).
// Echo bounds in Go duration form so they match the input the
// user typed (raw seconds like `86400s` reads as a different
// unit).
return 0, fmt.Errorf(
"idle-timeout must be 0 (clear) or between %s and %s, got %s",
formatDurationSecs(minIdleTimeoutSecs),
Expand Down
6 changes: 2 additions & 4 deletions cmd/lakebox/delete.go
Original file line number Diff line number Diff line change
Expand Up @@ -47,10 +47,8 @@ Examples:
return err
}

// Validate existence first so `delete <typo>` fails clearly
// instead of returning a confident "✓ Removed" on a sandbox
// the server never had — the DELETE endpoint treats 404 as
// idempotent success on the wire.
// DELETE returns success on 404, so pre-check existence to
// surface typos clearly.
entry, err := api.get(ctx, lakeboxID)
if err != nil {
if errors.Is(err, apierr.ErrNotFound) {
Expand Down
6 changes: 0 additions & 6 deletions cmd/lakebox/keyhash.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,13 +10,7 @@ import (
// the first 16 bytes and hex-encoded; the OpenSSH comment (anything after
// the second whitespace-separated token) is stripped before hashing, so
// registering the same key under different comments yields the same hash.
// Inputs that don't have a second token are hashed as-is.
//
// Useful for matching a locally-known key against entries in a
// GET /ssh-keys listing without sending the key contents back to the
// server.
func keyHash(publicKey string) string {
// Slice off the OpenSSH comment by stopping at the second space.
end := len(publicKey)
spaces := 0
for i, c := range publicKey {
Expand Down
Loading
Loading