Fix runtime spinner issues (#1779)

Co-authored-by: Pritesh Arora <[email protected]>
Co-authored-by: Greg Neiheisel <[email protected]>

airflow/runtimes/command.go

@@ -1,7 +1,10 @@
 package runtimes
 
 import (
+	"bufio"
 	"bytes"
+	"fmt"
+	"io"
 	"os/exec"
 )
 
@@ -20,3 +23,54 @@ func (p *Command) Execute() (string, error) {
 	err := cmd.Run()
 	return out.String(), err
 }
+
+func (p *Command) ExecuteWithProgress(progressHandler func(string)) error {
+	cmd := exec.Command(p.Command, p.Args...) //nolint:gosec
+
+	// Create pipes for stdout and stderr
+	stdoutPipe, err := cmd.StdoutPipe()
+	if err != nil {
+		return fmt.Errorf("error creating stdout pipe: %w", err)
+	}
+	stderrPipe, err := cmd.StderrPipe()
+	if err != nil {
+		return fmt.Errorf("error creating stderr pipe: %w", err)
+	}
+
+	// Start the command
+	if err := cmd.Start(); err != nil {
+		return fmt.Errorf("error setting up astro project: %w", err)
+	}
+
+	// Stream stdout and stderr concurrently
+	doneCh := make(chan error, 2)
+	go streamOutput(stdoutPipe, progressHandler, doneCh)
+	go streamOutput(stderrPipe, progressHandler, doneCh)
+
+	// Wait for both streams to finish
+	for i := 0; i < 2; i++ {
+		if err := <-doneCh; err != nil {
+			return err
+		}
+	}
+	// Wait for the command to complete
+	if err := cmd.Wait(); err != nil {
+		return fmt.Errorf("astro project execution failed: %w", err)
+	}
+
+	return nil
+}
+
+func streamOutput(pipe io.ReadCloser, handler func(string), doneCh chan<- error) {
+	scanner := bufio.NewScanner(pipe)
+	for scanner.Scan() {
+		line := scanner.Text()
+		handler(line)
+	}
+	// Notify completion or error
+	if err := scanner.Err(); err != nil {
+		doneCh <- fmt.Errorf("error reading output: %w", err)
+		return
+	}
+	doneCh <- nil
+}

airflow/runtimes/container_runtime.go

@@ -24,9 +24,8 @@ const (
 	containerRuntimeNotFoundErrMsg = "Failed to find a container runtime. " +
 		"See the Astro CLI prerequisites for more information. " +
 		"https://www.astronomer.io/docs/astro/cli/install-cli"
-	containerRuntimeInitMessage = " Astro uses container technology to run your Airflow project. " +
-		"Please wait while we get things started…"
-	spinnerRefresh = 100 * time.Millisecond
+	containerRuntimeInitMessage = "Astro uses containers to run your project. Please wait while we get started…"
+	spinnerRefresh              = 100 * time.Millisecond
 )
 
 // ContainerRuntime interface defines the methods that manage

airflow/runtimes/docker_runtime.go

@@ -97,7 +97,7 @@ func (rt DockerRuntime) initializeDocker(timeoutSeconds int) error {
 	timeout := time.After(time.Duration(timeoutSeconds) * time.Second)
 	ticker := time.NewTicker(time.Duration(tickNum) * time.Millisecond)
 	s := spinner.New(spinnerCharSet, spinnerRefresh)
-	s.Suffix = containerRuntimeInitMessage
+	s.Suffix = " " + containerRuntimeInitMessage
 	defer s.Stop()
 
 	// Execute `docker ps` to check if Docker is running.

airflow/runtimes/podman_engine.go

@@ -7,14 +7,13 @@ import (
 
 	"github.com/astronomer/astro-cli/airflow/runtimes/types"
 	"github.com/astronomer/astro-cli/config"
+	"github.com/briandowns/spinner"
 )
 
 const (
-	podmanStatusRunning   = "running"
-	podmanStatusStopped   = "stopped"
-	composeProjectLabel   = "com.docker.compose.project"
-	podmanInitSlowMessage = " Sorry for the wait, this is taking a bit longer than expected. " +
-		"This initial download will be cached once finished."
+	podmanStatusRunning               = "running"
+	podmanStatusStopped               = "stopped"
+	composeProjectLabel               = "com.docker.compose.project"
 	podmanMachineAlreadyRunningErrMsg = "astro needs a podman machine to run your project, " +
 		"but it looks like a machine is already running. " +
 		"Mac hosts are limited to one running machine at a time. " +
@@ -25,7 +24,7 @@ const (
 type podmanEngine struct{}
 
 // InitializeMachine initializes our astro Podman machine.
-func (e podmanEngine) InitializeMachine(name string) error {
+func (e podmanEngine) InitializeMachine(name string, s *spinner.Spinner) error {
 	// Grab some optional configurations from the config file.
 	podmanCmd := Command{
 		Command: podman,
@@ -40,9 +39,18 @@ func (e podmanEngine) InitializeMachine(name string) error {
 			"--now",
 		},
 	}
-	output, err := podmanCmd.Execute()
+	err := podmanCmd.ExecuteWithProgress(func(line string) {
+		switch {
+		case strings.Contains(line, "Looking up Podman Machine image"):
+		case strings.Contains(line, "Getting image source signatures"):
+		case strings.Contains(line, "Copying blob"):
+			s.Suffix = " Downloading Astro machine image…"
+		default:
+			s.Suffix = " Starting Astro machine…"
+		}
+	})
 	if err != nil {
-		return ErrorFromOutput("error initializing machine: %s", output)
+		return fmt.Errorf("error initializing machine: %w", err)
 	}
 	return nil
 }

airflow/runtimes/podman_runtime.go

@@ -8,7 +8,7 @@ import (
 	"time"
 
 	"github.com/astronomer/astro-cli/airflow/runtimes/types"
-
+	sp "github.com/astronomer/astro-cli/pkg/spinner"
 	"github.com/briandowns/spinner"
 )
 
@@ -18,7 +18,7 @@ const (
 )
 
 type PodmanEngine interface {
-	InitializeMachine(name string) error
+	InitializeMachine(name string, s *spinner.Spinner) error
 	StartMachine(name string) error
 	StopMachine(name string) error
 	RemoveMachine(name string) error
@@ -118,16 +118,9 @@ func (rt PodmanRuntime) Kill() error {
 
 func (rt PodmanRuntime) ensureMachine() error {
 	// Show a spinner message while we're initializing the machine.
-	s := spinner.New(spinnerCharSet, spinnerRefresh)
-	s.Suffix = containerRuntimeInitMessage
+	s := sp.NewSpinner(containerRuntimeInitMessage)
 	defer s.Stop()
 
-	// Update the message after a bit if it's still running.
-	go func() {
-		<-time.After(1 * time.Minute)
-		s.Suffix = podmanInitSlowMessage
-	}()
-
 	// Check if another, non-astro Podman machine is running
 	nonAstroMachineName := rt.isAnotherMachineRunning()
 	// If there is another machine running, and it has no running containers, stop it.
@@ -189,9 +182,12 @@ func (rt PodmanRuntime) ensureMachine() error {
 
 	// Otherwise, initialize the machine
 	s.Start()
-	if err := rt.Engine.InitializeMachine(podmanMachineName); err != nil {
+	// time delay of 1 second to display containerRuntimeInitMessage before initializing astro-machine
+	time.Sleep(1 * time.Second)
+	if err := rt.Engine.InitializeMachine(podmanMachineName, s); err != nil {
 		return err
 	}
+	sp.StopWithCheckmark(s, "Astro machine initialized")
 
 	return rt.getAndConfigureMachineForUsage(podmanMachineName)
 }

airflow/runtimes/podman_runtime_test.go

@@ -6,6 +6,8 @@ import (
 
 	"github.com/astronomer/astro-cli/airflow/runtimes/mocks"
 	"github.com/astronomer/astro-cli/airflow/runtimes/types"
+	"github.com/briandowns/spinner"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
 
 	"github.com/stretchr/testify/assert"
@@ -87,7 +89,10 @@ func (s *PodmanRuntimeSuite) TestPodmanRuntimeInitialize() {
 	s.Run("No machines running on mac, initialize podman", func() {
 		// Set up mocks.
 		mockPodmanEngine.On("ListMachines").Return(mockListedMachines, nil)
-		mockPodmanEngine.On("InitializeMachine", podmanMachineName).Return(nil)
+		mockPodmanEngine.On("InitializeMachine", podmanMachineName, mock.MatchedBy(func(s interface{}) bool {
+			_, ok := s.(*spinner.Spinner)
+			return ok
+		})).Return(nil)
 		mockPodmanEngine.On("InspectMachine", podmanMachineName).Return(mockInspectedAstroMachine, nil)
 		mockPodmanEngine.On("SetMachineAsDefault", podmanMachineName).Return(nil)
 		mockPodmanOSChecker.On("IsWindows").Return(false)
@@ -105,7 +110,10 @@ func (s *PodmanRuntimeSuite) TestPodmanRuntimeInitializeWindows() {
 	s.Run("No machines running on windows, initialize podman", func() {
 		// Set up mocks.
 		mockPodmanEngine.On("ListMachines").Return(mockListedMachines, nil)
-		mockPodmanEngine.On("InitializeMachine", podmanMachineName).Return(nil)
+		mockPodmanEngine.On("InitializeMachine", podmanMachineName, mock.MatchedBy(func(s interface{}) bool {
+			_, ok := s.(*spinner.Spinner)
+			return ok
+		})).Return(nil)
 		mockPodmanEngine.On("InspectMachine", podmanMachineName).Return(mockInspectedAstroMachine, nil)
 		mockPodmanEngine.On("SetMachineAsDefault", podmanMachineName).Return(nil)
 		mockPodmanOSChecker.On("IsWindows").Return(true)
@@ -134,7 +142,10 @@ func (s *PodmanRuntimeSuite) TestPodmanRuntimeInitializeWithAnotherMachineRunnin
 		mockPodmanEngine.On("ListContainers").Return(mockListedContainers, nil)
 		mockPodmanEngine.On("StopMachine", mockListedMachines[0].Name).Return(nil)
 		mockPodmanEngine.On("ListMachines").Return([]types.ListedMachine{}, nil).Once()
-		mockPodmanEngine.On("InitializeMachine", podmanMachineName).Return(nil)
+		mockPodmanEngine.On("InitializeMachine", podmanMachineName, mock.MatchedBy(func(s interface{}) bool {
+			_, ok := s.(*spinner.Spinner)
+			return ok
+		})).Return(nil)
 		mockPodmanEngine.On("InspectMachine", podmanMachineName).Return(mockInspectedAstroMachine, nil).Once()
 		mockPodmanEngine.On("SetMachineAsDefault", podmanMachineName).Return(nil).Once()
 		mockPodmanOSChecker.On("IsMac").Return(true)