diff --git a/cmd/lakebox/completion.go b/cmd/lakebox/completion.go new file mode 100644 index 0000000000..c078832b8b --- /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 55cc47f0f6..e6cda5c7bc 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 cd96df172d..e637165e79 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 001a382252..83d8d75189 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 d8d567271c..82b91cf09d 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 708d5e4b9d..005a26d871 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 9005670d34..6b8b0010a3 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 10aef7daa6..16a60a140a 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 f40d632729..3c837f9aa3 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)