feat(comp): Dynamic completion for plugins

For each plugin, helm registers a ValidArgsFunction which the completion
script will call when it needs dynamic completion.  This
ValidArgsFunction will setup the plugin environment and then call the
executable `plugin.complete`, if it is provided by the plugin.

The output of the call to `plugin.complete` will be used as completion
choices.  The last line of the output can optionally be used by the
plugin to specify a completion directive.

Signed-off-by: Marc Khouzam <marc.khouzam@montreal.ca>

cmd/helm/load_plugins.go

@@ -16,6 +16,7 @@ limitations under the License.
 package main
 
 import (
+	"bytes"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -23,6 +24,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"strconv"
 	"strings"
 	"syscall"
 
@@ -30,10 +32,14 @@ import (
 	"github.com/spf13/cobra"
 	"sigs.k8s.io/yaml"
 
+	"helm.sh/helm/v3/internal/completion"
 	"helm.sh/helm/v3/pkg/plugin"
 )
 
-const pluginStaticCompletionFile = "completion.yaml"
+const (
+	pluginStaticCompletionFile        = "completion.yaml"
+	pluginDynamicCompletionExecutable = "plugin.complete"
+)
 
 type pluginError struct {
 	error
@@ -81,6 +87,33 @@ func loadPlugins(baseCmd *cobra.Command, out io.Writer) {
 			md.Usage = fmt.Sprintf("the %q plugin", md.Name)
 		}
 
+		// This function is used to setup the environment for the plugin and then
+		// call the executable specified by the parameter 'main'
+		callPluginExecutable := func(cmd *cobra.Command, main string, argv []string, out io.Writer) error {
+			env := os.Environ()
+			for k, v := range settings.EnvVars() {
+				env = append(env, fmt.Sprintf("%s=%s", k, v))
+			}
+
+			prog := exec.Command(main, argv...)
+			prog.Env = env
+			prog.Stdin = os.Stdin
+			prog.Stdout = out
+			prog.Stderr = os.Stderr
+			if err := prog.Run(); err != nil {
+				if eerr, ok := err.(*exec.ExitError); ok {
+					os.Stderr.Write(eerr.Stderr)
+					status := eerr.Sys().(syscall.WaitStatus)
+					return pluginError{
+						error: errors.Errorf("plugin %q exited with error", md.Name),
+						code:  status.ExitStatus(),
+					}
+				}
+				return err
+			}
+			return nil
+		}
+
 		c := &cobra.Command{
 			Use:   md.Name,
 			Short: md.Usage,
@@ -101,33 +134,59 @@ func loadPlugins(baseCmd *cobra.Command, out io.Writer) {
 					return errors.Errorf("plugin %q exited with error", md.Name)
 				}
 
-				env := os.Environ()
-				for k, v := range settings.EnvVars() {
-					env = append(env, fmt.Sprintf("%s=%s", k, v))
-				}
-
-				prog := exec.Command(main, argv...)
-				prog.Env = env
-				prog.Stdin = os.Stdin
-				prog.Stdout = out
-				prog.Stderr = os.Stderr
-				if err := prog.Run(); err != nil {
-					if eerr, ok := err.(*exec.ExitError); ok {
-						os.Stderr.Write(eerr.Stderr)
-						status := eerr.Sys().(syscall.WaitStatus)
-						return pluginError{
-							error: errors.Errorf("plugin %q exited with error", md.Name),
-							code:  status.ExitStatus(),
-						}
-					}
-					return err
-				}
-				return nil
+				return callPluginExecutable(cmd, main, argv, out)
 			},
 			// This passes all the flags to the subcommand.
 			DisableFlagParsing: true,
 		}
 
+		// Setup dynamic completion for the plugin
+		completion.RegisterValidArgsFunc(c, func(cmd *cobra.Command, args []string, toComplete string) ([]string, completion.BashCompDirective) {
+			u, err := processParent(cmd, args)
+			if err != nil {
+				return nil, completion.BashCompDirectiveError
+			}
+
+			// We will call the dynamic completion script of the plugin
+			main := strings.Join([]string{plug.Dir, pluginDynamicCompletionExecutable}, string(filepath.Separator))
+
+			argv := []string{}
+			if !md.IgnoreFlags {
+				argv = append(argv, u...)
+				argv = append(argv, toComplete)
+			}
+			plugin.SetupPluginEnv(settings, md.Name, plug.Dir)
+
+			completion.CompDebugln(fmt.Sprintf("calling %s with args %v", main, argv))
+			buf := new(bytes.Buffer)
+			if err := callPluginExecutable(cmd, main, argv, buf); err != nil {
+				return nil, completion.BashCompDirectiveError
+			}
+
+			var completions []string
+			for _, comp := range strings.Split(buf.String(), "\n") {
+				// Remove any empty lines
+				if len(comp) > 0 {
+					completions = append(completions, comp)
+				}
+			}
+
+			// Check if the last line of output is of the form :<integer>, which
+			// indicates the BashCompletionDirective.
+			directive := completion.BashCompDirectiveDefault
+			if len(completions) > 0 {
+				lastLine := completions[len(completions)-1]
+				if len(lastLine) > 1 && lastLine[0] == ':' {
+					if strInt, err := strconv.Atoi(lastLine[1:]); err == nil {
+						directive = completion.BashCompDirective(strInt)
+						completions = completions[:len(completions)-1]
+					}
+				}
+			}
+
+			return completions, directive
+		})
+
 		// TODO: Make sure a command with this name does not already exist.
 		baseCmd.AddCommand(c)
 	}

