Merge pull request #1142 from plebioda/completion-scripts

Add support for shell completion

Author: boris-ilijic

Makefile

@@ -1,4 +1,4 @@
-.PHONY: build-pbm build-agent build install install-pbm install-agent test
+.PHONY: build-pbm build-agent build install install-pbm install-agent test completion completion-bash completion-zsh
 
 GOOS?=$(shell go env GOOS)
 GOMOD?=on
@@ -129,3 +129,28 @@ install-agent-cover:
 	$(ENVS) go install -cover -ldflags="$(LDFLAGS)" $(BUILD_FLAGS) ./cmd/pbm-agent
 install-stest-cover:
 	$(ENVS) go install -cover -ldflags="$(LDFLAGS)" $(BUILD_FLAGS) ./cmd/pbm-speed-test
+
+# COMPLETION SCRIPTS
+completion: completion-bash completion-zsh
+
+completion-dir-bash:
+	mkdir -p ./bin/completions/bash
+
+completion-dir-zsh:
+	mkdir -p ./bin/completions/zsh
+
+completion-bash: completion-dir-bash completion-bash-pbm completion-bash-agent completion-bash-stest
+completion-bash-pbm:
+	./bin/pbm completion bash > ./bin/completions/bash/pbm
+completion-bash-agent:
+	./bin/pbm-agent completion bash > ./bin/completions/bash/pbm-agent
+completion-bash-stest:
+	./bin/pbm-speed-test completion bash > ./bin/completions/bash/pbm-speed-test
+
+completion-zsh: completion-dir-zsh completion-zsh-pbm completion-zsh-agent completion-zsh-stest
+completion-zsh-pbm:
+	./bin/pbm completion zsh > ./bin/completions/zsh/_pbm
+completion-zsh-agent:
+	./bin/pbm-agent completion zsh > ./bin/completions/zsh/_pbm-agent
+completion-zsh-stest:
+	./bin/pbm-speed-test completion zsh > ./bin/completions/zsh/_pbm-speed-test

cmd/pbm-agent/main.go

@@ -17,6 +17,7 @@ import (
 	"github.com/percona/percona-backup-mongodb/pbm/connect"
 	"github.com/percona/percona-backup-mongodb/pbm/errors"
 	"github.com/percona/percona-backup-mongodb/pbm/log"
+	"github.com/percona/percona-backup-mongodb/pbm/util"
 	"github.com/percona/percona-backup-mongodb/pbm/version"
 )
 
@@ -25,6 +26,7 @@ const mongoConnFlag = "mongodb-uri"
 func main() {
 	rootCmd := rootCommand()
 	rootCmd.AddCommand(versionCommand())
+	rootCmd.AddCommand(util.CompletionCommand())
 
 	if err := rootCmd.Execute(); err != nil {
 		fmt.Println(err)

cmd/pbm-speed-test/main.go

@@ -158,6 +158,7 @@ func main() {
 	versionCmd.Flags().StringVar(&versionFormat, "format", "", "Output format <json or \"\">")
 
 	rootCmd.AddCommand(versionCmd)
+	rootCmd.AddCommand(util.CompletionCommand())
 
 	if err := rootCmd.Execute(); err != nil {
 		os.Exit(1)

cmd/pbm/main.go

@@ -20,6 +20,7 @@ import (
 	"github.com/percona/percona-backup-mongodb/pbm/log"
 	"github.com/percona/percona-backup-mongodb/pbm/oplog"
 	"github.com/percona/percona-backup-mongodb/pbm/topo"
+	"github.com/percona/percona-backup-mongodb/pbm/util"
 	"github.com/percona/percona-backup-mongodb/pbm/version"
 	"github.com/percona/percona-backup-mongodb/sdk"
 )
@@ -148,14 +149,15 @@ func newPbmApp() *pbmApp {
 	app.rootCmd.AddCommand(app.buildRestoreFinishCmd())
 	app.rootCmd.AddCommand(app.buildStatusCmd())
 	app.rootCmd.AddCommand(app.buildVersionCmd())
+	app.rootCmd.AddCommand(util.CompletionCommand())
 
 	return app
 }
 
 func (app *pbmApp) persistentPreRun(cmd *cobra.Command, args []string) error {
 	app.pbmOutF = outFormat(viper.GetString("out"))
 
-	if cmd.Name() == "help" || cmd.Name() == "version" {
+	if cmd.Name() == "help" || cmd.Name() == "version" || util.IsCompletionCommand(cmd.Name()) {
 		return nil
 	}
 

e2e-tests/docker/pbm.dockerfile

@@ -9,7 +9,7 @@ WORKDIR /build
 RUN mkdir -p /data/db
 
 COPY --from=mongo_image /bin/mongod /bin/
-RUN dnf install epel-release && dnf update && dnf install make gcc krb5-devel iproute-tc libfaketime
+RUN dnf install epel-release && dnf update && dnf install make gcc krb5-devel iproute-tc libfaketime bash-completion
 
 RUN arch=$(arch | sed s/aarch64/arm64/ | sed s/x86_64/amd64/) && \
 curl -sL -o /tmp/golang.tar.gz https://go.dev/dl/go1.23.2.linux-${arch}.tar.gz && \
@@ -21,3 +21,5 @@ FROM base-build
 COPY . .
 
 RUN make build-tests && cp /build/bin/* /bin/
+RUN pbm completion bash > /etc/bash_completion.d/pbm
+RUN useradd -u 1001 user

pbm/util/completion.go

@@ -0,0 +1,31 @@
+package util
+
+import (
+	"os"
+
+	"github.com/spf13/cobra"
+)
+
+func CompletionCommand() *cobra.Command {
+	return &cobra.Command{
+		Use:                   "completion [bash|zsh]",
+		Short:                 "Generate completion script for the specified shell",
+		DisableFlagsInUseLine: true,
+		Args:                  cobra.MatchAll(cobra.ExactArgs(1), cobra.OnlyValidArgs),
+		ValidArgs:             []string{"bash", "zsh"},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			switch args[0] {
+			case "bash":
+				return cmd.Root().GenBashCompletion(os.Stdout)
+			case "zsh":
+				return cmd.Root().GenZshCompletion(os.Stdout)
+			default:
+				return cmd.Help()
+			}
+		},
+	}
+}
+
+func IsCompletionCommand(name string) bool {
+	return name == "completion" || name == "__complete"
+}