Add bash shell completion command

dgraph/cmd/root.go

@@ -18,6 +18,7 @@ package cmd
 
 import (
 	"flag"
+	"os"
 	"strings"
 
 	"github.com/dgraph-io/dgraph/dgraph/cmd/alpha"
@@ -105,6 +106,9 @@ func initCmds() {
 		// environment variable.
 		sc.Conf.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
 	}
+	// For bash shell completion
+	RootCmd.AddCommand(shellCompletionCmd())
+
 	cobra.OnInitialize(func() {
 		cfg := rootConf.GetString("config")
 		if cfg == "" {
@@ -145,3 +149,50 @@ func setGlogFlags(conf *viper.Viper) {
 		}
 	}
 }
+
+func shellCompletionCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "completion",
+		Short: "Generates shell completion scripts for bash or zsh",
+	}
+
+	// bash subcommand
+	cmd.AddCommand(&cobra.Command{
+		Use:   "bash",
+		Short: "bash shell completion",
+		Long: `To load bash completion run:
+. < (dgraph completion bash)
+
+To configure your bash shell to load completions for each session,
+add to your bashrc:
+
+# ~/.bashrc or ~/.profile
+. < (dgraph completion bash)
+`,
+		Args: cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return RootCmd.GenBashCompletion(os.Stdout)
+		},
+	})
+
+	// zsh subcommand
+	cmd.AddCommand(&cobra.Command{
+		Use:   "zsh",
+		Short: "zsh shell completion",
+		Long: `To load bash completion run:
+. < (dgraph completion zsh)
+
+To configure your zsh shell to load completions for each session,
+add to your zshrc:
+
+# ~/.zshrc or ~/.profile
+. < (dgraph completion zsh)
+`,
+		Args: cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return RootCmd.GenZshCompletion(os.Stdout)
+		},
+	})
+
+	return cmd
+}