Skip to content

[CJ-75152][lakebox] CLI: fix CreateSandbox wire format + add --name + surface name/timestamps #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
akshaysingla-db wants to merge 1 commit into
base: lakebox-cli
Choose a base branch
from
Open

[CJ-75152][lakebox] CLI: fix CreateSandbox wire format + add --name + surface name/timestamps #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
42 changes: 26 additions & 16 deletions cmd/lakebox/api.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,15 +23,18 @@ type lakeboxAPI struct {
w *databricks.WorkspaceClient
}

// sandboxCreateBody is the inner `Sandbox` message in the create payload.
// Only `name` is caller-settable today (see `LakeboxApi::create_sandbox` in
// `lakebox/src/api/service.rs`); 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.
//
// The proto-defined `CreateSandboxRequest` carries a `Sandbox sandbox = 1`
// field today (every member is server-chosen), but JSON transcoding accepts
// the unwrapped form for forward-compatible callers. Keep `public_key` here
// as a no-op compat shim so older `lakebox create --public-key-file=...`
// invocations don't error — the manager ignores it on the wire.
// `CreateSandboxRequest { Sandbox sandbox = 1 }` has `body: "*"`, so the
// wire body is the full request with a `sandbox` wrapper.
type createRequest struct {
PublicKey string `json:"public_key,omitempty"`
Sandbox sandboxCreateBody `json:"sandbox"`
}

// createResponse is the JSON body returned by POST /api/2.0/lakebox/sandboxes.
Expand All @@ -53,11 +56,14 @@ type createResponse struct {
// 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"`
IdleTimeout *string `json:"idleTimeout,omitempty"`
NoAutostop *bool `json:"noAutostop,omitempty"`
SandboxID string `json:"sandboxId"`
Status string `json:"status"`
FQDN string `json:"fqdn"`
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
Expand Down Expand Up @@ -140,10 +146,11 @@ 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)
// 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) {
jsonBody, err := json.Marshal(createRequest{Sandbox: sandboxCreateBody{Name: name}})
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
}
Expand Down Expand Up @@ -216,6 +223,7 @@ func (a *lakeboxAPI) get(ctx context.Context, id string) (*sandboxEntry, error)
// `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"`
}
Expand All @@ -227,6 +235,7 @@ type updateBody struct {
func (a *lakeboxAPI) update(
ctx context.Context,
id string,
name *string,
idleTimeoutSecs *int64,
noAutostop *bool,
) (*sandboxEntry, error) {
Expand All @@ -237,6 +246,7 @@ func (a *lakeboxAPI) update(
}
body := updateBody{
SandboxID: id,
Name: name,
IdleTimeout: idleTimeout,
NoAutostop: noAutostop,
}
Expand Down
28 changes: 22 additions & 6 deletions cmd/lakebox/config.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,17 @@ const (
func newConfigCommand() *cobra.Command {
var idleTimeoutFlag string
var noAutostopFlag bool
var nameFlag string

cmd := &cobra.Command{
Use: "config <lakebox-id>",
Short: "Configure a Lakebox's auto-stop policy",
Long: `Configure a Lakebox's auto-stop policy.
Short: "Configure a Lakebox's name and auto-stop policy",
Long: `Configure a Lakebox's name and auto-stop policy.

Two knobs are independent — pass either or both:
Three knobs are independent — pass any combination:

--name <label> Display label for the lakebox (max 256 bytes).
Pass an empty string to clear.

--idle-timeout <duration> Per-sandbox idle timeout. The watchdog reaps
the sandbox after this much idle time. Pass
Expand All @@ -42,6 +46,7 @@ Two knobs are independent — pass either or both:
stops on explicit 'lakebox delete'.

Examples:
lakebox config happy-panda-1234 --name my-project
lakebox config happy-panda-1234 --idle-timeout 15m
lakebox config happy-panda-1234 --idle-timeout 1h30m
lakebox config happy-panda-1234 --idle-timeout 0 # clear, use default
Expand Down Expand Up @@ -75,23 +80,34 @@ Examples:
noAutostop = &p
}

if idleSecs == nil && noAutostop == nil {
return fmt.Errorf("nothing to update — pass --idle-timeout and/or --no-autostop")
var name *string
if cmd.Flags().Changed("name") {
n := nameFlag
name = &n
}

if idleSecs == nil && noAutostop == nil && name == nil {
return fmt.Errorf("nothing to update — pass --name, --idle-timeout, and/or --no-autostop")
}

updated, err := api.update(ctx, id, idleSecs, noAutostop)
updated, err := api.update(ctx, id, name, idleSecs, noAutostop)
if err != nil {
return fmt.Errorf("failed to update lakebox %s: %w", id, err)
}

blank(out)
field(out, "id", bold(updated.SandboxID))
if updated.Name != "" {
field(out, "name", updated.Name)
}
field(out, "autostop", dim(updated.autoStopLabel()))
blank(out)
return nil
},
}

cmd.Flags().StringVar(&nameFlag, "name", "",
"Display label for the lakebox (max 256 bytes). Pass --name= to clear.")
cmd.Flags().StringVar(&idleTimeoutFlag, "idle-timeout", "",
"Idle timeout (e.g. 15m, 1h30m, 90s). Pass 0 to clear and revert to the manager's default.")
cmd.Flags().BoolVar(&noAutostopFlag, "no-autostop", false,
Expand Down
21 changes: 6 additions & 15 deletions cmd/lakebox/create.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,13 @@ package lakebox

import (
"fmt"
"os"

"github.com/databricks/cli/libs/cmdctx"
"github.com/spf13/cobra"
)

func newCreateCommand() *cobra.Command {
var publicKeyFile string
var name string

cmd := &cobra.Command{
Use: "create",
Expand All @@ -19,27 +18,19 @@ 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.

Example:
lakebox create`,
Examples:
lakebox create
lakebox create --name my-project`,
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 != "" {
data, err := os.ReadFile(publicKeyFile)
if err != nil {
return fmt.Errorf("failed to read public key file %s: %w", publicKeyFile, err)
}
publicKey = string(data)
}

s := spin(stderr, "Provisioning your lakebox…")

result, err := api.create(ctx, publicKey)
result, err := api.create(ctx, name)
if err != nil {
s.fail("Failed to create lakebox")
return fmt.Errorf("failed to create lakebox: %w", err)
Expand Down Expand Up @@ -73,7 +64,7 @@ Example:
},
}

cmd.Flags().StringVar(&publicKeyFile, "public-key-file", "", "Path to SSH public key file to install in the lakebox")
cmd.Flags().StringVar(&name, "name", "", "Display label for the lakebox (max 256 bytes)")

return cmd
}
6 changes: 1 addition & 5 deletions cmd/lakebox/ssh.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,13 +90,9 @@ Examples:
lakeboxID = def
} 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)
}

s := spin(stderr, "Provisioning your lakebox…")
result, err := api.create(ctx, string(pubKeyData))
result, err := api.create(ctx, "")
if err != nil {
s.fail("Failed to create lakebox")
return fmt.Errorf("failed to create lakebox: %w", err)
Expand Down
3 changes: 3 additions & 0 deletions cmd/lakebox/status.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,9 @@ Example:
out := cmd.OutOrStdout()
blank(out)
field(out, "id", bold(entry.SandboxID))
if entry.Name != "" {
field(out, "name", entry.Name)
}
field(out, "status", status(entry.Status))
if entry.FQDN != "" {
field(out, "fqdn", dim(entry.FQDN))
Expand Down