Adds autocomplete for tridentctl plugins

Co-authored-by: rushi-net <rushikek@netapp.com>

cli/cmd/completion.go

@@ -0,0 +1,176 @@
+package cmd
+
+import (
+	"bytes"
+	"fmt"
+	"os"
+	"os/exec"
+	"strings"
+
+	"github.com/spf13/cobra"
+)
+
+const (
+	cobraOutputSuffix = `Completion ended with directive: ShellCompDirectiveNoFileComp`
+)
+
+// CLI flags
+var skipPluginDetection bool
+
+// The cobra library comes by default with a completion command that generates completion scripts for bash, zsh, fish, and powershell.
+// But these "standard" completion scripts are not aware of the plugins that are available in the system.
+// This command here overwrites the default cobra completion command and generates bash and zsh completion scripts that are aware of the plugins
+//  and redirect the completion requests to the appropriate plugin binaries.
+
+var completionCmd = &cobra.Command{
+	Use:   "completion [command] ",
+	Short: "Generate completion script ",
+	Long: `Generate the autocompletion script for tridentctl for the specified shell.
+See each sub-command's help for details on how to use the generated script.`,
+}
+
+var completionBashCmd = &cobra.Command{
+	Use:   "bash",
+	Short: "Generate the autocompletion script for bash",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		cobraPoweredPlugins := GetCobraPoweredPlugins()
+
+		if skipPluginDetection || len(cobraPoweredPlugins) == 0 {
+			return cmd.Root().GenBashCompletionV2(os.Stdout, true)
+		}
+
+		var buf bytes.Buffer
+		err := cmd.Root().GenBashCompletionV2(&buf, true)
+		if err != nil {
+			return err
+		}
+
+		patchedCompletionScript := GenPatchedBashCompletionScript(cmd, cobraPoweredPlugins, buf.String())
+		fmt.Print(patchedCompletionScript)
+		return nil
+	},
+}
+
+var completionZshCmd = &cobra.Command{
+	Use:   "zsh",
+	Short: "Generate the autocompletion script for zsh ",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		cobraPoweredPlugins := GetCobraPoweredPlugins()
+
+		if skipPluginDetection || len(cobraPoweredPlugins) == 0 {
+			return cmd.Root().GenZshCompletion(os.Stdout)
+		}
+
+		var buf bytes.Buffer
+		err := cmd.Root().GenZshCompletion(&buf)
+		if err != nil {
+			return err
+		}
+
+		patchedCompletionScript := GenPatchedZshCompletionScript(cmd, cobraPoweredPlugins, buf.String())
+		fmt.Print(patchedCompletionScript)
+		return nil
+	},
+}
+
+var completionFishCmd = &cobra.Command{
+	Use:   "fish",
+	Short: "Generate the autocompletion script for fish",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		return cmd.Root().GenFishCompletion(os.Stdout, true)
+	},
+}
+
+var completionPowershellCmd = &cobra.Command{
+	Use:   "powershell",
+	Short: "Generate the autocompletion script for powershell",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		return cmd.Root().GenPowerShellCompletionWithDesc(os.Stdout)
+	},
+}
+
+func GenPatchedBashCompletionScript(cmd *cobra.Command, cobraPoweredPlugins []string, originalScript string) string {
+	searchString := `requestComp="${words[0]} __complete ${args[*]}"`
+	ifReplaceString := `if [[ ${args[0]} == "%s" ]] ; then
+        requestComp="tridentctl-%s __complete ${args[*]:1}"`
+	elseReplaceString := `else
+        requestComp="${words[0]} __complete ${args[*]}"
+    fi`
+
+	return patchFunction(originalScript, searchString, ifReplaceString, elseReplaceString, cobraPoweredPlugins)
+}
+
+func GenPatchedZshCompletionScript(cmd *cobra.Command, cobraPoweredPlugins []string, originalScript string) string {
+	searchString := `requestComp="${words[1]} __complete ${words[2,-1]}"`
+	ifReplaceString := `if [[ ${#words[@]} -ge 3 && "${words[2]}" == "%s" ]] ; then
+        requestComp="tridentctl-%s __complete ${words[3,-1]}"`
+
+	elseReplaceString := `else
+        requestComp="${words[1]} __complete ${words[2,-1]}"
+    fi`
+
+	return patchFunction(originalScript, searchString, ifReplaceString, elseReplaceString, cobraPoweredPlugins)
+}
+
+// patchFunction creates if-elif-...-else block in the completion script
+// seems to be generic enough to work for both zsh and bash
+func patchFunction(originalScript, searchString, ifReplaceString, elseReplaceString string, cobraPoweredPlugins []string) string {
+	fullReplaceString := "# Patch for cobra-powered plugins"
+
+	i := 0
+	for _, plugin := range cobraPoweredPlugins {
+		prefix := "    "
+		if i > 0 {
+			// make it an elif
+			prefix += "el"
+		}
+		fullReplaceString += "\n" + prefix + fmt.Sprintf(ifReplaceString, plugin, plugin)
+		i++
+	}
+	fullReplaceString += "\n    " + elseReplaceString
+
+	return strings.Replace(originalScript, searchString, fullReplaceString, 1)
+}
+
+// GetCobraPoweredPlugins collect cobra-powered plugins
+// to check if the plugin is a cobra plugin we run the plugin with argument "__complete -- @@@@"
+// and compare the output with the cobraOutputSuffix string
+func GetCobraPoweredPlugins() []string {
+	if skipPluginDetection {
+		return nil
+	}
+
+	cobraPoweredPlugins := make([]string, 0)
+	for plugin := range pluginsFound {
+		// run command and compare output
+		pluginCommand := tridentctlPrefix + plugin
+
+		cmd := exec.Command(pluginCommand, "__complete", "--", "@@@@") // @@@@ is random string unlikely to be a valid prefix
+
+		// Run command and capture output
+		output, err := cmd.CombinedOutput()
+		if err != nil {
+			fmt.Fprintln(os.Stderr, err)
+			return nil
+		}
+
+		// Compare output with string
+		if strings.HasSuffix(strings.TrimSpace(string(output)), cobraOutputSuffix) {
+			cobraPoweredPlugins = append(cobraPoweredPlugins, plugin)
+		} else {
+			fmt.Fprintf(os.Stderr, "Plugin %s does not seem to be a cobra plugin\n", plugin)
+		}
+	}
+	return cobraPoweredPlugins
+}
+
+func init() {
+	completionCmd.AddCommand(completionBashCmd)
+	completionCmd.AddCommand(completionZshCmd)
+	completionCmd.AddCommand(completionFishCmd)
+	completionCmd.AddCommand(completionPowershellCmd)
+
+	RootCmd.AddCommand(completionCmd)
+	completionCmd.PersistentFlags().BoolVar(&skipPluginDetection, "skip-plugin-detection", false,
+		"Create completion script without checking for cobra-powered plugins")
+}

