From a2748d531b31de50f7e46d729937d380c3c5c4af Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Mon, 3 Mar 2025 19:19:39 +0100
Subject: [PATCH] cleanup

---
 acceptance/bundle/exec/basic/output.txt |  3 ++
 acceptance/bundle/exec/basic/script     |  5 +++-
 cmd/bundle/exec.go                      | 37 +++++++++++++++++++------
 3 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/acceptance/bundle/exec/basic/output.txt b/acceptance/bundle/exec/basic/output.txt
index 3584582db..9283131d6 100644
--- a/acceptance/bundle/exec/basic/output.txt
+++ b/acceptance/bundle/exec/basic/output.txt
@@ -12,3 +12,6 @@ hello, world
 Error: Running "bash -c exit 5" failed with exit code: 5
 
 Exit code: 1
+
+>>> [CLI] bundle exec -- bash -c echo hello > /dev/stderr
+hello
diff --git a/acceptance/bundle/exec/basic/script b/acceptance/bundle/exec/basic/script
index 41c57725f..ebe76884b 100644
--- a/acceptance/bundle/exec/basic/script
+++ b/acceptance/bundle/exec/basic/script
@@ -6,4 +6,7 @@ trace $CLI bundle exec -- pwd
 trace $CLI bundle exec -- echo --help
 
 # The error message should include the exit code.
-trace $CLI bundle exec -- bash -c "exit 5"
+errcode trace $CLI bundle exec -- bash -c "exit 5"
+
+# stderr should also be shown in the output.
+trace $CLI bundle exec -- bash -c "echo hello > /dev/stderr"
diff --git a/cmd/bundle/exec.go b/cmd/bundle/exec.go
index bd96b41de..fbb4d2547 100644
--- a/cmd/bundle/exec.go
+++ b/cmd/bundle/exec.go
@@ -27,11 +27,12 @@ func newExecCommand() *cobra.Command {
 		Use:   "exec",
 		Short: "Execute a command using the same authentication context as the bundle",
 		Args:  cobra.MinimumNArgs(1),
-		// TODO: format once we have all the documentation here.
-		Long: `
-Note: This command executes scripts
+		Long: `Execute a command using the same authentication context as the bundle
 
-Examples:
+The current working directory of the provided command will be set to the root
+of the bundle.
+
+Example usage:
 1. databricks bundle exec -- echo hello
 2. databricks bundle exec -- /bin/bash -c "echo hello""
 3. databricks bundle exec -- uv run pytest"`,
@@ -81,8 +82,8 @@ Examples:
 			// adding support for the scripts section.
 			childCmd.Dir = b.BundleRootPath
 
-			// Create pipes for stdout and stderr.
-			// TODO: Test streaming of this? Is there a way?
+			// Create pipes to stream the stdout and stderr output from the child
+			// process.
 			stdout, err := childCmd.StdoutPipe()
 			if err != nil {
 				return fmt.Errorf("Error creating stdout pipe: %w", err)
@@ -98,15 +99,21 @@ Examples:
 				return fmt.Errorf("Error starting command: %s\n", err)
 			}
 
-			// Stream both stdout and stderr to the user.
+			// Stream both stdout and stderr to the user. We do this so that the user
+			// does not have to wait for the command to finish before seeing the output.
 			var wg sync.WaitGroup
+			var stdoutErr, stderrErr error
 			wg.Add(2)
 
 			go func() {
 				defer wg.Done()
 				scanner := bufio.NewScanner(stdout)
 				for scanner.Scan() {
-					fmt.Println(scanner.Text())
+					_, err = cmd.OutOrStdout().Write([]byte(scanner.Text() + "\n"))
+					if err != nil {
+						stdoutErr = fmt.Errorf("Error writing to stdout: %w", err)
+						return
+					}
 				}
 			}()
 
@@ -114,10 +121,22 @@ Examples:
 				defer wg.Done()
 				scanner := bufio.NewScanner(stderr)
 				for scanner.Scan() {
-					fmt.Println(scanner.Text())
+					_, err := cmd.ErrOrStderr().Write([]byte(scanner.Text() + "\n"))
+					if err != nil {
+						stderrErr = fmt.Errorf("Error writing to stderr: %w", err)
+						return
+					}
 				}
 			}()
 
+			if stdoutErr != nil {
+				return stdoutErr
+			}
+
+			if stderrErr != nil {
+				return stderrErr
+			}
+
 			// Wait for the command to finish.
 			err = childCmd.Wait()
 			if exitErr, ok := err.(*exec.ExitError); ok {