feat: stream output for custom workflows #2261
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,161 @@ | ||
| package models | ||
|
|
||
| import ( | ||
| "bufio" | ||
| "io" | ||
| "os/exec" | ||
| "strings" | ||
| "sync" | ||
|
|
||
| "github.com/pkg/errors" | ||
| "github.com/runatlantis/atlantis/server/events/command" | ||
| "github.com/runatlantis/atlantis/server/events/terraform/ansi" | ||
| "github.com/runatlantis/atlantis/server/jobs" | ||
| ) | ||
|
|
||
| // Setting the buffer size to 10mb | ||
| const BufioScannerBufferSize = 10 * 1024 * 1024 | ||
|
|
||
| // Line represents a line that was output from a shell command. | ||
| type Line struct { | ||
| // Line is the contents of the line (without the newline). | ||
| Line string | ||
| // Err is set if there was an error. | ||
| Err error | ||
| } | ||
|
|
||
| // ShellCommandRunner runs a command via `exec.Command` and streams output to the | ||
| // `ProjectCommandOutputHandler`. | ||
| type ShellCommandRunner struct { | ||
| command string | ||
| workingDir string | ||
| outputHandler jobs.ProjectCommandOutputHandler | ||
| streamOutput bool | ||
| cmd *exec.Cmd | ||
| } | ||
|
|
||
| func NewShellCommandRunner(command string, environ []string, workingDir string, streamOutput bool, outputHandler jobs.ProjectCommandOutputHandler) *ShellCommandRunner { | ||
| cmd := exec.Command("sh", "-c", command) // #nosec | ||
| cmd.Env = environ | ||
| cmd.Dir = workingDir | ||
|
|
||
| return &ShellCommandRunner{ | ||
| command: command, | ||
| workingDir: workingDir, | ||
| outputHandler: outputHandler, | ||
| streamOutput: streamOutput, | ||
| cmd: cmd, | ||
| } | ||
| } | ||
|
|
||
| func (s *ShellCommandRunner) Run(ctx command.ProjectContext) (string, error) { | ||
| _, outCh := s.RunCommandAsync(ctx) | ||
|
|
||
| outbuf := new(strings.Builder) | ||
| var err error | ||
| for line := range outCh { | ||
| if line.Err != nil { | ||
| err = line.Err | ||
| break | ||
| } | ||
| outbuf.WriteString(line.Line) | ||
| outbuf.WriteString("\n") | ||
| } | ||
|
|
||
| // sanitize output by stripping out any ansi characters. | ||
| output := ansi.Strip(outbuf.String()) | ||
| return output, err | ||
| } | ||
|
|
||
| // RunCommandAsync runs terraform with args. It immediately returns an | ||
| // input and output channel. Callers can use the output channel to | ||
| // get the realtime output from the command. | ||
| // Callers can use the input channel to pass stdin input to the command. | ||
| // If any error is passed on the out channel, there will be no | ||
| // further output (so callers are free to exit). | ||
| func (s *ShellCommandRunner) RunCommandAsync(ctx command.ProjectContext) (chan<- string, <-chan Line) { | ||
|
There was a problem hiding this comment. This code was extracted as-is from |
||
| outCh := make(chan Line) | ||
| inCh := make(chan string) | ||
|
|
||
| // We start a goroutine to do our work asynchronously and then immediately | ||
| // return our channels. | ||
| go func() { | ||
| // Ensure we close our channels when we exit. | ||
| defer func() { | ||
| close(outCh) | ||
| close(inCh) | ||
| }() | ||
|
|
||
| stdout, _ := s.cmd.StdoutPipe() | ||
| stderr, _ := s.cmd.StderrPipe() | ||
| stdin, _ := s.cmd.StdinPipe() | ||
|
|
||
| ctx.Log.Debug("starting %q in %q", s.command, s.workingDir) | ||
| err := s.cmd.Start() | ||
| if err != nil { | ||
| err = errors.Wrapf(err, "running %q in %q", s.command, s.workingDir) | ||
| ctx.Log.Err(err.Error()) | ||
| outCh <- Line{Err: err} | ||
| return | ||
| } | ||
|
|
||
| // If we get anything on inCh, write it to stdin. | ||
| // This function will exit when inCh is closed which we do in our defer. | ||
| go func() { | ||
| for line := range inCh { | ||
| ctx.Log.Debug("writing %q to remote command's stdin", line) | ||
| _, err := io.WriteString(stdin, line) | ||
| if err != nil { | ||
| ctx.Log.Err(errors.Wrapf(err, "writing %q to process", line).Error()) | ||
| } | ||
| } | ||
| }() | ||
|
|
||
| wg := new(sync.WaitGroup) | ||
| wg.Add(2) | ||
| // Asynchronously copy from stdout/err to outCh. | ||
| go func() { | ||
| scanner := bufio.NewScanner(stdout) | ||
| buf := []byte{} | ||
| scanner.Buffer(buf, BufioScannerBufferSize) | ||
|
|
||
| for scanner.Scan() { | ||
| message := scanner.Text() | ||
| outCh <- Line{Line: message} | ||
| if s.streamOutput { | ||
| s.outputHandler.Send(ctx, message, false) | ||
| } | ||
| } | ||
| wg.Done() | ||
| }() | ||
| go func() { | ||
| scanner := bufio.NewScanner(stderr) | ||
| for scanner.Scan() { | ||
| message := scanner.Text() | ||
| outCh <- Line{Line: message} | ||
| if s.streamOutput { | ||
| s.outputHandler.Send(ctx, message, false) | ||
| } | ||
| } | ||
| wg.Done() | ||
| }() | ||
|
|
||
| // Wait for our copying to complete. This *must* be done before | ||
| // calling cmd.Wait(). (see https://github.com/golang/go/issues/19685) | ||
| wg.Wait() | ||
|
|
||
| // Wait for the command to complete. | ||
| err = s.cmd.Wait() | ||
|
|
||
| // We're done now. Send an error if there was one. | ||
| if err != nil { | ||
| err = errors.Wrapf(err, "running %q in %q", s.command, s.workingDir) | ||
| ctx.Log.Err(err.Error()) | ||
| outCh <- Line{Err: err} | ||
| } else { | ||
| ctx.Log.Info("successfully ran %q in %q", s.command, s.workingDir) | ||
| } | ||
| }() | ||
|
|
||
| return inCh, outCh | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,75 @@ | ||
| package models_test | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "os" | ||
| "strings" | ||
| "testing" | ||
|
|
||
| . "github.com/petergtz/pegomock" | ||
| "github.com/runatlantis/atlantis/server/core/runtime/mocks/matchers" | ||
| "github.com/runatlantis/atlantis/server/core/runtime/models" | ||
| "github.com/runatlantis/atlantis/server/events/command" | ||
| "github.com/runatlantis/atlantis/server/jobs/mocks" | ||
| "github.com/runatlantis/atlantis/server/logging" | ||
| . "github.com/runatlantis/atlantis/testing" | ||
| ) | ||
|
|
||
| func TestShellCommandRunner_Run(t *testing.T) { | ||
| cases := []struct { | ||
| Command string | ||
| ExpLines []string | ||
| Environ map[string]string | ||
| }{ | ||
| { | ||
| Command: "echo $HELLO", | ||
| Environ: map[string]string{ | ||
| "HELLO": "world", | ||
| }, | ||
| ExpLines: []string{"world"}, | ||
| }, | ||
| { | ||
| Command: ">&2 echo this is an error", | ||
| ExpLines: []string{"this is an error"}, | ||
| }, | ||
| } | ||
|
|
||
| for _, c := range cases { | ||
| t.Run(c.Command, func(t *testing.T) { | ||
| RegisterMockTestingT(t) | ||
| ctx := command.ProjectContext{ | ||
| Log: logging.NewNoopLogger(t), | ||
| Workspace: "default", | ||
| RepoRelDir: ".", | ||
| } | ||
| projectCmdOutputHandler := mocks.NewMockProjectCommandOutputHandler() | ||
|
|
||
| cwd, err := os.Getwd() | ||
| Ok(t, err) | ||
| environ := []string{} | ||
| for k, v := range c.Environ { | ||
| environ = append(environ, fmt.Sprintf("%s=%s", k, v)) | ||
| } | ||
| expectedOutput := fmt.Sprintf("%s\n", strings.Join(c.ExpLines, "\n")) | ||
|
|
||
| // Run once with streaming enabled | ||
| runner := models.NewShellCommandRunner(c.Command, environ, cwd, true, projectCmdOutputHandler) | ||
| output, err := runner.Run(ctx) | ||
| Ok(t, err) | ||
| Equals(t, expectedOutput, output) | ||
| for _, line := range c.ExpLines { | ||
| projectCmdOutputHandler.VerifyWasCalledOnce().Send(ctx, line, false) | ||
| } | ||
|
|
||
| // And again with streaming disabled. Everything should be the same except the | ||
| // command output handler should not have received anything | ||
|
|
||
| projectCmdOutputHandler = mocks.NewMockProjectCommandOutputHandler() | ||
| runner = models.NewShellCommandRunner(c.Command, environ, cwd, false, projectCmdOutputHandler) | ||
| output, err = runner.Run(ctx) | ||
| Ok(t, err) | ||
| Equals(t, expectedOutput, output) | ||
| projectCmdOutputHandler.VerifyWasCalled(Never()).Send(matchers.AnyModelsProjectCommandContext(), AnyString(), EqBool(false)) | ||
| }) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't think this is a particularly good place for this, but noticed it was where the other
execfunctionality was put. Not totally sure on the naming convention of this project, but it couldn't be inruntimebecause of import cycles with theterraformpackage.