Added process.Background() and process.Forwarded()

This PR adds higher-level wrappers for calling subprocesses

Author: nfx

bundle/config/artifact.go

@@ -4,11 +4,11 @@ import (
 	"bytes"
 	"context"
 	"fmt"
-	"os/exec"
 	"path"
 	"strings"
 
 	"github.com/databricks/cli/bundle/config/paths"
+	"github.com/databricks/cli/libs/process"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 )
 
@@ -56,13 +56,11 @@ func (a *Artifact) Build(ctx context.Context) ([]byte, error) {
 	commands := strings.Split(a.BuildCommand, " && ")
 	for _, command := range commands {
 		buildParts := strings.Split(command, " ")
-		cmd := exec.CommandContext(ctx, buildParts[0], buildParts[1:]...)
-		cmd.Dir = a.Path
-		res, err := cmd.CombinedOutput()
+		res, err := process.Background(ctx, buildParts, process.WithDir(a.Path))
 		if err != nil {
-			return res, err
+			return nil, err
 		}
-		out = append(out, res)
+		out = append(out, []byte(res))
 	}
 	return bytes.Join(out, []byte{}), nil
 }

bundle/scripts/scripts.go

@@ -61,6 +61,7 @@ func executeHook(ctx context.Context, b *bundle.Bundle, hook config.ScriptHook)
 		return nil, nil, err
 	}
 
+	// TODO: switch to process.Background(...)
 	cmd := exec.CommandContext(ctx, interpreter, "-c", string(command))
 	cmd.Dir = b.Config.Path
 

libs/git/clone.go

@@ -1,13 +1,14 @@
 package git
 
 import (
-	"bytes"
 	"context"
 	"errors"
 	"fmt"
 	"os/exec"
 	"regexp"
 	"strings"
+
+	"github.com/databricks/cli/libs/process"
 )
 
 // source: https://stackoverflow.com/questions/59081778/rules-for-special-characters-in-github-repository-name
@@ -42,24 +43,14 @@ func (opts cloneOptions) args() []string {
 }
 
 func (opts cloneOptions) clone(ctx context.Context) error {
-	cmd := exec.CommandContext(ctx, "git", opts.args()...)
-	var cmdErr bytes.Buffer
-	cmd.Stderr = &cmdErr
-
-	// start git clone
-	err := cmd.Start()
+	// start and wait for git clone to complete
+	_, err := process.Background(ctx, append([]string{"git"}, opts.args()...))
 	if errors.Is(err, exec.ErrNotFound) {
 		return fmt.Errorf("please install git CLI to clone a repository: %w", err)
 	}
 	if err != nil {
 		return fmt.Errorf("git clone failed: %w", err)
 	}
-
-	// wait for git clone to complete
-	err = cmd.Wait()
-	if err != nil {
-		return fmt.Errorf("git clone failed: %w. %s", err, cmdErr.String())
-	}
 	return nil
 }
 

libs/process/background.go

@@ -0,0 +1,32 @@
+package process
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"os"
+	"os/exec"
+	"strings"
+
+	"github.com/databricks/cli/libs/log"
+)
+
+func Background(ctx context.Context, args []string, opts ...execOption) (string, error) {
+	commandStr := strings.Join(args, " ")
+	log.Debugf(ctx, "running: %s", commandStr)
+	cmd := exec.CommandContext(ctx, args[0], args[1:]...)
+	stdout := &bytes.Buffer{}
+	cmd.Stdin = os.Stdin
+	cmd.Stdout = stdout
+	cmd.Stderr = stdout
+	for _, o := range opts {
+		err := o(cmd)
+		if err != nil {
+			return "", err
+		}
+	}
+	if err := cmd.Run(); err != nil {
+		return "", fmt.Errorf("%s: %s %w", commandStr, stdout.String(), err)
+	}
+	return strings.TrimSpace(stdout.String()), nil
+}

libs/process/background_test.go

