diff --git a/cli/cmd/completion.go b/cli/cmd/completion.go new file mode 100644 index 000000000..bf75d4e7d --- /dev/null +++ b/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") +} diff --git a/cli/cmd/completion_test.go b/cli/cmd/completion_test.go new file mode 100644 index 000000000..1d88ab8b1 --- /dev/null +++ b/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) +} diff --git a/cli/cmd/plugins.go b/cli/cmd/plugins.go index f67466da5..67585a1bf 100644 --- a/cli/cmd/plugins.go +++ b/cli/cmd/plugins.go @@ -2,7 +2,6 @@ package cmd import ( "fmt" - "log" "os" "os/exec" "path/filepath" @@ -17,6 +16,25 @@ 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) @@ -24,6 +42,14 @@ func init() { 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 {