From 7b1efe542896a122b941f89ddc02d416521328d6 Mon Sep 17 00:00:00 2001
From: Johnny Steenbergen <jonathansteenbergen@gmail.com>
Date: Fri, 13 Mar 2020 16:53:31 -0700
Subject: [PATCH] feat(influx): add completion command for users to generate
 the completions for the influx cli

---
 CHANGELOG.md             |  2 ++
 cmd/influx/completion.go | 51 ++++++++++++++++++++++++++++++++++++++++
 cmd/influx/main.go       |  3 +++
 3 files changed, 56 insertions(+)
 create mode 100644 cmd/influx/completion.go

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b89bf5fbfbd..8a866f3e737 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ### Features
 
+1. [17273](https://github.com/influxdata/influxdb/pull/17273): Add shell completions command for the influx cli
+
 ### Bug Fixes
 
 1. [17240](https://github.com/influxdata/influxdb/pull/17240): NodeJS logo displays properly in Firefox
diff --git a/cmd/influx/completion.go b/cmd/influx/completion.go
new file mode 100644
index 00000000000..1c379814b2e
--- /dev/null
+++ b/cmd/influx/completion.go
@@ -0,0 +1,51 @@
+package main
+
+import (
+	"io"
+
+	"github.com/spf13/cobra"
+)
+
+func completionCmd(rootCmd *cobra.Command) *cobra.Command {
+	writeZSH := func(w io.Writer) error {
+		if err := rootCmd.GenZshCompletion(w); err != nil {
+			return err
+		}
+		_, err := io.WriteString(w, "\ncompdef _influx influx\n")
+		return err
+	}
+
+	return &cobra.Command{
+		Use:       "completion [bash|zsh]",
+		Short:     "Generates completion scripts",
+		Args:      cobra.ExactValidArgs(1),
+		ValidArgs: []string{"bash", "zsh", "powershell"},
+		Long: `
+	Outputs shell completion for the given shell (bash or zsh)
+
+	OS X:
+		$ source $(brew --prefix)/etc/bash_completion	# for bash users
+		$ source <(influx completion bash)		# for bash users
+		$ source <(influx completion zsh)		# for zsh users
+
+	Ubuntu:
+		$ source /etc/bash-completion	   # for bash users
+		$ source <(influx completion bash) # for bash users
+		$ source <(influx completion zsh)  # for zsh users
+
+	Additionally, you may want to output the completion to a file and source in your .bashrc/.zshrc
+`,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			writer := rootCmd.OutOrStdout()
+			switch args[0] {
+			case "bash":
+				return rootCmd.GenBashCompletion(writer)
+			case "powershell":
+				return rootCmd.GenPowerShellCompletion(writer)
+			case "zsh":
+				return writeZSH(writer)
+			}
+			return nil
+		},
+	}
+}
diff --git a/cmd/influx/main.go b/cmd/influx/main.go
index 2093e8bfa0d..3754692adb9 100644
--- a/cmd/influx/main.go
+++ b/cmd/influx/main.go
@@ -192,6 +192,9 @@ func (b *cmdInfluxBuilder) cmd(childCmdFns ...func(f *globalFlags, opt genericCL
 		c.Flags().BoolP("help", "h", false, fmt.Sprintf("Help for the %s command ", c.Name()))
 	})
 
+	// completion command goes last, after the walk, so that all
+	// commands have every flag listed in the bash|zsh completions.
+	cmd.AddCommand(completionCmd(cmd))
 	return cmd
 }