cli/cmd/completion_test.go

@@ -0,0 +1,57 @@
+package cmd
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGenPatchedBashCompletionScript(t *testing.T) {
+	dummyOriginalScript := `    args=("${words[@]:1}")
+    requestComp="${words[0]} __complete ${args[*]}"
+
+    lastParam=${words[$((${#words[@]}-1))]}
+`
+
+	expectedResult := `    args=("${words[@]:1}")
+    # Patch for cobra-powered plugins
+    if [[ ${args[0]} == "plugin1" ]] ; then
+        requestComp="tridentctl-plugin1 __complete ${args[*]:1}"
+    elif [[ ${args[0]} == "plugin2" ]] ; then
+        requestComp="tridentctl-plugin2 __complete ${args[*]:1}"
+    else
+        requestComp="${words[0]} __complete ${args[*]}"
+    fi
+
+    lastParam=${words[$((${#words[@]}-1))]}
+`
+
+	cobraPoweredPlugins := []string{"plugin1", "plugin2"}
+	result := GenPatchedBashCompletionScript(completionCmd, cobraPoweredPlugins, dummyOriginalScript)
+
+	assert.Equal(t, expectedResult, result)
+}
+
+func TestGenPatchedZshCompletionScript(t *testing.T) {
+	dummyOriginalScript := `    # Prepare the command to obtain completions
+    requestComp="${words[1]} __complete ${words[2,-1]}"
+    if [ "${lastChar}" = "" ]; then
+`
+
+	expectedResult := `    # Prepare the command to obtain completions
+    # Patch for cobra-powered plugins
+    if [[ ${#words[@]} -ge 3 && "${words[2]}" == "plugin1" ]] ; then
+        requestComp="tridentctl-plugin1 __complete ${words[3,-1]}"
+    elif [[ ${#words[@]} -ge 3 && "${words[2]}" == "plugin2" ]] ; then
+        requestComp="tridentctl-plugin2 __complete ${words[3,-1]}"
+    else
+        requestComp="${words[1]} __complete ${words[2,-1]}"
+    fi
+    if [ "${lastChar}" = "" ]; then
+`
+
+	cobraPoweredPlugins := []string{"plugin1", "plugin2"}
+	result := GenPatchedZshCompletionScript(completionCmd, cobraPoweredPlugins, dummyOriginalScript)
+
+	assert.Equal(t, expectedResult, result)
+}

cli/cmd/plugins.go

@@ -2,7 +2,6 @@ package cmd
 
 import (
 	"fmt"
-	"log"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -17,13 +16,40 @@ const (
 	pluginPathEnvVar = "TRIDENTCTL_PLUGIN_PATH"
 )
 
+var (
+	pluginsFound       = make(map[string]struct{})
+	invalidPluginNames = map[string]struct{}{
+		"completion": {},
+		"create":     {},
+		"delete":     {},
+		"get":        {},
+		"help":       {},
+		"images":     {},
+		"import":     {},
+		"install":    {},
+		"logs":       {},
+		"send":       {},
+		"uninstall":  {},
+		"update":     {},
+		"version":    {},
+	}
+)
+
 func init() {
 	// search PATH for all binaries starting with "tridentctl-"
 	binaries := findBinaries(tridentctlPrefix)
 
 	plugins := make([]*cobra.Command, 0, len(binaries))
 	for _, binary := range binaries {
 		pluginName := strings.TrimPrefix(filepath.Base(binary), tridentctlPrefix)
+		if _, ok := invalidPluginNames[pluginName]; ok {
+			continue
+		}
+		if _, ok := pluginsFound[pluginName]; ok {
+			// Skip duplicates
+			continue
+		}
+
 		pluginCmd := &cobra.Command{
 			Use:                pluginName,
 			Short:              "Run the " + pluginName + " plugin",
@@ -46,7 +72,7 @@ func init() {
 			},
 		}
 		plugins = append(plugins, pluginCmd)
-
+		pluginsFound[pluginName] = struct{}{}
 	}
 
 	if len(plugins) == 0 {
@@ -74,7 +100,7 @@ func findBinaries(prefix string) []string {
 	for _, dir := range strings.Split(path, ":") {
 		files, err := filepath.Glob(filepath.Join(dir, prefix+"*"))
 		if err != nil {
-			log.Println(err)
+			fmt.Fprintln(os.Stderr, err)
 			continue
 		}
 		for _, file := range files {