@@ -0,0 +1,31 @@
+package process
+
+import (
+	"context"
+	"fmt"
+	"os/exec"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBackground(t *testing.T) {
+	ctx := context.Background()
+	res, err := Background(ctx, []string{"echo", "1"}, WithDir("/"))
+	assert.NoError(t, err)
+	assert.Equal(t, "1", res)
+}
+
+func TestBackgroundFails(t *testing.T) {
+	ctx := context.Background()
+	_, err := Background(ctx, []string{"ls", "/dev/null/x"})
+	assert.NotNil(t, err)
+}
+
+func TestBackgroundFailsOnOption(t *testing.T) {
+	ctx := context.Background()
+	_, err := Background(ctx, []string{"ls", "/dev/null/x"}, func(c *exec.Cmd) error {
+		return fmt.Errorf("nope")
+	})
+	assert.EqualError(t, err, "nope")
+}

libs/process/forwarded.go

@@ -0,0 +1,48 @@
+package process
+
+import (
+	"context"
+	"io"
+	"os/exec"
+	"strings"
+
+	"github.com/databricks/cli/libs/log"
+)
+
+func Forwarded(ctx context.Context, args []string, src io.Reader, dst io.Writer, opts ...execOption) error {
+	commandStr := strings.Join(args, " ")
+	log.Debugf(ctx, "starting: %s", commandStr)
+	cmd := exec.CommandContext(ctx, args[0], args[1:]...)
+
+	// make sure to sync on writing to stdout
+	reader, writer := io.Pipe()
+	go io.CopyBuffer(dst, reader, make([]byte, 128))
+	defer reader.Close()
+	defer writer.Close()
+	cmd.Stdout = writer
+	cmd.Stderr = writer
+
+	// apply common options
+	for _, o := range opts {
+		err := o(cmd)
+		if err != nil {
+			return err
+		}
+	}
+
+	// pipe standard input to the child process, so that we can allow terminal UX
+	// see the PoC at https://github.com/databricks/cli/pull/637
+	stdin, err := cmd.StdinPipe()
+	if err != nil {
+		return err
+	}
+	go io.CopyBuffer(stdin, src, make([]byte, 128))
+	defer stdin.Close()
+
+	err = cmd.Start()
+	if err != nil {
+		return err
+	}
+
+	return cmd.Wait()
+}

libs/process/forwarded_test.go

@@ -0,0 +1,43 @@
+package process
+
+import (
+	"bytes"
+	"context"
+	"os/exec"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestForwarded(t *testing.T) {
+	ctx := context.Background()
+	buf := bytes.NewBufferString("")
+	err := Forwarded(ctx, []string{
+		"python3", "-c", "print(input('input: '))",
+	}, strings.NewReader("abc\n"), buf)
+	assert.NoError(t, err)
+
+	assert.Equal(t, "input: abc\n", buf.String())
+}
+
+func TestForwardedFails(t *testing.T) {
+	ctx := context.Background()
+	buf := bytes.NewBufferString("")
+	err := Forwarded(ctx, []string{
+		"_non_existent_",
+	}, strings.NewReader("abc\n"), buf)
+	assert.NotNil(t, err)
+}
+
+func TestForwardedFailsOnStdinPipe(t *testing.T) {
+	ctx := context.Background()
+	buf := bytes.NewBufferString("")
+	err := Forwarded(ctx, []string{
+		"_non_existent_",
+	}, strings.NewReader("abc\n"), buf, func(c *exec.Cmd) error {
+		c.Stdin = strings.NewReader("x")
+		return nil
+	})
+	assert.NotNil(t, err)
+}

libs/process/opts.go

@@ -0,0 +1,51 @@
+package process
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"os/exec"
+)
+
+type execOption func(*exec.Cmd) error
+
+func WithEnv(key, value string) execOption {
+	return func(c *exec.Cmd) error {
+		if c.Env == nil {
+			c.Env = os.Environ()
+		}
+		v := fmt.Sprintf("%s=%s", key, value)
+		c.Env = append(c.Env, v)
+		return nil
+	}
+}
+
+func WithEnvs(envs map[string]string) execOption {
+	return func(c *exec.Cmd) error {
+		for k, v := range envs {
+			err := WithEnv(k, v)(c)
+			if err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+}
+
+func WithDir(dir string) execOption {
+	return func(c *exec.Cmd) error {
+		c.Dir = dir
+		return nil
+	}
+}
+
+func WithStdoutPipe(dst *io.ReadCloser) execOption {
+	return func(c *exec.Cmd) error {
+		outPipe, err := c.StdoutPipe()
+		if err != nil {
+			return err
+		}
+		*dst = outPipe
+		return nil
+	}
+}

libs/process/opts_test.go

@@ -0,0 +1,18 @@
+package process
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWithEnvs(t *testing.T) {
+	ctx := context.Background()
+	res, err := Background(ctx, []string{"/bin/sh", "-c", "echo $FOO $BAR"}, WithEnvs(map[string]string{
+		"FOO": "foo",
+		"BAR": "delirium",
+	}))
+	assert.NoError(t, err)
+	assert.Equal(t, "foo delirium", res)
+}

python/runner.go

@@ -8,6 +8,8 @@ import (
 	"os/exec"
 	"runtime"
 	"strings"
+
+	"github.com/databricks/cli/libs/process"
 )
 
 func PyInline(ctx context.Context, inlinePy string) (string, error) {
@@ -88,8 +90,8 @@ func DetectExecutable(ctx context.Context) (string, error) {
 
 func execAndPassErr(ctx context.Context, name string, args ...string) ([]byte, error) {
 	// TODO: move out to a separate package, once we have Maven integration
-	out, err := exec.CommandContext(ctx, name, args...).Output()
-	return out, nicerErr(err)
+	out, err := process.Background(ctx, append([]string{name}, args...))
+	return []byte(out), nicerErr(err)
 }
 
 func getFirstMatch(out string) string {