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

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,14 @@ 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()
+		var buf bytes.Buffer
+		_, err := process.Background(ctx, buildParts,
+			process.WithCombinedOutput(&buf),
+			process.WithDir(a.Path))
 		if err != nil {
-			return res, err
+			return buf.Bytes(), err
 		}
-		out = append(out, res)
+		out = append(out, buf.Bytes())
 	}
 	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/env/context.go

@@ -3,6 +3,7 @@ package env
 import (
 	"context"
 	"os"
+	"strings"
 )
 
 var envContextKey int
@@ -61,3 +62,21 @@ func Set(ctx context.Context, key, value string) context.Context {
 	m[key] = value
 	return setMap(ctx, m)
 }
+
+// All returns environment variables that are defined in both os.Environ
+// and this package. `env.Set(ctx, x, y)` will override x from os.Environ.
+func All(ctx context.Context) map[string]string {
+	m := map[string]string{}
+	for _, line := range os.Environ() {
+		split := strings.SplitN(line, "=", 2)
+		if len(split) != 2 {
+			continue
+		}
+		m[split[0]] = split[1]
+	}
+	// override existing environment variables with the ones we set
+	for k, v := range copyMap(getMap(ctx)) {
+		m[k] = v
+	}
+	return m
+}

libs/env/context_test.go

@@ -38,4 +38,11 @@ func TestContext(t *testing.T) {
 	assert.Equal(t, "qux", Get(ctx2, "FOO"))
 	assert.Equal(t, "baz", Get(ctx1, "FOO"))
 	assert.Equal(t, "bar", Get(ctx0, "FOO"))
+
+	ctx3 := Set(ctx2, "BAR", "x=y")
+
+	all := All(ctx3)
+	assert.NotNil(t, all)
+	assert.Equal(t, "qux", all["FOO"])
+	assert.Equal(t, "x=y", all["BAR"])
 }

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,23 +43,17 @@ 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)
+	var processErr *process.ProcessError
+	if errors.As(err, &processErr) {
+		return fmt.Errorf("git clone failed: %w. %s", err, processErr.Stderr)
 	}
-
-	// wait for git clone to complete
-	err = cmd.Wait()
 	if err != nil {
-		return fmt.Errorf("git clone failed: %w. %s", err, cmdErr.String())
+		return fmt.Errorf("git clone failed: %w", err)
 	}
 	return nil
 }

libs/process/background.go

@@ -0,0 +1,52 @@
+package process
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"os/exec"
+	"strings"
+
+	"github.com/databricks/cli/libs/log"
+)
+
+type ProcessError struct {
+	Command string
+	Err     error
+	Stdout  string
+	Stderr  string
+}
+
+func (perr *ProcessError) Unwrap() error {
+	return perr.Err
+}
+
+func (perr *ProcessError) Error() string {
+	return fmt.Sprintf("%s: %s %s", perr.Command, perr.Stderr, perr.Err)
+}
+
+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:]...)
+	var stdout, stderr bytes.Buffer
+	// For background processes, there's no standard input
+	cmd.Stdin = nil
+	cmd.Stdout = &stdout
+	cmd.Stderr = &stderr
+	for _, o := range opts {
+		err := o(ctx, cmd)
+		if err != nil {
+			return "", err
+		}
+	}
+	if err := cmd.Run(); err != nil {
+		return stdout.String(), &ProcessError{
+			Err:     err,
+			Command: commandStr,
+			Stdout:  stdout.String(),
+			Stderr:  stderr.String(),
+		}
+	}
+	return stdout.String(), nil
+}

libs/process/background_test.go

@@ -0,0 +1,81 @@
+package process
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"os"
+	"os/exec"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBackgroundUnwrapsNotFound(t *testing.T) {
+	ctx := context.Background()
+	_, err := Background(ctx, []string{"/bin/meeecho", "1"})
+	assert.ErrorIs(t, err, os.ErrNotExist)
+}
+
+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 TestBackgroundOnlyStdoutGetsoutOnSuccess(t *testing.T) {
+	ctx := context.Background()
+	res, err := Background(ctx, []string{
+		"python3", "-c", "import sys; sys.stderr.write('1'); sys.stdout.write('2')",
+	})
+	assert.NoError(t, err)
+	assert.Equal(t, "2", res)
+}
+
+func TestBackgroundCombinedOutput(t *testing.T) {
+	ctx := context.Background()
+	var buf bytes.Buffer
+	res, err := Background(ctx, []string{
+		"python3", "-c", "import sys; sys.stderr.write('1'); sys.stdout.write('2')",
+	}, WithCombinedOutput(&buf))
+	assert.NoError(t, err)
+	assert.Equal(t, "2", res)
+	assert.Equal(t, "12", buf.String())
+}
+
+func TestBackgroundCombinedOutputFailure(t *testing.T) {
+	ctx := context.Background()
+	var buf bytes.Buffer
+	res, err := Background(ctx, []string{
+		"python3", "-c", "import sys; sys.stderr.write('1'); sys.stdout.write('2'); sys.exit(42)",
+	}, WithCombinedOutput(&buf))
+	var processErr *ProcessError
+	if assert.ErrorAs(t, err, &processErr) {
+		assert.Equal(t, "1", processErr.Stderr)
+		assert.Equal(t, "2", processErr.Stdout)
+	}
+	assert.Equal(t, "2", res)
+	assert.Equal(t, "12", buf.String())
+}
+
+func TestBackgroundNoStdin(t *testing.T) {
+	ctx := context.Background()
+	res, err := Background(ctx, []string{"cat"})
+	assert.NoError(t, err)
+	assert.Equal(t, "", 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(_ context.Context, c *exec.Cmd) error {
+		return fmt.Errorf("nope")
+	})
+	assert.EqualError(t, err, "nope")
+}

libs/process/forwarded.go

@@ -0,0 +1,42 @@
+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()
+
+	// empirical tests showed buffered copies being more responsive
+	go io.CopyBuffer(dst, reader, make([]byte, 128))
+	defer reader.Close()
+	defer writer.Close()
+	cmd.Stdout = writer
+	cmd.Stderr = writer
+	cmd.Stdin = src
+
+	// apply common options
+	for _, o := range opts {
+		err := o(ctx, cmd)
+		if err != nil {
+			return err
+		}
+	}
+
+	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()
+	var buf bytes.Buffer
+	err := Forwarded(ctx, []string{
+		"python3", "-c", "print(input('input: '))",
+	}, strings.NewReader("abc\n"), &buf)
+	assert.NoError(t, err)
+
+	assert.Equal(t, "input: abc", strings.TrimSpace(buf.String()))
+}
+
+func TestForwardedFails(t *testing.T) {
+	ctx := context.Background()
+	var buf bytes.Buffer
+	err := Forwarded(ctx, []string{
+		"_non_existent_",
+	}, strings.NewReader("abc\n"), &buf)
+	assert.NotNil(t, err)
+}
+
+func TestForwardedFailsOnStdinPipe(t *testing.T) {
+	ctx := context.Background()
+	var buf bytes.Buffer
+	err := Forwarded(ctx, []string{
+		"_non_existent_",
+	}, strings.NewReader("abc\n"), &buf, func(_ context.Context, c *exec.Cmd) error {
+		c.Stdin = strings.NewReader("x")
+		return nil
+	})
+	assert.NotNil(t, err)
+}