Infrakit CLI improvements (#450)

Signed-off-by: David Chung <[email protected]>

Author: David Chung

pkg/cli/examples/CreateCluster renamed to examples/cli/CreateCluster

@@ -9,6 +9,7 @@ import (
 	"strings"
 
 	"github.com/docker/infrakit/pkg/template"
+	"github.com/docker/infrakit/pkg/util/exec"
 	"github.com/spf13/cobra"
 )
 
@@ -160,17 +161,31 @@ func (c *Context) loadBackend() error {
 	}
 	t.AddFunc("print",
 		func() string {
-			c.run = func(view string) error {
-				fmt.Println(view)
+			c.run = func(script string) error {
+				fmt.Println(script)
 				return nil
 			}
 			return ""
 		})
 	t.AddFunc("sh",
-		func() string {
-			c.run = func(view string) error {
-				fmt.Println(view)
-				return nil
+		func(opts ...string) string {
+			c.run = func(script string) error {
+
+				cmd := strings.Join(append([]string{"/bin/sh"}, opts...), " ")
+				log.Debug("sh", "cmd", cmd)
+
+				return exec.Command(cmd).
+					InheritEnvs(true).StartWithStreams(
+
+					exec.Do(exec.SendInput(
+						func(stdin io.WriteCloser) error {
+							_, err := stdin.Write([]byte(script))
+							return err
+						})).Then(
+						exec.RedirectStdout(os.Stdout)).Then(
+						exec.RedirectStderr(os.Stderr),
+					).Done(),
+				)
 			}
 			return ""
 		})

pkg/cli/examples/README.md renamed to examples/cli/README.md

@@ -23,7 +23,7 @@ func TestContext(t *testing.T) {
 	// the cobra command
 
 	script := `
-{{/* The directive here tells infrakit to run this script with sh:  =% sh %=  */}}
+{{/* The directive here tells infrakit to run this script with sh:  =% print %=  */}}
 
 {{/* The function 'flag' will create a flag in the CLI; the function 'prompt' will ask user for input */}}
 
@@ -89,3 +89,38 @@ func TestContext(t *testing.T) {
 		"instanceType": "large",
 	}).String(), types.AnyValueMust(m).String())
 }
+
+func TestContextRunShell(t *testing.T) {
+
+	script := `#!/bin/bash
+{{/* The directive here tells infrakit to run this script with sh:  =% sh "-s" "--"  %=  */}}
+{{ $lines := flag "lines" "int" "the number of lines" 5 }}
+
+for i in $(seq {{$lines}}); do
+echo line $i
+done
+`
+
+	c := &Context{
+		cmd: &cobra.Command{
+			Use:   "test",
+			Short: "test",
+		},
+		src: "str://" + script,
+	}
+
+	c.exec = false
+	err := c.buildFlags()
+	require.NoError(t, err)
+
+	err = c.cmd.Flags().Parse(strings.Split("--lines 3", " "))
+	require.NoError(t, err)
+
+	err = c.loadBackend()
+	require.NoError(t, err)
+	require.NotNil(t, c.run)
+
+	err = c.execute()
+	require.NoError(t, err)
+
+}

pkg/cli/examples/build/.short renamed to examples/cli/build/.short

@@ -1,13 +1,17 @@
 package exec
 
 import (
+	"io"
+	"os"
 	"os/exec"
 	"strings"
 
+	logutil "github.com/docker/infrakit/pkg/log"
 	"github.com/docker/infrakit/pkg/template"
-	log "github.com/golang/glog"
 )
 
+var log = logutil.New("module", "util/exec")
+
 // Command is the template which is rendered before it's executed
 type Command string
 
@@ -26,35 +30,50 @@ func (c Command) Run(args ...string) error {
 	return c.builder().Run(args...)
 }
 
+// String returns the interpolated version of the command
+func (c Command) String(args ...string) (string, error) {
+	p, err := c.builder().generate(args...)
+	if err == nil {
+		return strings.Join(p, " "), nil
+	}
+	return string(c), err
+}
+
 // WithOptions adds the template options
 func (c Command) WithOptions(options template.Options) *Builder {
-	b := c.builder()
-	b.options = options
-	return b
+	return c.builder().WithOptions(options)
 }
 
 // WithFunc adds a function that can be used in the template
 func (c Command) WithFunc(name string, function interface{}) *Builder {
-	b := c.builder()
-	b.funcs[name] = function
-	return b
+	return c.builder().WithFunc(name, function)
 }
 
 // WithContext sets the context for the template
 func (c Command) WithContext(context interface{}) *Builder {
-	b := c.builder()
-	b.context = context
-	return b
+	return c.builder().WithContext(context)
+}
+
+// InheritEnvs determines whether the process should inherit the envs of the parent
+func (c Command) InheritEnvs(v bool) *Builder {
+	return c.builder().InheritEnvs(v)
+}
+
+// NewCommand creates an instance of the command builder to allow detailed configuration
+func NewCommand(s string) *Builder {
+	return Command(s).builder()
 }
 
 // Builder collects options until it's run
 type Builder struct {
-	command  Command
-	options  template.Options
-	funcs    map[string]interface{}
-	context  interface{}
-	rendered string // rendered command string
-	cmd      *exec.Cmd
+	command     Command
+	options     template.Options
+	inheritEnvs bool
+	envs        []string
+	funcs       map[string]interface{}
+	context     interface{}
+	rendered    string // rendered command string
+	cmd         *exec.Cmd
 }
 
 func (c Command) builder() *Builder {
@@ -64,6 +83,18 @@ func (c Command) builder() *Builder {
 	}
 }
 
+// InheritEnvs determines whether the process should inherit the envs of the parent
+func (b *Builder) InheritEnvs(v bool) *Builder {
+	b.inheritEnvs = v
+	return b
+}
+
+// WithEnvs adds environment variables for the exec, in format of key=value
+func (b *Builder) WithEnvs(kv ...string) *Builder {
+	b.envs = append(b.envs, kv...)
+	return b
+}
+
 // WithOptions adds the template options
 func (b *Builder) WithOptions(options template.Options) *Builder {
 	b.options = options
@@ -100,6 +131,104 @@ func (b *Builder) Start(args ...string) error {
 	return run.Start()
 }
 
+// Step is something you do with the processes streams
+type Step func(stdin io.WriteCloser, stdout io.ReadCloser, stderr io.ReadCloser) error
+
+// Thenable is a fluent builder for chaining tasks
+type Thenable struct {
+	steps []Step
+}
+
+// Do creates a thenable
+func Do(f Step) *Thenable {
+	return &Thenable{
+		steps: []Step{f},
+	}
+}
+
+// Then adds another step
+func (t *Thenable) Then(then Step) *Thenable {
+	t.steps = append(t.steps, then)
+	return t
+}
+
+// Done returns the final function
+func (t *Thenable) Done() Step {
+	steps := t.steps
+	return func(stdin io.WriteCloser, stdout, stderr io.ReadCloser) error {
+		for _, step := range steps {
+			if err := step(stdin, stdout, stderr); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+}
+
+// SendInput is a convenience function for writing to the exec process's stdin. When the function completes, the
+// stdin is closed.
+func SendInput(f func(io.WriteCloser) error) Step {
+	return func(stdin io.WriteCloser, stdout, stderr io.ReadCloser) error {
+		defer stdin.Close()
+		return f(stdin)
+	}
+}
+
+// RedirectStdout sends stdout to given writer
+func RedirectStdout(out io.Writer) Step {
+	return func(stdin io.WriteCloser, stdout, stderr io.ReadCloser) error {
+		_, err := io.Copy(out, stdout)
+		return err
+	}
+}
+
+// RedirectStderr sends stdout to given writer
+func RedirectStderr(out io.Writer) Step {
+	return func(stdin io.WriteCloser, stdout, stderr io.ReadCloser) error {
+		_, err := io.Copy(out, stderr)
+		return err
+	}
+}
+
+// MergeOutput combines the stdout and stderr into the given stream
+func MergeOutput(out io.Writer) Step {
+	return func(stdin io.WriteCloser, stdout, stderr io.ReadCloser) error {
+		_, err := io.Copy(out, io.MultiReader(stdout, stderr))
+		return err
+	}
+}
+
+// StartWithStreams starts the the process and then calls the function which allows
+// the streams to be wired.  Calling the provided function blocks.
+func (b *Builder) StartWithStreams(f Step,
+	args ...string) error {
+
+	_, err := b.exec(args...)
+	if err != nil {
+		return err
+	}
+
+	pOut, err := b.cmd.StdoutPipe()
+	if err != nil {
+		return err
+	}
+	pErr, err := b.cmd.StderrPipe()
+	if err != nil {
+		return err
+	}
+	pIn, err := b.cmd.StdinPipe()
+	if err != nil {
+		return err
+	}
+
+	err = b.cmd.Start()
+	if err != nil {
+		return err
+	}
+
+	return f(pIn, pOut, pErr)
+}
+
 // Run does a Cmd.Run on the command
 func (b *Builder) Run(args ...string) error {
 	run, err := b.exec(args...)
@@ -143,7 +272,11 @@ func (b *Builder) exec(args ...string) (*exec.Cmd, error) {
 	if err != nil {
 		return nil, err
 	}
-	log.V(50).Infoln("exec:", command)
+	log.Debug("exec", "command", command)
 	b.cmd = exec.Command(command[0], command[1:]...)
+	if b.inheritEnvs {
+		b.cmd.Env = append(os.Environ(), b.envs...)
+	}
+
 	return b.cmd, nil
 }

pkg/cli/examples/build/README.md renamed to examples/cli/build/README.md

@@ -1,6 +1,8 @@
 package exec
 
 import (
+	"io"
+	"os"
 	"strings"
 	"testing"
 	"time"
@@ -21,10 +23,11 @@ docker run --rm \
 	busyboxLs Command = `docker run --rm \
        busybox ls {{ arg 1 }}
 `
-	busyboxDateStream Command = `docker run --rm --name {{ .container }} \
-busybox /bin/sh -c 'while true; do date; sleep {{ .sleep }}; done'`
+	busyboxSh Command = `docker run --rm -ti --name {{ .container }} busybox /bin/sh`
 
 	dockerStop Command = `docker stop {{ arg 1 }}`
+
+	dateStream Command = `docker run --rm --name {{ arg 1 }} chungers/timer streamer`
 )
 
 func TestBuilder(t *testing.T) {
@@ -44,11 +47,10 @@ func TestBuilder(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, []string{"docker", "run", "--rm", "busybox", "ls", "sys"}, cmd)
 
-	b = busyboxDateStream.builder().WithContext(map[string]interface{}{"container": "bob", "sleep": "1"})
+	b = busyboxSh.builder().WithContext(map[string]interface{}{"container": "bob"})
 	cmd, err = b.generate()
 	require.NoError(t, err)
-	require.Equal(t, []string{"docker", "run", "--rm", "--name", "bob", "busybox", "/bin/sh", "-c",
-		"'while", "true;", "do", "date;", "sleep", "1;", "done'"}, cmd)
+	require.Equal(t, []string{"docker", "run", "--rm", "-ti", "--name", "bob", "busybox", "/bin/sh"}, cmd)
 
 }
 
@@ -71,11 +73,28 @@ func TestRun(t *testing.T) {
 	require.Equal(t, []string{"spool", "www"}, strings.Split(strings.Trim(string(output), " \n"), "\n"))
 
 	name := "stream-test"
-	err = busyboxDateStream.WithContext(map[string]interface{}{"container": name, "sleep": 1}).Start()
+	go func() {
+		<-time.After(2 * time.Second)
+		err := dockerStop.Run(name)
+		if err != nil {
+			panic(err)
+		}
+	}()
+
+	err = dateStream.InheritEnvs(true).StartWithStreams(MergeOutput(os.Stderr),
+		name, // arg 1 for container name
+	)
 	require.NoError(t, err)
 
-	time.Sleep(1 * time.Second)
-
-	err = dockerStop.Run(name)
+	// testing with stdin
+	err = Command("/bin/sh").InheritEnvs(true).StartWithStreams(
+		Do(SendInput(
+			func(stdin io.WriteCloser) error {
+				stdin.Write([]byte(`for i in $(seq 10); do echo $i; sleep 1; done`))
+				return nil
+			})).Then(MergeOutput(os.Stderr)).Done(),
+		name, // arg 1 for container name
+	)
 	require.NoError(t, err)
+
 }