diff --git a/airflow/runtimes/command.go b/airflow/runtimes/command.go index 458f9fb61..5dc03d6c1 100644 --- a/airflow/runtimes/command.go +++ b/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 +} diff --git a/airflow/runtimes/container_runtime.go b/airflow/runtimes/container_runtime.go index 7fe9d2483..91b41d351 100644 --- a/airflow/runtimes/container_runtime.go +++ b/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 diff --git a/airflow/runtimes/docker_runtime.go b/airflow/runtimes/docker_runtime.go index a59d5df2c..0a0419657 100644 --- a/airflow/runtimes/docker_runtime.go +++ b/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. diff --git a/airflow/runtimes/mocks/PodmanEngine.go b/airflow/runtimes/mocks/PodmanEngine.go index 61fc78b03..6c1f268b1 100644 --- a/airflow/runtimes/mocks/PodmanEngine.go +++ b/airflow/runtimes/mocks/PodmanEngine.go @@ -5,6 +5,8 @@ package mocks import ( mock "github.com/stretchr/testify/mock" + spinner "github.com/briandowns/spinner" + types "github.com/astronomer/astro-cli/airflow/runtimes/types" ) @@ -13,17 +15,17 @@ type PodmanEngine struct { mock.Mock } -// InitializeMachine provides a mock function with given fields: name -func (_m *PodmanEngine) InitializeMachine(name string) error { - ret := _m.Called(name) +// InitializeMachine provides a mock function with given fields: name, s +func (_m *PodmanEngine) InitializeMachine(name string, s *spinner.Spinner) error { + ret := _m.Called(name, s) if len(ret) == 0 { panic("no return value specified for InitializeMachine") } var r0 error - if rf, ok := ret.Get(0).(func(string) error); ok { - r0 = rf(name) + if rf, ok := ret.Get(0).(func(string, *spinner.Spinner) error); ok { + r0 = rf(name, s) } else { r0 = ret.Error(0) } diff --git a/airflow/runtimes/podman_engine.go b/airflow/runtimes/podman_engine.go index 5952d8649..0aa3da4a5 100644 --- a/airflow/runtimes/podman_engine.go +++ b/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 } diff --git a/airflow/runtimes/podman_runtime.go b/airflow/runtimes/podman_runtime.go index 95129474c..e18dc9fcf 100644 --- a/airflow/runtimes/podman_runtime.go +++ b/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) } diff --git a/airflow/runtimes/podman_runtime_test.go b/airflow/runtimes/podman_runtime_test.go index 1ce575cf9..448d774c5 100644 --- a/airflow/runtimes/podman_runtime_test.go +++ b/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)