cli: provide zsh autocompletion script gen command

pkg/cli/cli_test.go

@@ -2039,18 +2039,43 @@ func TestGenAutocomplete(t *testing.T) {
 		}
 	}()
 
-	const minsize = 25000
-	acpath := filepath.Join(acdir, "cockroach.bash")
+	for _, tc := range []struct {
+		shell  string
+		expErr string
+	}{
+		{shell: ""},
+		{shell: "bash"},
+		{shell: "zsh"},
+		{shell: "bad", expErr: `invalid argument "bad" for "cockroach gen autocomplete"`},
+	} {
+		t.Run("shell="+tc.shell, func(t *testing.T) {
+			const minsize = 1000
+			acpath := filepath.Join(acdir, "output-"+tc.shell)
 
-	if err := Run([]string{"gen", "autocomplete", "--out=" + acpath}); err != nil {
-		t.Fatal(err)
-	}
-	info, err := os.Stat(acpath)
-	if err != nil {
-		t.Fatal(err)
-	}
-	if size := info.Size(); size < minsize {
-		t.Fatalf("autocomplete file size (%d) < minimum (%d)", size, minsize)
+			args := []string{"gen", "autocomplete", "--out=" + acpath}
+			if len(tc.shell) > 0 {
+				args = append(args, tc.shell)
+			}
+			err := Run(args)
+			if tc.expErr == "" {
+				if err != nil {
+					t.Fatal(err)
+				}
+			} else {
+				if !testutils.IsError(err, tc.expErr) {
+					t.Fatalf("expected error %s, found %v", tc.expErr, err)
+				}
+				return
+			}
+
+			info, err := os.Stat(acpath)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if size := info.Size(); size < minsize {
+				t.Fatalf("autocomplete file size (%d) < minimum (%d)", size, minsize)
+			}
+		})
 	}
 }
 
@@ -2064,7 +2089,6 @@ func TestJunkPositionalArguments(t *testing.T) {
 		"start",
 		"sql",
 		"gen man",
-		"gen autocomplete",
 		"gen example-data intro",
 	} {
 		const junk = "junk"

pkg/cli/gen.go

@@ -31,9 +31,10 @@ var genManCmd = &cobra.Command{
 	Short: "generate man pages for CockroachDB",
 	Long: `This command generates man pages for CockroachDB.
 
-By default, this places man pages into the "man/man1" directory under the current directory.
-Use "--path=PATH" to override the output directory. For example, to install man pages globally on
-many Unix-like systems, use "--path=/usr/local/share/man/man1".
+By default, this places man pages into the "man/man1" directory under the
+current directory. Use "--path=PATH" to override the output directory. For
+example, to install man pages globally on many Unix-like systems,
+use "--path=/usr/local/share/man/man1".
 `,
 	Args: cobra.NoArgs,
 	RunE: MaybeDecorateGRPCError(runGenManCmd),
@@ -76,26 +77,49 @@ func runGenManCmd(cmd *cobra.Command, args []string) error {
 var autoCompletePath string
 
 var genAutocompleteCmd = &cobra.Command{
-	Use:   "autocomplete",
-	Short: "generate bash autocompletion script for CockroachDB",
-	Long: `Generate bash autocompletion script for CockroachDB. This takes an optional parameter
-to specify the path for the completion file.
+	Use:   "autocomplete [shell]",
+	Short: "generate autocompletion script for CockroachDB",
+	Long: `Generate autocompletion script for CockroachDB.
 
-By default, completion file is written to ./cockroach.bash. Use "--out=/path/to/file" to
-override the file location.
+If no arguments are passed, or if 'bash' is passed, a bash completion file is
+written to ./cockroach.bash. If 'zsh' is passed, a zsh completion file is written
+to ./_cockroach. Use "--out=/path/to/file" to override the output file location.
 
-Note that for the generated file to work on OS X, you'll need to install Homebrew's bash-completion
-package (or an equivalent) and follow the post-install instructions.
+Note that for the generated file to work on OS X with bash, you'll need to install
+Homebrew's bash-completion package (or an equivalent) and follow the post-install
+instructions.
 `,
-	Args: cobra.NoArgs,
-	RunE: MaybeDecorateGRPCError(runGenAutocompleteCmd),
+	Args:      cobra.OnlyValidArgs,
+	ValidArgs: []string{"bash", "zsh"},
+	RunE:      MaybeDecorateGRPCError(runGenAutocompleteCmd),
 }
 
 func runGenAutocompleteCmd(cmd *cobra.Command, args []string) error {
-	if err := cmd.Root().GenBashCompletionFile(autoCompletePath); err != nil {
-		return err
+	var shell string
+	if len(args) > 0 {
+		shell = args[0]
+	} else {
+		shell = "bash"
+	}
+
+	var err error
+	switch shell {
+	case "bash":
+		if autoCompletePath == "" {
+			autoCompletePath = "cockroach.bash"
+		}
+		err = cmd.Root().GenBashCompletionFile(autoCompletePath)
+	case "zsh":
+		if autoCompletePath == "" {
+			autoCompletePath = "_cockroach"
+		}
+		err = cmd.Root().GenZshCompletionFile(autoCompletePath)
 	}
-	fmt.Println("Generated bash completion file", autoCompletePath)
+	if err != nil {
+		return nil
+	}
+
+	fmt.Printf("Generated %s completion file: %s\n", shell, autoCompletePath)
 	return nil
 }
 
@@ -118,7 +142,7 @@ var genCmds = []*cobra.Command{
 func init() {
 	genManCmd.PersistentFlags().StringVar(&manPath, "path", "man/man1",
 		"path where man pages will be outputted")
-	genAutocompleteCmd.PersistentFlags().StringVar(&autoCompletePath, "out", "cockroach.bash",
+	genAutocompleteCmd.PersistentFlags().StringVar(&autoCompletePath, "out", "",
 		"path to generated autocomplete file")
 	genHAProxyCmd.PersistentFlags().StringVar(&haProxyPath, "out", "haproxy.cfg",
 		"path to generated haproxy configuration file")