-
Notifications
You must be signed in to change notification settings - Fork 15
/
cmd.go
111 lines (93 loc) · 2.27 KB
/
cmd.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
package util
import (
"bufio"
"context"
"os"
"os/exec"
"strings"
log "github.com/sirupsen/logrus"
)
// Printer is a function that prints out a string, e.g. to stdout.
type Printer func(input string)
// LogrusInfoPrinter returns a printer that prints via logrus at the info level.
func LogrusInfoPrinter(prefix string) Printer {
return func(input string) {
log.Infof("%s %s", prefix, input)
}
}
// LogrusWarnPrinter returns a printer that prints via logrus at the warn level.
func LogrusWarnPrinter(prefix string) Printer {
return func(input string) {
log.Warnf("%s %s", prefix, input)
}
}
// LogrusDebugPrinter returns a printer that prints via logrus at the debug level.
func LogrusDebugPrinter(prefix string) Printer {
return func(input string) {
log.Debugf("%s %s", prefix, input)
}
}
// RunCmdWithPrinters runs a command with output streamed via custom printer functions.
// Adapted from example in https://github.com/golang/go/issues/19685#issuecomment-288949629.
func RunCmdWithPrinters(
ctx context.Context,
command string,
args []string,
extraEnv []string,
blockedEnv map[string]struct{},
stdoutPrinter Printer,
stderrPrinter Printer,
) error {
log.Debugf("Running %s with args %+v", command, args)
cmd := exec.CommandContext(ctx, command, args...)
envVars := os.Environ()
cmdEnvVars := []string{}
for _, envVar := range envVars {
components := strings.SplitN(envVar, "=", 2)
if _, ok := blockedEnv[components[0]]; !ok {
cmdEnvVars = append(cmdEnvVars, envVar)
}
}
cmdEnvVars = append(
cmdEnvVars,
extraEnv...,
)
cmd.Env = cmdEnvVars
cmd.Stdin = os.Stdin
done := make(chan struct{}, 2)
stdoutPipe, err := cmd.StdoutPipe()
if err != nil {
return err
}
stdoutScanner := bufio.NewScanner(stdoutPipe)
go func() {
for stdoutScanner.Scan() {
stdoutPrinter(stdoutScanner.Text())
}
done <- struct{}{}
}()
stderrPipe, err := cmd.StderrPipe()
if err != nil {
return err
}
stderrScanner := bufio.NewScanner(stderrPipe)
go func() {
for stderrScanner.Scan() {
stderrPrinter(stderrScanner.Text())
}
done <- struct{}{}
}()
err = cmd.Start()
if err != nil {
return err
}
// Wait for all scanner goroutines to finish
for i := 0; i < 2; i++ {
select {
case <-done:
case <-ctx.Done():
break
}
}
return cmd.Wait()
}