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 getMap(ctx) {
+ m[k] = v
+	}
+	return m
+}

libs/env/context_test.go

@@ -38,4 +38,12 @@ 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"])
+	assert.NotEmpty(t, all["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,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,59 @@
+package process
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"os/exec"
+	"strings"
+
+	"github.com/databricks/cli/libs/env"
+	"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", perr.Command, 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:]...)
+	stdout := bytes.Buffer{}
+	stderr := bytes.Buffer{}
+	// For background processes, there's no standard input
+	cmd.Stdin = nil
+	cmd.Stdout = &stdout
+	cmd.Stderr = &stderr
+	// we pull the env through lib/env such that we can run
+	// parallel tests with anything using libs/process.
+	for k, v := range env.All(ctx) {
+ cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, v))
+	}
+	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,91 @@
+package process
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"os"
+	"os/exec"
+	"strings"
+	"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", strings.TrimSpace(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()
+	buf := bytes.Buffer{}
+	res, err := Background(ctx, []string{
+ "python3", "-c", "import sys, time; " +
+ `sys.stderr.write("1\n"); sys.stderr.flush(); ` +
+ "time.sleep(0.001); " +
+ "print('2', flush=True); sys.stdout.flush(); " +
+ "time.sleep(0.001)",
+	}, WithCombinedOutput(&buf))
+	assert.NoError(t, err)
+	assert.Equal(t, "2", strings.TrimSpace(res))
+	assert.Equal(t, "1\n2\n", strings.ReplaceAll(buf.String(), "\r", ""))
+}
+
+func TestBackgroundCombinedOutputFailure(t *testing.T) {
+	ctx := context.Background()
+	buf := bytes.Buffer{}
+	res, err := Background(ctx, []string{
+ "python3", "-c", "import sys, time; " +
+ `sys.stderr.write("1\n"); sys.stderr.flush(); ` +
+ "time.sleep(0.001); " +
+ "print('2', flush=True); sys.stdout.flush(); " +
+ "time.sleep(0.001); " +
+ "sys.exit(42)",
+	}, WithCombinedOutput(&buf))
+	var processErr *ProcessError
+	if assert.ErrorAs(t, err, &processErr) {
+ assert.Equal(t, "1", strings.TrimSpace(processErr.Stderr))
+ assert.Equal(t, "2", strings.TrimSpace(processErr.Stdout))
+	}
+	assert.Equal(t, "2", strings.TrimSpace(res))
+	assert.Equal(t, "1\n2\n", strings.ReplaceAll(buf.String(), "\r", ""))
+}
+
+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,43 @@
+package process
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"os/exec"
+	"strings"
+
+	"github.com/databricks/cli/libs/env"
+	"github.com/databricks/cli/libs/log"
+)
+
+func Forwarded(ctx context.Context, args []string, src io.Reader, outWriter, errWriter io.Writer, opts ...execOption) error {
+	commandStr := strings.Join(args, " ")
+	log.Debugf(ctx, "starting: %s", commandStr)
+	cmd := exec.CommandContext(ctx, args[0], args[1:]...)
+
+	// empirical tests showed buffered copies being more responsive
+	cmd.Stdout = outWriter
+	cmd.Stderr = errWriter
+	cmd.Stdin = src
+	// we pull the env through lib/env such that we can run
+	// parallel tests with anything using libs/process.
+	for k, v := range env.All(ctx) {
+ cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, v))
+	}
+
+	// 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, &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, &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, &buf, func(_ context.Context, c *exec.Cmd) error {
+ c.Stdin = strings.NewReader("x")
+ return nil
+	})
+	assert.NotNil(t, err)
+}