From 3e5561daca4128d8e310b22a777052ff0f1044f8 Mon Sep 17 00:00:00 2001 From: Derek Nola Date: Fri, 29 Apr 2022 12:53:34 -0700 Subject: [PATCH] Add new `k3s completion` command for shell completion (#5461) * Add shell completion CLI Signed-off-by: Derek Nola --- cmd/completion/main.go | 23 +++++++ cmd/k3s/main.go | 6 ++ cmd/server/main.go | 2 + main.go | 2 + pkg/cli/cmds/completion.go | 20 ++++++ pkg/cli/completion/completion.go | 103 +++++++++++++++++++++++++++++++ scripts/build | 2 + scripts/package-cli | 2 +- 8 files changed, 159 insertions(+), 1 deletion(-) create mode 100644 cmd/completion/main.go create mode 100644 pkg/cli/cmds/completion.go create mode 100644 pkg/cli/completion/completion.go diff --git a/cmd/completion/main.go b/cmd/completion/main.go new file mode 100644 index 0000000000..a49107c308 --- /dev/null +++ b/cmd/completion/main.go @@ -0,0 +1,23 @@ +package main + +import ( + "context" + "errors" + "os" + + "github.com/k3s-io/k3s/pkg/cli/cmds" + "github.com/k3s-io/k3s/pkg/cli/completion" + "github.com/sirupsen/logrus" + "github.com/urfave/cli" +) + +func main() { + app := cmds.NewApp() + app.Commands = []cli.Command{ + cmds.NewCompletionCommand(completion.Run), + } + + if err := app.Run(os.Args); err != nil && !errors.Is(err, context.Canceled) { + logrus.Fatal(err) + } +} diff --git a/cmd/k3s/main.go b/cmd/k3s/main.go index 73aee27d96..a37b8e3daa 100644 --- a/cmd/k3s/main.go +++ b/cmd/k3s/main.go @@ -40,6 +40,7 @@ func main() { // Handle subcommand invocation (k3s server, k3s crictl, etc) app := cmds.NewApp() + app.EnableBashCompletion = true app.Commands = []cli.Command{ cmds.NewServerCommand(internalCLIAction(version.Program+"-server", dataDir, os.Args)), cmds.NewAgentCommand(internalCLIAction(version.Program+"-agent", dataDir, os.Args)), @@ -67,6 +68,7 @@ func main() { cmds.NewCertSubcommands( certCommand), ), + cmds.NewCompletionCommand(internalCLIAction(version.Program+"-completion", dataDir, os.Args)), } if err := app.Run(os.Args); err != nil && !errors.Is(err, context.Canceled) { @@ -135,6 +137,10 @@ func externalCLI(cli, dataDir string, args []string) error { // internalCLIAction returns a function that will call a K3s internal command, be used as the Action of a cli.Command. func internalCLIAction(cmd, dataDir string, args []string) func(ctx *cli.Context) error { return func(ctx *cli.Context) error { + // We don't want the Info logs seen when printing the autocomplete script + if cmd == "k3s-completion" { + logrus.SetLevel(logrus.ErrorLevel) + } return stageAndRunCLI(ctx, cmd, dataDir, args) } } diff --git a/cmd/server/main.go b/cmd/server/main.go index e96d214039..f5e2c3cda8 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -10,6 +10,7 @@ import ( "github.com/k3s-io/k3s/pkg/cli/agent" "github.com/k3s-io/k3s/pkg/cli/cert" "github.com/k3s-io/k3s/pkg/cli/cmds" + "github.com/k3s-io/k3s/pkg/cli/completion" "github.com/k3s-io/k3s/pkg/cli/crictl" "github.com/k3s-io/k3s/pkg/cli/ctr" "github.com/k3s-io/k3s/pkg/cli/etcdsnapshot" @@ -67,6 +68,7 @@ func main() { cmds.NewCertSubcommands( cert.Run), ), + cmds.NewCompletionCommand(completion.Run), } if err := app.Run(configfilearg.MustParse(os.Args)); err != nil && !errors.Is(err, context.Canceled) { diff --git a/main.go b/main.go index dfd910fabc..8fbeeeb923 100644 --- a/main.go +++ b/main.go @@ -14,6 +14,7 @@ import ( "github.com/k3s-io/k3s/pkg/cli/agent" "github.com/k3s-io/k3s/pkg/cli/cert" "github.com/k3s-io/k3s/pkg/cli/cmds" + "github.com/k3s-io/k3s/pkg/cli/completion" "github.com/k3s-io/k3s/pkg/cli/crictl" "github.com/k3s-io/k3s/pkg/cli/etcdsnapshot" "github.com/k3s-io/k3s/pkg/cli/kubectl" @@ -51,6 +52,7 @@ func main() { cmds.NewCertSubcommands( cert.Run), ), + cmds.NewCompletionCommand(completion.Run), } if err := app.Run(configfilearg.MustParse(os.Args)); err != nil && !errors.Is(err, context.Canceled) { diff --git a/pkg/cli/cmds/completion.go b/pkg/cli/cmds/completion.go new file mode 100644 index 0000000000..2a5fdee678 --- /dev/null +++ b/pkg/cli/cmds/completion.go @@ -0,0 +1,20 @@ +package cmds + +import ( + "github.com/urfave/cli" +) + +func NewCompletionCommand(action func(*cli.Context) error) cli.Command { + return cli.Command{ + Name: "completion", + Usage: "Install shell completion script", + UsageText: appName + " completion [SHELL] (valid shells: bash, zsh)", + Action: action, + Flags: []cli.Flag{ + cli.BoolFlag{ + Name: "i", + Usage: "Install source line to rc file", + }, + }, + } +} diff --git a/pkg/cli/completion/completion.go b/pkg/cli/completion/completion.go new file mode 100644 index 0000000000..29e54158a4 --- /dev/null +++ b/pkg/cli/completion/completion.go @@ -0,0 +1,103 @@ +package completion + +import ( + "fmt" + "os" + + "github.com/k3s-io/k3s/pkg/version" + + "github.com/urfave/cli" +) + +func Run(ctx *cli.Context) error { + if ctx.NArg() < 1 { + return fmt.Errorf("must provide a valid SHELL argument") + } + shell := ctx.Args()[0] + completetionScript, err := genCompletionScript(shell) + if err != nil { + return err + } + if ctx.Bool("i") { + return writeToRC(shell) + } + fmt.Println(completetionScript) + return nil +} + +func genCompletionScript(shell string) (string, error) { + var completionScript string + if shell == "bash" { + completionScript = fmt.Sprintf(`#! /bin/bash +_cli_bash_autocomplete() { +if [[ "${COMP_WORDS[0]}" != "source" ]]; then + local cur opts base + COMPREPLY=() + cur="${COMP_WORDS[COMP_CWORD]}" + if [[ "$cur" == "-"* ]]; then + opts=$( ${COMP_WORDS[@]:0:$COMP_CWORD} ${cur} --generate-bash-completion ) + else + opts=$( ${COMP_WORDS[@]:0:$COMP_CWORD} --generate-bash-completion ) + fi + COMPREPLY=( $(compgen -W "${opts}" -- ${cur}) ) + return 0 +fi +} + +complete -o bashdefault -o default -o nospace -F _cli_bash_autocomplete %s +`, version.Program) + } else if shell == "zsh" { + completionScript = fmt.Sprintf(`#compdef %[1]s +_cli_zsh_autocomplete() { + + local -a opts + local cur + cur=${words[-1]} + if [[ "$cur" == "-"* ]]; then + opts=("${(@f)$(_CLI_ZSH_AUTOCOMPLETE_HACK=1 ${words[@]:0:#words[@]-1} ${cur} --generate-bash-completion)}") + else + opts=("${(@f)$(_CLI_ZSH_AUTOCOMPLETE_HACK=1 ${words[@]:0:#words[@]-1} --generate-bash-completion)}") + fi + + if [[ "${opts[1]}" != "" ]]; then + _describe 'values' opts + else + _files + fi + + return +} + +compdef _cli_zsh_autocomplete %[1]s`, version.Program) + } else { + return "", fmt.Errorf("unknown shell: %s", shell) + } + + return completionScript, nil +} + +func writeToRC(shell string) error { + rcFileName := "" + if shell == "bash" { + rcFileName = "/.bashrc" + } else if shell == "zsh" { + rcFileName = "/.zshrc" + } + + home, err := os.UserHomeDir() + if err != nil { + return nil + } + rcFileName = home + rcFileName + f, err := os.OpenFile(rcFileName, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) + if err != nil { + return err + } + defer f.Close() + bashEntry := fmt.Sprintf("# >> %[1]s command completion (start)\n. <(%[1]s completion %s)\n# >> %[1]s command completion (end)", version.Program, shell) + if _, err := f.WriteString(bashEntry); err != nil { + return err + } + fmt.Printf("Autocomplete for %s added to: %s\n", shell, rcFileName) + return nil +} diff --git a/scripts/build b/scripts/build index d3ccd84286..d8c27c1400 100755 --- a/scripts/build +++ b/scripts/build @@ -82,6 +82,7 @@ rm -f \ bin/k3s-etcd-snapshot \ bin/k3s-secrets-encrypt \ bin/k3s-certificate \ + bin/k3s-completion \ bin/kubectl \ bin/crictl \ bin/ctr \ @@ -116,6 +117,7 @@ ln -s k3s ./bin/k3s-server ln -s k3s ./bin/k3s-etcd-snapshot ln -s k3s ./bin/k3s-secrets-encrypt ln -s k3s ./bin/k3s-certificate +ln -s k3s ./bin/k3s-completion ln -s k3s ./bin/kubectl ln -s k3s ./bin/crictl ln -s k3s ./bin/ctr diff --git a/scripts/package-cli b/scripts/package-cli index c96a2cae51..2a78d172ae 100755 --- a/scripts/package-cli +++ b/scripts/package-cli @@ -7,7 +7,7 @@ cd $(dirname $0)/.. GO=${GO-go} -for i in crictl kubectl k3s-agent k3s-server k3s-etcd-snapshot k3s-secrets-encrypt k3s-certificate; do +for i in crictl kubectl k3s-agent k3s-server k3s-etcd-snapshot k3s-secrets-encrypt k3s-certificate k3s-completion; do rm -f bin/$i ln -s k3s bin/$i done