Add new `k3s completion` command for shell completion (#5461)

cmd/completion/main.go

+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)
+	}
+}

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)
 	}
 }

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) {

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) {

pkg/cli/cmds/completion.go

+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",
+			},
+		},
+	}
+}

pkg/cli/completion/completion.go

+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
+}

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

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