cmd/helm/plugin_test.go

@@ -25,6 +25,8 @@ import (
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
+
+	"helm.sh/helm/v3/pkg/release"
 )
 
 func TestManuallyProcessArgs(t *testing.T) {
@@ -242,6 +244,45 @@ func checkCommand(t *testing.T, plugins []*cobra.Command, tests []staticCompleti
 	}
 }
 
+func TestPluginDynamicCompletion(t *testing.T) {
+
+	tests := []cmdTestCase{{
+		name:   "completion for plugin",
+		cmd:    "__complete args ''",
+		golden: "output/plugin_args_comp.txt",
+		rels:   []*release.Release{},
+	}, {
+		name:   "completion for plugin with flag",
+		cmd:    "__complete args --myflag ''",
+		golden: "output/plugin_args_flag_comp.txt",
+		rels:   []*release.Release{},
+	}, {
+		name:   "completion for plugin with global flag",
+		cmd:    "__complete args --namespace mynamespace ''",
+		golden: "output/plugin_args_ns_comp.txt",
+		rels:   []*release.Release{},
+	}, {
+		name:   "completion for plugin with multiple args",
+		cmd:    "__complete args --myflag --namespace mynamespace start",
+		golden: "output/plugin_args_many_args_comp.txt",
+		rels:   []*release.Release{},
+	}, {
+		name:   "completion for plugin no directive",
+		cmd:    "__complete echo -n mynamespace ''",
+		golden: "output/plugin_echo_no_directive.txt",
+		rels:   []*release.Release{},
+	}, {
+		name:   "completion for plugin bad directive",
+		cmd:    "__complete echo ''",
+		golden: "output/plugin_echo_bad_directive.txt",
+		rels:   []*release.Release{},
+	}}
+	for _, test := range tests {
+		settings.PluginsDirectory = "testdata/helmhome/helm/plugins"
+		runTestCmd(t, []cmdTestCase{test})
+	}
+}
+
 func TestLoadPlugins_HelmNoPlugins(t *testing.T) {
 	settings.PluginsDirectory = "testdata/helmhome/helm/plugins"
 	settings.RepositoryConfig = "testdata/helmhome/helm/repository"

cmd/helm/testdata/helmhome/helm/plugins/args/plugin.complete

@@ -0,0 +1,13 @@
+#!/usr/bin/env sh
+
+echo "plugin.complete was called"
+echo "Namespace: ${HELM_NAMESPACE:-NO_NS}"
+echo "Num args received: ${#}"
+echo "Args received: ${@}"
+
+# Final printout is the optional completion directive of the form :<directive>
+if [ "$HELM_NAMESPACE" = "default" ]; then
+    echo ":4"
+else
+    echo ":2"
+fi

cmd/helm/testdata/helmhome/helm/plugins/echo/plugin.complete

@@ -0,0 +1,14 @@
+#!/usr/bin/env sh
+
+echo "echo plugin.complete was called"
+echo "Namespace: ${HELM_NAMESPACE:-NO_NS}"
+echo "Num args received: ${#}"
+echo "Args received: ${@}"
+
+# Final printout is the optional completion directive of the form :<directive>
+if [ "$HELM_NAMESPACE" = "default" ]; then
+    # Output an invalid directive, which should be ignored
+    echo ":2222"
+# else
+    # Don't include the directive, to test it is really optional
+fi

cmd/helm/testdata/output/plugin_args_comp.txt

@@ -0,0 +1,5 @@
+plugin.complete was called
+Namespace: default
+Num args received: 1
+Args received: 
+:4

cmd/helm/testdata/output/plugin_args_flag_comp.txt

@@ -0,0 +1,5 @@
+plugin.complete was called
+Namespace: default
+Num args received: 2
+Args received: --myflag 
+:4

cmd/helm/testdata/output/plugin_args_many_args_comp.txt

@@ -0,0 +1,5 @@
+plugin.complete was called
+Namespace: mynamespace
+Num args received: 2
+Args received: --myflag start
+:2

cmd/helm/testdata/output/plugin_args_ns_comp.txt

@@ -0,0 +1,5 @@
+plugin.complete was called
+Namespace: mynamespace
+Num args received: 1
+Args received: 
+:2

cmd/helm/testdata/output/plugin_echo_bad_directive.txt

@@ -0,0 +1,5 @@
+echo plugin.complete was called
+Namespace: default
+Num args received: 1
+Args received: 
+:0

cmd/helm/testdata/output/plugin_echo_no_directive.txt

@@ -0,0 +1,5 @@
+echo plugin.complete was called
+Namespace: mynamespace
+Num args received: 1
+Args received: 
+:0

internal/completion/complete.go

@@ -188,29 +188,47 @@ func NewCompleteCmd(settings *cli.EnvSettings, out io.Writer) *cobra.Command {
 		Run: func(cmd *cobra.Command, args []string) {
 			CompDebugln(fmt.Sprintf("%s was called with args %v", cmd.Name(), args))
 
-			flag, trimmedArgs, toComplete, err := checkIfFlagCompletion(cmd.Root(), args[:len(args)-1], args[len(args)-1])
-			if err != nil {
-				// Error while attempting to parse flags
-				CompErrorln(err.Error())
-				return
-			}
+			// The last argument, which is not complete, should not be part of the list of arguments
+			toComplete := args[len(args)-1]
+			trimmedArgs := args[:len(args)-1]
+
 			// Find the real command for which completion must be performed
 			finalCmd, finalArgs, err := cmd.Root().Find(trimmedArgs)
 			if err != nil {
 				// Unable to find the real command. E.g., helm invalidCmd <TAB>
+				CompDebugln(fmt.Sprintf("Unable to find a command for arguments: %v", trimmedArgs))
 				return
 			}
 
 			CompDebugln(fmt.Sprintf("Found final command '%s', with finalArgs %v", finalCmd.Name(), finalArgs))
 
+			var flag *pflag.Flag
+			if !finalCmd.DisableFlagParsing {
+				// We only do flag completion if we are allowed to parse flags
+				// This is important for helm plugins which need to do their own flag completion.
+				flag, finalArgs, toComplete, err = checkIfFlagCompletion(finalCmd, finalArgs, toComplete)
+				if err != nil {
+					// Error while attempting to parse flags
+					CompErrorln(err.Error())
+					return
+				}
+			}
+
 			// Parse the flags and extract the arguments to prepare for calling the completion function
 			if err = finalCmd.ParseFlags(finalArgs); err != nil {
 				CompErrorln(fmt.Sprintf("Error while parsing flags from args %v: %s", finalArgs, err.Error()))
 				return
 			}
-			argsWoFlags := finalCmd.Flags().Args()
-			CompDebugln(fmt.Sprintf("Args without flags are '%v' with length %d", argsWoFlags, len(argsWoFlags)))
 
+			// We only remove the flags from the arguments if DisableFlagParsing is not set.
+			// This is important for helm plugins, which need to receive all flags.
+			// The plugin completion code will do its own flag parsing.
+			if !finalCmd.DisableFlagParsing {
+				finalArgs = finalCmd.Flags().Args()
+				CompDebugln(fmt.Sprintf("Args without flags are '%v' with length %d", finalArgs, len(finalArgs)))
+			}
+
+			// Find completion function for the flag or command
 			var key interface{}
 			var keyStr string
 			if flag != nil {
@@ -220,21 +238,23 @@ func NewCompleteCmd(settings *cli.EnvSettings, out io.Writer) *cobra.Command {
 				key = finalCmd
 				keyStr = finalCmd.Name()
 			}
-
-			// Find completion function for the flag or command
 			completionFn, ok := validArgsFunctions[key]
 			if !ok {
 				CompErrorln(fmt.Sprintf("Dynamic completion not supported/needed for flag or command: %s", keyStr))
 				return
 			}
 
-			CompDebugln(fmt.Sprintf("Calling completion method for subcommand '%s' with args '%v' and toComplete '%s'", finalCmd.Name(), argsWoFlags, toComplete))
-			completions, directive := completionFn(finalCmd, argsWoFlags, toComplete)
+			CompDebugln(fmt.Sprintf("Calling completion method for subcommand '%s' with args '%v' and toComplete '%s'", finalCmd.Name(), finalArgs, toComplete))
+			completions, directive := completionFn(finalCmd, finalArgs, toComplete)
 			for _, comp := range completions {
 				// Print each possible completion to stdout for the completion script to consume.
 				fmt.Fprintln(out, comp)
 			}
 
+			if directive > BashCompDirectiveError+BashCompDirectiveNoSpace+BashCompDirectiveNoFileComp {
+				directive = BashCompDirectiveDefault
+			}
+
 			// As the last printout, print the completion directive for the
 			// completion script to parse.
 			// The directive integer must be that last character following a single :
@@ -252,7 +272,7 @@ func isFlag(arg string) bool {
 	return len(arg) > 0 && arg[0] == '-'
 }
 
-func checkIfFlagCompletion(rootCmd *cobra.Command, args []string, lastArg string) (*pflag.Flag, []string, string, error) {
+func checkIfFlagCompletion(finalCmd *cobra.Command, args []string, lastArg string) (*pflag.Flag, []string, string, error) {
 	var flagName string
 	trimmedArgs := args
 	flagWithEqual := false
@@ -287,19 +307,10 @@ func checkIfFlagCompletion(rootCmd *cobra.Command, args []string, lastArg string
 		return nil, trimmedArgs, lastArg, nil
 	}
 
-	// Find the real command for which completion must be performed
-	finalCmd, _, err := rootCmd.Find(trimmedArgs)
-	if err != nil {
-		// Unable to find the real command. E.g., helm invalidCmd <TAB>
-		return nil, nil, "", errors.New("Unable to find final command for completion")
-	}
-
-	CompDebugln(fmt.Sprintf("checkIfFlagCompletion: found final command '%s'", finalCmd.Name()))
-
 	flag := findFlag(finalCmd, flagName)
 	if flag == nil {
 		// Flag not supported by this command, nothing to complete
-		err = fmt.Errorf("Subcommand '%s' does not support flag '%s'", finalCmd.Name(), flagName)
+		err := fmt.Errorf("Subcommand '%s' does not support flag '%s'", finalCmd.Name(), flagName)
 		return nil, nil, "", err
 	}