From 36073ca26857c393ebe95cdc505a123e846c2b01 Mon Sep 17 00:00:00 2001 From: Akshay Singla Date: Fri, 22 May 2026 04:57:50 +0000 Subject: [PATCH] lakebox: tab completion for sandbox IDs and ssh-key hashes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire `ValidArgsFunction` on every lakebox subcommand that takes a positional argument so shells with cobra completion sourced suggest real IDs instead of falling back to filename completion: - ssh / status / stop / start / delete / config / set-default — sandbox IDs - ssh-key delete — registered key hashes Surfaced through the existing `databricks completion bash|zsh|fish` subcommand. Completion is best-effort — silently returns no suggestions on any error so the shell stays usable. Stacked on top of # (start + ssh fail-fast) since `start` needs to participate in completion too. Co-authored-by: Isaac --- cmd/lakebox/completion.go | 63 +++++++++++++++++++++++++++++++++++++++ cmd/lakebox/config.go | 5 ++-- cmd/lakebox/default.go | 5 ++-- cmd/lakebox/delete.go | 5 ++-- cmd/lakebox/ssh.go | 5 ++-- cmd/lakebox/sshkey.go | 5 ++-- cmd/lakebox/start.go | 5 ++-- cmd/lakebox/status.go | 5 ++-- cmd/lakebox/stop.go | 5 ++-- 9 files changed, 87 insertions(+), 16 deletions(-) create mode 100644 cmd/lakebox/completion.go diff --git a/cmd/lakebox/completion.go b/cmd/lakebox/completion.go new file mode 100644 index 00000000000..c078832b8b8 --- /dev/null +++ b/cmd/lakebox/completion.go @@ -0,0 +1,63 @@ +package lakebox + +import ( + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/spf13/cobra" +) + +// completeSandboxIDs is a Cobra ValidArgsFunction returning the caller's +// sandbox IDs for tab completion. Cobra runs this in a separate process +// from the main command, so we need to bootstrap the workspace client +// ourselves (PreRunE is skipped during completion). +// +// Best-effort: any failure (no auth, no network, lakebox not deployed) +// returns no suggestions instead of an error so the shell stays usable +// and the user can still type the ID by hand. +func completeSandboxIDs(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) { + if len(args) > 0 { + return nil, cobra.ShellCompDirectiveNoFileComp + } + if err := root.MustWorkspaceClient(cmd, args); err != nil { + return nil, cobra.ShellCompDirectiveNoFileComp + } + ctx := cmd.Context() + api, err := newLakeboxAPI(cmdctx.WorkspaceClient(ctx)) + if err != nil { + return nil, cobra.ShellCompDirectiveNoFileComp + } + entries, err := api.list(ctx) + if err != nil { + return nil, cobra.ShellCompDirectiveNoFileComp + } + ids := make([]string, 0, len(entries)) + for _, e := range entries { + ids = append(ids, e.SandboxID) + } + return ids, cobra.ShellCompDirectiveNoFileComp +} + +// completeSSHKeyHashes is the equivalent for `ssh-key delete `, +// returning the hashes of registered keys. Same best-effort semantics. +func completeSSHKeyHashes(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) { + if len(args) > 0 { + return nil, cobra.ShellCompDirectiveNoFileComp + } + if err := root.MustWorkspaceClient(cmd, args); err != nil { + return nil, cobra.ShellCompDirectiveNoFileComp + } + ctx := cmd.Context() + api, err := newLakeboxAPI(cmdctx.WorkspaceClient(ctx)) + if err != nil { + return nil, cobra.ShellCompDirectiveNoFileComp + } + keys, err := api.listKeys(ctx) + if err != nil { + return nil, cobra.ShellCompDirectiveNoFileComp + } + hashes := make([]string, 0, len(keys)) + for _, k := range keys { + hashes = append(hashes, k.KeyHash) + } + return hashes, cobra.ShellCompDirectiveNoFileComp +} diff --git a/cmd/lakebox/config.go b/cmd/lakebox/config.go index 55cc47f0f60..e6cda5c7bc4 100644 --- a/cmd/lakebox/config.go +++ b/cmd/lakebox/config.go @@ -56,8 +56,9 @@ Examples: databricks lakebox config happy-panda-1234 --no-autostop # never auto-stop databricks lakebox config happy-panda-1234 --no-autostop=false # back to timeout path databricks lakebox config happy-panda-1234 --idle-timeout 30m --no-autostop=false`, - Args: cobra.ExactArgs(1), - PreRunE: root.MustWorkspaceClient, + Args: cobra.ExactArgs(1), + PreRunE: root.MustWorkspaceClient, + ValidArgsFunction: completeSandboxIDs, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) diff --git a/cmd/lakebox/default.go b/cmd/lakebox/default.go index cd96df172d1..e637165e796 100644 --- a/cmd/lakebox/default.go +++ b/cmd/lakebox/default.go @@ -18,8 +18,9 @@ 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, + Args: cobra.ExactArgs(1), + PreRunE: root.MustWorkspaceClient, + ValidArgsFunction: completeSandboxIDs, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) diff --git a/cmd/lakebox/delete.go b/cmd/lakebox/delete.go index 001a3822524..83d8d751892 100644 --- a/cmd/lakebox/delete.go +++ b/cmd/lakebox/delete.go @@ -19,8 +19,9 @@ Permanently terminates and removes the specified lakebox. Example: databricks lakebox delete happy-panda-1234`, - Args: cobra.ExactArgs(1), - PreRunE: root.MustWorkspaceClient, + Args: cobra.ExactArgs(1), + PreRunE: root.MustWorkspaceClient, + ValidArgsFunction: completeSandboxIDs, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go index d8d567271ce..82b91cf09df 100644 --- a/cmd/lakebox/ssh.go +++ b/cmd/lakebox/ssh.go @@ -51,8 +51,9 @@ Examples: databricks lakebox ssh -- ls -la /home # run command on default lakebox databricks lakebox ssh happy-panda-1234 -- cat /etc/os-release # run command on specific lakebox databricks lakebox ssh -- -L 8080:localhost:8080 # port forwarding on default lakebox`, - Args: cobra.ArbitraryArgs, - PreRunE: root.MustWorkspaceClient, + Args: cobra.ArbitraryArgs, + PreRunE: root.MustWorkspaceClient, + ValidArgsFunction: completeSandboxIDs, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) diff --git a/cmd/lakebox/sshkey.go b/cmd/lakebox/sshkey.go index 708d5e4b9d9..005a26d871c 100644 --- a/cmd/lakebox/sshkey.go +++ b/cmd/lakebox/sshkey.go @@ -35,8 +35,9 @@ private key will fail until the key is re-registered. Example: databricks lakebox ssh-key delete a1b2c3d4e5f6...`, - Args: cobra.ExactArgs(1), - PreRunE: root.MustWorkspaceClient, + Args: cobra.ExactArgs(1), + PreRunE: root.MustWorkspaceClient, + ValidArgsFunction: completeSSHKeyHashes, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) diff --git a/cmd/lakebox/start.go b/cmd/lakebox/start.go index 9005670d347..6b8b0010a3b 100644 --- a/cmd/lakebox/start.go +++ b/cmd/lakebox/start.go @@ -24,8 +24,9 @@ Starting an already-running sandbox is a no-op. Example: databricks lakebox start happy-panda-1234`, - Args: cobra.ExactArgs(1), - PreRunE: root.MustWorkspaceClient, + Args: cobra.ExactArgs(1), + PreRunE: root.MustWorkspaceClient, + ValidArgsFunction: completeSandboxIDs, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go index 10aef7daa69..16a60a140a2 100644 --- a/cmd/lakebox/status.go +++ b/cmd/lakebox/status.go @@ -21,8 +21,9 @@ func newStatusCommand() *cobra.Command { Example: lakebox status happy-panda-1234 lakebox status happy-panda-1234 --json`, - Args: cobra.ExactArgs(1), - PreRunE: root.MustWorkspaceClient, + Args: cobra.ExactArgs(1), + PreRunE: root.MustWorkspaceClient, + ValidArgsFunction: completeSandboxIDs, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) diff --git a/cmd/lakebox/stop.go b/cmd/lakebox/stop.go index f40d632729a..3c837f9aa33 100644 --- a/cmd/lakebox/stop.go +++ b/cmd/lakebox/stop.go @@ -23,8 +23,9 @@ Stopping an already-stopped sandbox is a no-op. Example: databricks lakebox stop happy-panda-1234`, - Args: cobra.ExactArgs(1), - PreRunE: root.MustWorkspaceClient, + Args: cobra.ExactArgs(1), + PreRunE: root.MustWorkspaceClient, + ValidArgsFunction: completeSandboxIDs, RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx)