ctlv3: support ETCDCTL_WATCH_KEY, ETCDCTL_WATCH_RANGE_END

Signed-off-by: Gyuho Lee <gyuhox@gmail.com>

etcdctl/ctlv3/command/watch_command.go

@@ -30,6 +30,7 @@ import (
 
 var (
 	errBadArgsNum              = errors.New("bad number of arguments")
+	errBadArgsNumConflictEnv   = errors.New("bad number of arguments (found conflicting environment key)")
 	errBadArgsNumSeparator     = errors.New("bad number of arguments (found separator --, but no commands)")
 	errBadArgsInteractiveWatch = errors.New("args[0] must be 'watch' for interactive calls")
 )
@@ -59,12 +60,17 @@ func NewWatchCommand() *cobra.Command {
 
 // watchCommandFunc executes the "watch" command.
 func watchCommandFunc(cmd *cobra.Command, args []string) {
+	envKey, envRange := os.Getenv("ETCDCTL_WATCH_KEY"), os.Getenv("ETCDCTL_WATCH_RANGE_END")
+	if envKey == "" && envRange != "" {
+		ExitWithError(ExitBadArgs, fmt.Errorf("ETCDCTL_WATCH_KEY is empty but got ETCDCTL_WATCH_RANGE_END=%q", envRange))
+	}
+
 	if watchInteractive {
-		watchInteractiveFunc(cmd, os.Args)
+		watchInteractiveFunc(cmd, os.Args, envKey, envRange)
 		return
 	}
 
-	watchArgs, execArgs, err := parseWatchArgs(os.Args, args, false)
+	watchArgs, execArgs, err := parseWatchArgs(os.Args, args, envKey, envRange, false)
 	if err != nil {
 		ExitWithError(ExitBadArgs, err)
 	}
@@ -82,7 +88,7 @@ func watchCommandFunc(cmd *cobra.Command, args []string) {
 	ExitWithError(ExitInterrupted, fmt.Errorf("watch is canceled by the server"))
 }
 
-func watchInteractiveFunc(cmd *cobra.Command, osArgs []string) {
+func watchInteractiveFunc(cmd *cobra.Command, osArgs []string, envKey, envRange string) {
 	c := mustClientFromCmd(cmd)
 
 	reader := bufio.NewReader(os.Stdin)
@@ -95,7 +101,7 @@ func watchInteractiveFunc(cmd *cobra.Command, osArgs []string) {
 		l = strings.TrimSuffix(l, "\n")
 
 		args := argify(l)
-		if len(args) < 2 {
+		if len(args) < 2 && envKey == "" {
 			fmt.Fprintf(os.Stderr, "Invalid command %s (command type or key is not provided)\n", l)
 			continue
 		}
@@ -105,7 +111,7 @@ func watchInteractiveFunc(cmd *cobra.Command, osArgs []string) {
 			continue
 		}
 
-		watchArgs, execArgs, perr := parseWatchArgs(osArgs, args, true)
+		watchArgs, execArgs, perr := parseWatchArgs(osArgs, args, envKey, envRange, true)
 		if perr != nil {
 			ExitWithError(ExitBadArgs, perr)
 		}
@@ -165,7 +171,7 @@ func printWatchCh(c *clientv3.Client, ch clientv3.WatchChan, execArgs []string)
 // (e.g. ./bin/etcdctl watch foo --rev 1 bar).
 // "--" characters are invalid arguments for "spf13/cobra" library,
 // so no need to handle such cases.
-func parseWatchArgs(osArgs, commandArgs []string, interactive bool) (watchArgs []string, execArgs []string, err error) {
+func parseWatchArgs(osArgs, commandArgs []string, envKey, envRange string, interactive bool) (watchArgs []string, execArgs []string, err error) {
 	watchArgs = commandArgs
 
 	// remove preceding commands (e.g. "watch foo bar" in interactive mode)
@@ -175,12 +181,54 @@ func parseWatchArgs(osArgs, commandArgs []string, interactive bool) (watchArgs [
 			break
 		}
 	}
+	if idx < len(watchArgs)-1 || envKey != "" {
 		if idx < len(watchArgs)-1 {
 			watchArgs = watchArgs[idx+1:]
+		}
+
+		execIdx, execExist := 0, false
+		for execIdx = range osArgs {
+			v := osArgs[execIdx]
+			if v == "--" && execIdx != len(osArgs)-1 {
+				execExist = true
+				break
+			}
+		}
+
+		if idx == len(watchArgs)-1 && envKey != "" {
+			if len(watchArgs) > 0 && !interactive {
+				// "watch --rev 1 -- echo Hello World" has no conflict
+				if !execExist {
+					// "watch foo" with ETCDCTL_WATCH_KEY=foo
+					// (watchArgs==["foo"])
+					return nil, nil, errBadArgsNumConflictEnv
+				}
+			}
+			// otherwise, watch with no argument and environment key is set
+			// if interactive, first "watch" command string should be removed
+			if interactive {
+				watchArgs = []string{}
+			}
+		}
+
+		// "watch foo -- echo hello" with ETCDCTL_WATCH_KEY=foo
+		// (watchArgs==["foo","echo","hello"])
+		if envKey != "" && execExist {
+			widx, oidx := 0, len(osArgs)-1
+			for widx = len(watchArgs) - 1; widx >= 0; widx-- {
+				if watchArgs[widx] == osArgs[oidx] {
+					oidx--
+					continue
+				}
+				if oidx == execIdx { // watchArgs has extra
+					return nil, nil, errBadArgsNumConflictEnv
+				}
+			}
+		}
 	} else if interactive { // "watch" not found
 		return nil, nil, errBadArgsInteractiveWatch
 	}
-	if len(watchArgs) < 1 {
+	if len(watchArgs) < 1 && envKey == "" {
 		return nil, nil, errBadArgsNum
 	}
 
@@ -192,7 +240,7 @@ func parseWatchArgs(osArgs, commandArgs []string, interactive bool) (watchArgs [
 	}
 	if idx < len(osArgs)-1 {
 		osArgs = osArgs[idx+1:]
-	} else {
+	} else if envKey == "" {
 		return nil, nil, errBadArgsNum
 	}
 
@@ -202,7 +250,7 @@ func parseWatchArgs(osArgs, commandArgs []string, interactive bool) (watchArgs [
 	}
 	foundSep := false
 	for idx = range argsWithSep {
-		if argsWithSep[idx] == "--" && idx > 0 {
+		if argsWithSep[idx] == "--" {
 			foundSep = true
 			break
 		}
@@ -214,6 +262,18 @@ func parseWatchArgs(osArgs, commandArgs []string, interactive bool) (watchArgs [
 		}
 		watchArgs = flagset.Args()
 	}
+
+	// "watch -- echo hello" with ETCDCTL_WATCH_KEY=foo
+	// should be translated to "watch foo -- echo hello"
+	// (watchArgs=["echo","hello"] should be ["foo","echo","hello"])
+	if envKey != "" {
+		tmp := []string{envKey}
+		if envRange != "" {
+			tmp = append(tmp, envRange)
+		}
+		watchArgs = append(tmp, watchArgs...)
+	}
+
 	if !foundSep {
 		return watchArgs, nil, nil
 	}

etcdctl/ctlv3/command/watch_command_test.go

@@ -23,6 +23,7 @@ func Test_parseWatchArgs(t *testing.T) {
 	tt := []struct {
 		osArgs           []string // raw arguments to "watch" command
 		commandArgs      []string // arguments after "spf13/cobra" preprocessing
+		envKey, envRange string
 		interactive      bool
 
 		watchArgs []string
@@ -45,9 +46,66 @@ func Test_parseWatchArgs(t *testing.T) {
 			execArgs:    nil,
 			err:         errBadArgsNumSeparator,
 		},
+		{
+			osArgs:      []string{"./bin/etcdctl", "watch"},
+			commandArgs: nil,
+			envKey:      "foo",
+			envRange:    "bar",
+			interactive: false,
+			watchArgs:   []string{"foo", "bar"},
+			execArgs:    nil,
+			err:         nil,
+		},
 		{
 			osArgs:      []string{"./bin/etcdctl", "watch", "foo"},
 			commandArgs: []string{"foo"},
+			envKey:      "foo",
+			envRange:    "",
+			interactive: false,
+			watchArgs:   nil,
+			execArgs:    nil,
+			err:         errBadArgsNumConflictEnv,
+		},
+		{
+			osArgs:      []string{"./bin/etcdctl", "watch", "foo", "bar"},
+			commandArgs: []string{"foo", "bar"},
+			envKey:      "foo",
+			envRange:    "",
+			interactive: false,
+			watchArgs:   nil,
+			execArgs:    nil,
+			err:         errBadArgsNumConflictEnv,
+		},
+		{
+			osArgs:      []string{"./bin/etcdctl", "watch", "foo", "bar"},
+			commandArgs: []string{"foo", "bar"},
+			envKey:      "foo",
+			envRange:    "bar",
+			interactive: false,
+			watchArgs:   nil,
+			execArgs:    nil,
+			err:         errBadArgsNumConflictEnv,
+		},
+		{
+			osArgs:      []string{"./bin/etcdctl", "watch", "foo"},
+			commandArgs: []string{"foo"},
+			interactive: false,
+			watchArgs:   []string{"foo"},
+			execArgs:    nil,
+			err:         nil,
+		},
+		{
+			osArgs:      []string{"./bin/etcdctl", "watch"},
+			commandArgs: nil,
+			envKey:      "foo",
+			interactive: false,
+			watchArgs:   []string{"foo"},
+			execArgs:    nil,
+			err:         nil,
+		},
+		{
+			osArgs:      []string{"./bin/etcdctl", "watch", "--rev", "1", "foo"},
+			commandArgs: []string{"foo"},
 			interactive: false,
 			watchArgs:   []string{"foo"},
 			execArgs:    nil,
@@ -56,6 +114,16 @@ func Test_parseWatchArgs(t *testing.T) {
 		{
 			osArgs:      []string{"./bin/etcdctl", "watch", "--rev", "1", "foo"},
 			commandArgs: []string{"foo"},
+			envKey:      "foo",
+			interactive: false,
+			watchArgs:   nil,
+			execArgs:    nil,
+			err:         errBadArgsNumConflictEnv,
+		},
+		{
+			osArgs:      []string{"./bin/etcdctl", "watch", "--rev", "1"},
+			commandArgs: nil,
+			envKey:      "foo",
 			interactive: false,
 			watchArgs:   []string{"foo"},
 			execArgs:    nil,
@@ -117,6 +185,35 @@ func Test_parseWatchArgs(t *testing.T) {
 			execArgs:    []string{"echo", "Hello", "World"},
 			err:         nil,
 		},
+		{
+			osArgs:      []string{"./bin/etcdctl", "watch", "--rev", "1", "--", "echo", "Hello", "World"},
+			commandArgs: []string{"echo", "Hello", "World"},
+			envKey:      "foo",
+			envRange:    "",
+			interactive: false,
+			watchArgs:   []string{"foo"},
+			execArgs:    []string{"echo", "Hello", "World"},
+			err:         nil,
+		},
+		{
+			osArgs:      []string{"./bin/etcdctl", "watch", "--rev", "1", "--", "echo", "Hello", "World"},
+			commandArgs: []string{"echo", "Hello", "World"},
+			envKey:      "foo",
+			envRange:    "bar",
+			interactive: false,
+			watchArgs:   []string{"foo", "bar"},
+			execArgs:    []string{"echo", "Hello", "World"},
+			err:         nil,
+		},
+		{
+			osArgs:      []string{"./bin/etcdctl", "watch", "foo", "bar", "--rev", "1", "--", "echo", "Hello", "World"},
+			commandArgs: []string{"foo", "bar", "echo", "Hello", "World"},
+			envKey:      "foo",
+			interactive: false,
+			watchArgs:   nil,
+			execArgs:    nil,
+			err:         errBadArgsNumConflictEnv,
+		},
 		{
 			osArgs:      []string{"./bin/etcdctl", "watch", "-i"},
 			commandArgs: []string{"foo", "bar", "--", "echo", "Hello", "World"},
@@ -141,6 +238,26 @@ func Test_parseWatchArgs(t *testing.T) {
 			execArgs:    nil,
 			err:         nil,
 		},
+		{
+			osArgs:      []string{"./bin/etcdctl", "watch", "-i"},
+			commandArgs: []string{"watch"},
+			envKey:      "foo",
+			envRange:    "bar",
+			interactive: true,
+			watchArgs:   []string{"foo", "bar"},
+			execArgs:    nil,
+			err:         nil,
+		},
+		{
+			osArgs:      []string{"./bin/etcdctl", "watch", "-i"},
+			commandArgs: []string{"watch"},
+			envKey:      "hello world!",
+			envRange:    "bar",
+			interactive: true,
+			watchArgs:   []string{"hello world!", "bar"},
+			execArgs:    nil,
+			err:         nil,
+		},
 		{
 			osArgs:      []string{"./bin/etcdctl", "watch", "-i"},
 			commandArgs: []string{"watch", "foo", "--rev", "1"},
@@ -165,6 +282,25 @@ func Test_parseWatchArgs(t *testing.T) {
 			execArgs:    []string{"echo", "Hello", "World"},
 			err:         nil,
 		},
+		{
+			osArgs:      []string{"./bin/etcdctl", "watch", "-i"},
+			commandArgs: []string{"watch", "--", "echo", "Hello", "World"},
+			envKey:      "foo",
+			interactive: true,
+			watchArgs:   []string{"foo"},
+			execArgs:    []string{"echo", "Hello", "World"},
+			err:         nil,
+		},
+		{
+			osArgs:      []string{"./bin/etcdctl", "watch", "-i"},
+			commandArgs: []string{"watch", "--", "echo", "Hello", "World"},
+			envKey:      "foo",
+			envRange:    "bar",
+			interactive: true,
+			watchArgs:   []string{"foo", "bar"},
+			execArgs:    []string{"echo", "Hello", "World"},
+			err:         nil,
+		},
 		{
 			osArgs:      []string{"./bin/etcdctl", "watch", "-i"},
 			commandArgs: []string{"watch", "foo", "bar", "--", "echo", "Hello", "World"},
@@ -181,6 +317,16 @@ func Test_parseWatchArgs(t *testing.T) {
 			execArgs:    []string{"echo", "Hello", "World"},
 			err:         nil,
 		},
+		{
+			osArgs:      []string{"./bin/etcdctl", "watch", "-i"},
+			commandArgs: []string{"watch", "--rev", "1", "--", "echo", "Hello", "World"},
+			envKey:      "foo",
+			envRange:    "bar",
+			interactive: true,
+			watchArgs:   []string{"foo", "bar"},
+			execArgs:    []string{"echo", "Hello", "World"},
+			err:         nil,
+		},
 		{
 			osArgs:      []string{"./bin/etcdctl", "watch", "-i"},
 			commandArgs: []string{"watch", "foo", "--rev", "1", "bar", "--", "echo", "Hello", "World"},
@@ -199,7 +345,7 @@ func Test_parseWatchArgs(t *testing.T) {
 		},
 	}
 	for i, ts := range tt {
-		watchArgs, execArgs, err := parseWatchArgs(ts.osArgs, ts.commandArgs, ts.interactive)
+		watchArgs, execArgs, err := parseWatchArgs(ts.osArgs, ts.commandArgs, ts.envKey, ts.envRange, ts.interactive)
 		if err != ts.err {
 			t.Fatalf("#%d: error expected %v, got %v", i, ts.err, err)
 		}