From d42a29243270dabe42f70884fd51e6c7fbee87cb Mon Sep 17 00:00:00 2001 From: Greg Neiheisel <1036482+schnie@users.noreply.github.com> Date: Wed, 11 Dec 2024 10:40:50 -0500 Subject: [PATCH] Adds Podman Container Runtime (#1750) Co-authored-by: Pritesh Arora Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Neel Dalsania --- .mockery.yaml | 9 + airflow/docker_image.go | 3 + airflow/runtimes/command_test.go | 13 +- airflow/runtimes/container_runtime.go | 39 +- airflow/runtimes/container_runtime_test.go | 37 +- airflow/runtimes/docker.go | 84 ---- airflow/runtimes/docker_engine.go | 33 ++ airflow/runtimes/docker_runtime.go | 82 +++- airflow/runtimes/docker_runtime_test.go | 89 ++++ airflow/runtimes/docker_test.go | 77 ---- airflow/runtimes/file_checker.go | 25 ++ airflow/runtimes/mocks/ContainerRuntime.go | 80 ++++ airflow/runtimes/mocks/DockerEngine.go | 72 ++++ airflow/runtimes/mocks/FileChecker.go | 13 + airflow/runtimes/mocks/OSChecker.go | 52 +++ airflow/runtimes/mocks/PodmanEngine.go | 176 ++++++++ airflow/runtimes/mocks/RuntimeChecker.go | 14 + airflow/runtimes/os_checker.go | 26 ++ airflow/runtimes/podman_engine.go | 206 +++++++++ airflow/runtimes/podman_engine_test.go | 36 ++ airflow/runtimes/podman_runtime.go | 325 +++++++++++++- airflow/runtimes/podman_runtime_test.go | 466 +++++++++++++++++++++ airflow/runtimes/types/podman.go | 35 ++ airflow/runtimes/utils.go | 15 - cmd/airflow.go | 130 +++--- cmd/airflow_hooks.go | 50 ++- cmd/airflow_test.go | 40 +- config/config.go | 2 + config/types.go | 2 + 29 files changed, 1908 insertions(+), 323 deletions(-) delete mode 100644 airflow/runtimes/docker.go create mode 100644 airflow/runtimes/docker_engine.go create mode 100644 airflow/runtimes/docker_runtime_test.go delete mode 100644 airflow/runtimes/docker_test.go create mode 100644 airflow/runtimes/file_checker.go create mode 100644 airflow/runtimes/mocks/ContainerRuntime.go create mode 100644 airflow/runtimes/mocks/DockerEngine.go create mode 100644 airflow/runtimes/mocks/FileChecker.go create mode 100644 airflow/runtimes/mocks/OSChecker.go create mode 100644 airflow/runtimes/mocks/PodmanEngine.go create mode 100644 airflow/runtimes/mocks/RuntimeChecker.go create mode 100644 airflow/runtimes/os_checker.go create mode 100644 airflow/runtimes/podman_engine.go create mode 100644 airflow/runtimes/podman_engine_test.go create mode 100644 airflow/runtimes/podman_runtime_test.go create mode 100644 airflow/runtimes/types/podman.go delete mode 100644 airflow/runtimes/utils.go diff --git a/.mockery.yaml b/.mockery.yaml index 6e1c4ee37..3e4bef97d 100644 --- a/.mockery.yaml +++ b/.mockery.yaml @@ -51,3 +51,12 @@ packages: dir: pkg/azure/mocks interfaces: Azure: + github.com/astronomer/astro-cli/airflow/runtimes: + config: + dir: airflow/runtimes/mocks + outpkg: mocks + interfaces: + OSChecker: + ContainerRuntime: + PodmanEngine: + DockerEngine: diff --git a/airflow/docker_image.go b/airflow/docker_image.go index 7e1807043..922e88ea0 100644 --- a/airflow/docker_image.go +++ b/airflow/docker_image.go @@ -85,6 +85,9 @@ func (d *DockerImage) Build(dockerfilePath, buildSecretString string, buildConfi if err != nil { return fmt.Errorf("reading dockerfile: %w", err) } + if runtimes.IsPodman(containerRuntime) { + args = append(args, "--format", "docker") + } if addPullFlag { args = append(args, "--pull") } diff --git a/airflow/runtimes/command_test.go b/airflow/runtimes/command_test.go index 8a4daacdb..b852aece4 100644 --- a/airflow/runtimes/command_test.go +++ b/airflow/runtimes/command_test.go @@ -1,10 +1,21 @@ package runtimes import ( + "testing" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/suite" ) -func (s *ContainerRuntimeSuite) TestCommandExecution() { +type ContainerRuntimeCommandSuite struct { + suite.Suite +} + +func TestContainerRuntimeCommand(t *testing.T) { + suite.Run(t, new(ContainerRuntimeCommandSuite)) +} + +func (s *ContainerRuntimeCommandSuite) TestCommandExecution() { s.Run("Command executes successfully", func() { cmd := &Command{ Command: "echo", diff --git a/airflow/runtimes/container_runtime.go b/airflow/runtimes/container_runtime.go index 40e39e35e..38c2a6fe9 100644 --- a/airflow/runtimes/container_runtime.go +++ b/airflow/runtimes/container_runtime.go @@ -10,7 +10,6 @@ import ( "github.com/briandowns/spinner" "github.com/astronomer/astro-cli/config" - "github.com/astronomer/astro-cli/pkg/fileutil" "github.com/astronomer/astro-cli/pkg/util" "github.com/pkg/errors" ) @@ -32,6 +31,9 @@ const ( // the container runtime lifecycle. type ContainerRuntime interface { Initialize() error + Configure() error + ConfigureOrKill() error + Kill() error } // GetContainerRuntime creates a new container runtime based on the runtime string @@ -46,44 +48,26 @@ func GetContainerRuntime() (ContainerRuntime, error) { // Return the appropriate container runtime based on the binary discovered. switch containerRuntime { case docker: - return DockerRuntime{}, nil + return CreateDockerRuntime(new(dockerEngine), new(osChecker)), nil case podman: - return PodmanRuntime{}, nil + return CreatePodmanRuntime(new(podmanEngine), new(osChecker)), nil default: return nil, errors.New(containerRuntimeNotFoundErrMsg) } } -// FileChecker interface defines a method to check if a file exists. -// This is here mostly for testing purposes. This allows us to mock -// around actually checking for binaries on a live system as that -// would create inconsistencies across developer machines when -// working with the unit tests. -type FileChecker interface { - Exists(path string) bool -} - -// OSFileChecker is a concrete implementation of FileChecker. -type OSFileChecker struct{} - -// Exists checks if the file exists in the file system. -func (f OSFileChecker) Exists(path string) bool { - exists, _ := fileutil.Exists(path, nil) - return exists -} - // FindBinary searches for the specified binary name in the provided $PATH directories, // using the provided FileChecker. It searches each specific path within the systems // $PATH environment variable for the binary concurrently and returns a boolean result // indicating if the binary was found or not. -func FindBinary(pathEnv, binaryName string, checker FileChecker) bool { +func FindBinary(pathEnv, binaryName string, checker FileChecker, osChecker OSChecker) bool { // Split the $PATH variable into it's individual paths, // using the OS specific path separator character. paths := strings.Split(pathEnv, string(os.PathListSeparator)) // Although programs can be called without the .exe extension, // we need to append it here when searching the file system. - if isWindows() { + if osChecker.IsWindows() { binaryName += ".exe" } @@ -147,7 +131,7 @@ var GetContainerRuntimeBinary = func() (string, error) { // Get the $PATH environment variable. pathEnv := os.Getenv("PATH") for _, binary := range binaries { - if found := FindBinary(pathEnv, binary, OSFileChecker{}); found { + if found := FindBinary(pathEnv, binary, CreateFileChecker(), CreateOSChecker()); found { return binary, nil } } @@ -158,3 +142,10 @@ var GetContainerRuntimeBinary = func() (string, error) { "See the Astro CLI prerequisites for more information. " + "https://www.astronomer.io/docs/astro/cli/install-cli") } + +// IsPodman is just a small helper to avoid exporting the podman constant, +// and used in other places that haven't been refactored to use the runtime package. +// This could probably be removed in the future. +func IsPodman(binaryName string) bool { + return binaryName == podman +} diff --git a/airflow/runtimes/container_runtime_test.go b/airflow/runtimes/container_runtime_test.go index 0a3199c54..f16cf5fdc 100644 --- a/airflow/runtimes/container_runtime_test.go +++ b/airflow/runtimes/container_runtime_test.go @@ -4,9 +4,9 @@ import ( "errors" "testing" - "github.com/stretchr/testify/assert" + "github.com/astronomer/astro-cli/airflow/runtimes/mocks" - "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/suite" ) @@ -15,23 +15,13 @@ type ContainerRuntimeSuite struct { suite.Suite } -func TestConfig(t *testing.T) { +func TestContainerRuntime(t *testing.T) { suite.Run(t, new(ContainerRuntimeSuite)) } -// Mock for GetContainerRuntimeBinary -type MockRuntimeChecker struct { - mock.Mock -} - -func (m *MockRuntimeChecker) GetContainerRuntimeBinary() (string, error) { - args := m.Called() - return args.String(0), args.Error(1) -} - func (s *ContainerRuntimeSuite) TestGetContainerRuntime() { s.Run("GetContainerRuntime_Docker", func() { - mockChecker := new(MockRuntimeChecker) + mockChecker := new(mocks.RuntimeChecker) mockChecker.On("GetContainerRuntimeBinary").Return(docker, nil) // Inject the mock and make sure we restore after the test. @@ -47,7 +37,7 @@ func (s *ContainerRuntimeSuite) TestGetContainerRuntime() { }) s.Run("GetContainerRuntime_Podman", func() { - mockChecker := new(MockRuntimeChecker) + mockChecker := new(mocks.RuntimeChecker) mockChecker.On("GetContainerRuntimeBinary").Return(podman, nil) // Inject the mock and make sure we restore after the test. @@ -63,7 +53,7 @@ func (s *ContainerRuntimeSuite) TestGetContainerRuntime() { }) s.Run("GetContainerRuntime_Error", func() { - mockChecker := new(MockRuntimeChecker) + mockChecker := new(mocks.RuntimeChecker) mockChecker.On("GetContainerRuntimeBinary").Return("", errors.New(containerRuntimeNotFoundErrMsg)) // Inject the mock and make sure we restore after the test. @@ -80,17 +70,6 @@ func (s *ContainerRuntimeSuite) TestGetContainerRuntime() { }) } -// MockFileChecker is a mock implementation of FileChecker for tests. -type MockFileChecker struct { - existingFiles map[string]bool -} - -// Exists is just a mock for os.Stat(). In our test implementation, we just check -// if the file exists in the list of mocked files for a given test. -func (m MockFileChecker) Exists(path string) bool { - return m.existingFiles[path] -} - // TestGetContainerRuntimeBinary runs a suite of tests against GetContainerRuntimeBinary, // using the MockFileChecker defined above. func (s *ContainerRuntimeSuite) TestGetContainerRuntimeBinary() { @@ -163,8 +142,8 @@ func (s *ContainerRuntimeSuite) TestGetContainerRuntimeBinary() { for _, tt := range tests { s.Run(tt.name, func() { - mockChecker := MockFileChecker{existingFiles: tt.mockFiles} - result := FindBinary(tt.pathEnv, tt.binary, mockChecker) + mockChecker := mocks.FileChecker{ExistingFiles: tt.mockFiles} + result := FindBinary(tt.pathEnv, tt.binary, mockChecker, new(osChecker)) s.Equal(tt.expected, result) }) } diff --git a/airflow/runtimes/docker.go b/airflow/runtimes/docker.go deleted file mode 100644 index afafad9b6..000000000 --- a/airflow/runtimes/docker.go +++ /dev/null @@ -1,84 +0,0 @@ -package runtimes - -import ( - "fmt" - "time" - - "github.com/briandowns/spinner" -) - -const ( - defaultTimeoutSeconds = 60 - tickNum = 500 - open = "open" - timeoutErrMsg = "timed out waiting for docker" - dockerOpenNotice = "We couldn't start the docker engine automatically. Please start it manually and try again." -) - -// DockerInitializer is a struct that contains the functions needed to initialize Docker. -// The concrete implementation that we use is DefaultDockerInitializer below. -// When running the tests, we substitute the default implementation with a mock implementation. -type DockerInitializer interface { - CheckDockerCmd() (string, error) - OpenDockerCmd() (string, error) -} - -// DefaultDockerInitializer is the default implementation of DockerInitializer. -// The concrete functions defined here are called from the InitializeDocker function below. -type DefaultDockerInitializer struct{} - -func (d DefaultDockerInitializer) CheckDockerCmd() (string, error) { - checkDockerCmd := Command{ - Command: docker, - Args: []string{"ps"}, - } - return checkDockerCmd.Execute() -} - -func (d DefaultDockerInitializer) OpenDockerCmd() (string, error) { - openDockerCmd := Command{ - Command: open, - Args: []string{"-a", docker}, - } - return openDockerCmd.Execute() -} - -// InitializeDocker initializes the Docker runtime. -// It checks if Docker is running, and if it is not, it attempts to start it. -func InitializeDocker(d DockerInitializer, timeoutSeconds int) error { - // Initialize spinner. - timeout := time.After(time.Duration(timeoutSeconds) * time.Second) - ticker := time.NewTicker(time.Duration(tickNum) * time.Millisecond) - s := spinner.New(spinnerCharSet, spinnerRefresh) - s.Suffix = containerRuntimeInitMessage - defer s.Stop() - - // Execute `docker ps` to check if Docker is running. - _, err := d.CheckDockerCmd() - - // If we didn't get an error, Docker is running, so we can return. - if err == nil { - return nil - } - - // If we got an error, Docker is not running, so we attempt to start it. - _, err = d.OpenDockerCmd() - if err != nil { - return fmt.Errorf(dockerOpenNotice) //nolint:stylecheck - } - - // Wait for Docker to start. - s.Start() - for { - select { - case <-timeout: - return fmt.Errorf(timeoutErrMsg) - case <-ticker.C: - _, err := d.CheckDockerCmd() - if err != nil { - continue - } - return nil - } - } -} diff --git a/airflow/runtimes/docker_engine.go b/airflow/runtimes/docker_engine.go new file mode 100644 index 000000000..8900ddcc2 --- /dev/null +++ b/airflow/runtimes/docker_engine.go @@ -0,0 +1,33 @@ +package runtimes + +const ( + defaultTimeoutSeconds = 60 + tickNum = 500 + open = "open" + timeoutErrMsg = "timed out waiting for docker" + dockerOpenNotice = "We couldn't start the docker engine automatically. Please start it manually and try again." +) + +// dockerEngine is the default implementation of DockerEngine. +type dockerEngine struct{} + +func (d dockerEngine) IsRunning() (string, error) { + checkDockerCmd := Command{ + Command: docker, + Args: []string{ + "ps", + }, + } + return checkDockerCmd.Execute() +} + +func (d dockerEngine) Start() (string, error) { + openDockerCmd := Command{ + Command: open, + Args: []string{ + "-a", + docker, + }, + } + return openDockerCmd.Execute() +} diff --git a/airflow/runtimes/docker_runtime.go b/airflow/runtimes/docker_runtime.go index 4e45d6fb1..e82853f89 100644 --- a/airflow/runtimes/docker_runtime.go +++ b/airflow/runtimes/docker_runtime.go @@ -1,14 +1,88 @@ package runtimes +import ( + "fmt" + "time" + + "github.com/briandowns/spinner" +) + +// DockerEngine is a struct that contains the functions needed to initialize Docker. +// The concrete implementation that we use is dockerEngine. +// When running the tests, we substitute the default implementation with a mock implementation. +type DockerEngine interface { + IsRunning() (string, error) + Start() (string, error) +} + // DockerRuntime is a concrete implementation of the ContainerRuntime interface. // When the docker binary is chosen, this implementation is used. -type DockerRuntime struct{} +type DockerRuntime struct { + Engine DockerEngine + OSChecker OSChecker +} + +func CreateDockerRuntime(engine DockerEngine, osChecker OSChecker) DockerRuntime { + return DockerRuntime{Engine: engine, OSChecker: osChecker} +} // Initialize initializes the Docker runtime. // We only attempt to initialize Docker on Mac today. -func (p DockerRuntime) Initialize() error { - if !isMac() { +func (rt DockerRuntime) Initialize() error { + if !rt.OSChecker.IsMac() { return nil } - return InitializeDocker(new(DefaultDockerInitializer), defaultTimeoutSeconds) + return rt.initializeDocker(defaultTimeoutSeconds) +} + +func (rt DockerRuntime) Configure() error { + return nil +} + +func (rt DockerRuntime) ConfigureOrKill() error { + return nil +} + +func (rt DockerRuntime) Kill() error { + return nil +} + +// initializeDocker initializes the Docker runtime. +// It checks if Docker is running, and if it is not, it attempts to start it. +func (rt DockerRuntime) initializeDocker(timeoutSeconds int) error { + // Initialize spinner. + timeout := time.After(time.Duration(timeoutSeconds) * time.Second) + ticker := time.NewTicker(time.Duration(tickNum) * time.Millisecond) + s := spinner.New(spinnerCharSet, spinnerRefresh) + s.Suffix = containerRuntimeInitMessage + defer s.Stop() + + // Execute `docker ps` to check if Docker is running. + _, err := rt.Engine.IsRunning() + + // If we didn't get an error, Docker is running, so we can return. + if err == nil { + return nil + } + + // If we got an error, Docker is not running, so we attempt to start it. + _, err = rt.Engine.Start() + if err != nil { + return fmt.Errorf(dockerOpenNotice) //nolint:stylecheck + } + + // Wait for Docker to start. + s.Start() + for { + select { + case <-timeout: + return fmt.Errorf(timeoutErrMsg) + case <-ticker.C: + _, err := rt.Engine.IsRunning() + if err != nil { + continue + } + return nil + } + } } diff --git a/airflow/runtimes/docker_runtime_test.go b/airflow/runtimes/docker_runtime_test.go new file mode 100644 index 000000000..471c0a538 --- /dev/null +++ b/airflow/runtimes/docker_runtime_test.go @@ -0,0 +1,89 @@ +package runtimes + +import ( + "fmt" + "testing" + + "github.com/astronomer/astro-cli/airflow/runtimes/mocks" + "github.com/stretchr/testify/suite" + + "github.com/stretchr/testify/assert" +) + +var ( + mockDockerEngine *mocks.DockerEngine + mockDockerOSChecker *mocks.OSChecker +) + +type DockerRuntimeSuite struct { + suite.Suite +} + +func TestDockerRuntime(t *testing.T) { + suite.Run(t, new(DockerRuntimeSuite)) +} + +func (s *DockerRuntimeSuite) SetupTest() { + // Reset some variables to defaults. + mockDockerEngine = new(mocks.DockerEngine) + mockDockerOSChecker = new(mocks.OSChecker) +} + +func (s *DockerRuntimeSuite) TestStartDocker() { + s.Run("Docker is running, returns nil", func() { + // Simulate that the initial `docker ps` succeeds and we exit early. + mockDockerEngine.On("IsRunning").Return("", nil).Once() + mockDockerOSChecker.On("IsMac").Return(true) + // Create the runtime with our mock engine. + rt := CreateDockerRuntime(mockDockerEngine, mockDockerOSChecker) + // Run our test and assert expectations. + err := rt.Initialize() + assert.Nil(s.T(), err, "Expected no error when docker is running") + mockDockerEngine.AssertExpectations(s.T()) + }) + + s.Run("Docker is not running, tries to start and waits", func() { + // Simulate that the initial `docker ps` fails. + mockDockerEngine.On("IsRunning").Return("", fmt.Errorf("docker not running")).Once() + // Simulate that `open -a docker` succeeds. + mockDockerEngine.On("Start").Return("", nil).Once() + // Simulate that `docker ps` works after trying to open docker. + mockDockerEngine.On("IsRunning").Return("", nil).Once() + mockDockerOSChecker.On("IsMac").Return(true) + // Create the runtime with our mock engine. + rt := CreateDockerRuntime(mockDockerEngine, mockDockerOSChecker) + // Run our test and assert expectations. + err := rt.Initialize() + assert.Nil(s.T(), err, "Expected no error when docker starts after retry") + mockDockerEngine.AssertExpectations(s.T()) + }) + + s.Run("Docker fails to open", func() { + // Simulate `docker ps` failing. + mockDockerEngine.On("IsRunning").Return("", fmt.Errorf("docker not running")).Once() + // Simulate `open -a docker` failing. + mockDockerEngine.On("Start").Return("", fmt.Errorf("failed to open docker")).Once() + mockDockerOSChecker.On("IsMac").Return(true) + // Create the runtime with our mock engine. + rt := CreateDockerRuntime(mockDockerEngine, mockDockerOSChecker) + // Run our test and assert expectations. + err := rt.Initialize() + assert.Equal(s.T(), fmt.Errorf(dockerOpenNotice), err, "Expected timeout error") + mockDockerEngine.AssertExpectations(s.T()) + }) + + s.Run("Docker open succeeds but check times out", func() { + // Simulate `docker ps` failing continuously. + mockDockerEngine.On("IsRunning").Return("", fmt.Errorf("docker not running")) + // Simulate `open -a docker` failing. + mockDockerEngine.On("Start").Return("", nil).Once() + // Create the runtime with our mock engine. + rt := CreateDockerRuntime(mockDockerEngine, mockDockerOSChecker) + // Run our test and assert expectations. + // Call the helper method directly with custom timeout. + // Simulate the timeout after 1 second. + err := rt.initializeDocker(1) + assert.Equal(s.T(), fmt.Errorf(timeoutErrMsg), err, "Expected timeout error") + mockDockerEngine.AssertExpectations(s.T()) + }) +} diff --git a/airflow/runtimes/docker_test.go b/airflow/runtimes/docker_test.go deleted file mode 100644 index e6d018f02..000000000 --- a/airflow/runtimes/docker_test.go +++ /dev/null @@ -1,77 +0,0 @@ -package runtimes - -import ( - "fmt" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/mock" -) - -type MockDockerInitializer struct { - mock.Mock -} - -func (d *MockDockerInitializer) CheckDockerCmd() (string, error) { - args := d.Called() - return args.String(0), args.Error(1) -} - -func (d *MockDockerInitializer) OpenDockerCmd() (string, error) { - args := d.Called() - return args.String(0), args.Error(1) -} - -func (s *ContainerRuntimeSuite) TestStartDocker() { - s.Run("Docker is running, returns nil", func() { - // Create mock initializer. - mockInitializer := new(MockDockerInitializer) - // Simulate that the initial `docker ps` succeeds and we exit early. - mockInitializer.On("CheckDockerCmd").Return("", nil).Once() - // Run our test and assert expectations. - err := InitializeDocker(mockInitializer, defaultTimeoutSeconds) - assert.Nil(s.T(), err, "Expected no error when docker is running") - mockInitializer.AssertExpectations(s.T()) - }) - - s.Run("Docker is not running, tries to start and waits", func() { - // Create mock initializer. - mockInitializer := new(MockDockerInitializer) - // Simulate that the initial `docker ps` fails. - mockInitializer.On("CheckDockerCmd").Return("", fmt.Errorf("docker not running")).Once() - // Simulate that `open -a docker` succeeds. - mockInitializer.On("OpenDockerCmd").Return("", nil).Once() - // Simulate that `docker ps` works after trying to open docker. - mockInitializer.On("CheckDockerCmd").Return("", nil).Once() - // Run our test and assert expectations. - err := InitializeDocker(mockInitializer, defaultTimeoutSeconds) - assert.Nil(s.T(), err, "Expected no error when docker starts after retry") - mockInitializer.AssertExpectations(s.T()) - }) - - s.Run("Docker fails to open", func() { - // Create mock initializer. - mockInitializer := new(MockDockerInitializer) - // Simulate `docker ps` failing. - mockInitializer.On("CheckDockerCmd").Return("", fmt.Errorf("docker not running")).Once() - // Simulate `open -a docker` failing. - mockInitializer.On("OpenDockerCmd").Return("", fmt.Errorf("failed to open docker")).Once() - // Run our test and assert expectations. - err := InitializeDocker(mockInitializer, defaultTimeoutSeconds) - assert.Equal(s.T(), fmt.Errorf(dockerOpenNotice), err, "Expected timeout error") - mockInitializer.AssertExpectations(s.T()) - }) - - s.Run("Docker open succeeds but check times out", func() { - // Create mock initializer. - mockInitializer := new(MockDockerInitializer) - // Simulate `docker ps` failing continuously. - mockInitializer.On("CheckDockerCmd").Return("", fmt.Errorf("docker not running")) - // Simulate `open -a docker` failing. - mockInitializer.On("OpenDockerCmd").Return("", nil).Once() - // Run our test and assert expectations. - // Simulate the timeout after 1 second. - err := InitializeDocker(mockInitializer, 1) - assert.Equal(s.T(), fmt.Errorf(timeoutErrMsg), err, "Expected timeout error") - mockInitializer.AssertExpectations(s.T()) - }) -} diff --git a/airflow/runtimes/file_checker.go b/airflow/runtimes/file_checker.go new file mode 100644 index 000000000..66cefcf30 --- /dev/null +++ b/airflow/runtimes/file_checker.go @@ -0,0 +1,25 @@ +package runtimes + +import "github.com/astronomer/astro-cli/pkg/fileutil" + +// FileChecker interface defines a method to check if a file exists. +// This is here mostly for testing purposes. This allows us to mock +// around actually checking for binaries on a live system as that +// would create inconsistencies across developer machines when +// working with the unit tests. +type FileChecker interface { + Exists(path string) bool +} + +// fileChecker is a concrete implementation of FileChecker. +type fileChecker struct{} + +func CreateFileChecker() FileChecker { + return new(fileChecker) +} + +// Exists checks if the file exists in the file system. +func (f fileChecker) Exists(path string) bool { + exists, _ := fileutil.Exists(path, nil) + return exists +} diff --git a/airflow/runtimes/mocks/ContainerRuntime.go b/airflow/runtimes/mocks/ContainerRuntime.go new file mode 100644 index 000000000..c6453409a --- /dev/null +++ b/airflow/runtimes/mocks/ContainerRuntime.go @@ -0,0 +1,80 @@ +// Code generated by mockery v2.32.0. DO NOT EDIT. + +package mocks + +import mock "github.com/stretchr/testify/mock" + +// ContainerRuntime is an autogenerated mock type for the ContainerRuntime type +type ContainerRuntime struct { + mock.Mock +} + +// Configure provides a mock function with given fields: +func (_m *ContainerRuntime) Configure() error { + ret := _m.Called() + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// ConfigureOrKill provides a mock function with given fields: +func (_m *ContainerRuntime) ConfigureOrKill() error { + ret := _m.Called() + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Initialize provides a mock function with given fields: +func (_m *ContainerRuntime) Initialize() error { + ret := _m.Called() + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// Kill provides a mock function with given fields: +func (_m *ContainerRuntime) Kill() error { + ret := _m.Called() + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewContainerRuntime creates a new instance of ContainerRuntime. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewContainerRuntime(t interface { + mock.TestingT + Cleanup(func()) +}) *ContainerRuntime { + mock := &ContainerRuntime{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/airflow/runtimes/mocks/DockerEngine.go b/airflow/runtimes/mocks/DockerEngine.go new file mode 100644 index 000000000..491eeaee9 --- /dev/null +++ b/airflow/runtimes/mocks/DockerEngine.go @@ -0,0 +1,72 @@ +// Code generated by mockery v2.32.0. DO NOT EDIT. + +package mocks + +import mock "github.com/stretchr/testify/mock" + +// DockerEngine is an autogenerated mock type for the DockerEngine type +type DockerEngine struct { + mock.Mock +} + +// IsRunning provides a mock function with given fields: +func (_m *DockerEngine) IsRunning() (string, error) { + ret := _m.Called() + + var r0 string + var r1 error + if rf, ok := ret.Get(0).(func() (string, error)); ok { + return rf() + } + if rf, ok := ret.Get(0).(func() string); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(string) + } + + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// Start provides a mock function with given fields: +func (_m *DockerEngine) Start() (string, error) { + ret := _m.Called() + + var r0 string + var r1 error + if rf, ok := ret.Get(0).(func() (string, error)); ok { + return rf() + } + if rf, ok := ret.Get(0).(func() string); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(string) + } + + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// NewDockerEngine creates a new instance of DockerEngine. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewDockerEngine(t interface { + mock.TestingT + Cleanup(func()) +}) *DockerEngine { + mock := &DockerEngine{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/airflow/runtimes/mocks/FileChecker.go b/airflow/runtimes/mocks/FileChecker.go new file mode 100644 index 000000000..662115cef --- /dev/null +++ b/airflow/runtimes/mocks/FileChecker.go @@ -0,0 +1,13 @@ +package mocks + +// FileChecker is a mock implementation of FileChecker for tests. +// This is a manually created mock, not generated by mockery. +type FileChecker struct { + ExistingFiles map[string]bool +} + +// Exists is just a mock for os.Stat(). In our test implementation, we just check +// if the file exists in the list of mocked files for a given test. +func (m FileChecker) Exists(path string) bool { + return m.ExistingFiles[path] +} diff --git a/airflow/runtimes/mocks/OSChecker.go b/airflow/runtimes/mocks/OSChecker.go new file mode 100644 index 000000000..1de522b5a --- /dev/null +++ b/airflow/runtimes/mocks/OSChecker.go @@ -0,0 +1,52 @@ +// Code generated by mockery v2.32.0. DO NOT EDIT. + +package mocks + +import mock "github.com/stretchr/testify/mock" + +// OSChecker is an autogenerated mock type for the OSChecker type +type OSChecker struct { + mock.Mock +} + +// IsMac provides a mock function with given fields: +func (_m *OSChecker) IsMac() bool { + ret := _m.Called() + + var r0 bool + if rf, ok := ret.Get(0).(func() bool); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(bool) + } + + return r0 +} + +// IsWindows provides a mock function with given fields: +func (_m *OSChecker) IsWindows() bool { + ret := _m.Called() + + var r0 bool + if rf, ok := ret.Get(0).(func() bool); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(bool) + } + + return r0 +} + +// NewOSChecker creates a new instance of OSChecker. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewOSChecker(t interface { + mock.TestingT + Cleanup(func()) +}) *OSChecker { + mock := &OSChecker{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/airflow/runtimes/mocks/PodmanEngine.go b/airflow/runtimes/mocks/PodmanEngine.go new file mode 100644 index 000000000..bfadbc701 --- /dev/null +++ b/airflow/runtimes/mocks/PodmanEngine.go @@ -0,0 +1,176 @@ +// Code generated by mockery v2.32.0. DO NOT EDIT. + +package mocks + +import ( + mock "github.com/stretchr/testify/mock" + + types "github.com/astronomer/astro-cli/airflow/runtimes/types" +) + +// PodmanEngine is an autogenerated mock type for the PodmanEngine type +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) + + var r0 error + if rf, ok := ret.Get(0).(func(string) error); ok { + r0 = rf(name) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// InspectMachine provides a mock function with given fields: name +func (_m *PodmanEngine) InspectMachine(name string) (*types.InspectedMachine, error) { + ret := _m.Called(name) + + var r0 *types.InspectedMachine + var r1 error + if rf, ok := ret.Get(0).(func(string) (*types.InspectedMachine, error)); ok { + return rf(name) + } + if rf, ok := ret.Get(0).(func(string) *types.InspectedMachine); ok { + r0 = rf(name) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*types.InspectedMachine) + } + } + + if rf, ok := ret.Get(1).(func(string) error); ok { + r1 = rf(name) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// ListContainers provides a mock function with given fields: +func (_m *PodmanEngine) ListContainers() ([]types.ListedContainer, error) { + ret := _m.Called() + + var r0 []types.ListedContainer + var r1 error + if rf, ok := ret.Get(0).(func() ([]types.ListedContainer, error)); ok { + return rf() + } + if rf, ok := ret.Get(0).(func() []types.ListedContainer); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]types.ListedContainer) + } + } + + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// ListMachines provides a mock function with given fields: +func (_m *PodmanEngine) ListMachines() ([]types.ListedMachine, error) { + ret := _m.Called() + + var r0 []types.ListedMachine + var r1 error + if rf, ok := ret.Get(0).(func() ([]types.ListedMachine, error)); ok { + return rf() + } + if rf, ok := ret.Get(0).(func() []types.ListedMachine); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]types.ListedMachine) + } + } + + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// RemoveMachine provides a mock function with given fields: name +func (_m *PodmanEngine) RemoveMachine(name string) error { + ret := _m.Called(name) + + var r0 error + if rf, ok := ret.Get(0).(func(string) error); ok { + r0 = rf(name) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// SetMachineAsDefault provides a mock function with given fields: name +func (_m *PodmanEngine) SetMachineAsDefault(name string) error { + ret := _m.Called(name) + + var r0 error + if rf, ok := ret.Get(0).(func(string) error); ok { + r0 = rf(name) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// StartMachine provides a mock function with given fields: name +func (_m *PodmanEngine) StartMachine(name string) error { + ret := _m.Called(name) + + var r0 error + if rf, ok := ret.Get(0).(func(string) error); ok { + r0 = rf(name) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// StopMachine provides a mock function with given fields: name +func (_m *PodmanEngine) StopMachine(name string) error { + ret := _m.Called(name) + + var r0 error + if rf, ok := ret.Get(0).(func(string) error); ok { + r0 = rf(name) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// NewPodmanEngine creates a new instance of PodmanEngine. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewPodmanEngine(t interface { + mock.TestingT + Cleanup(func()) +}) *PodmanEngine { + mock := &PodmanEngine{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/airflow/runtimes/mocks/RuntimeChecker.go b/airflow/runtimes/mocks/RuntimeChecker.go new file mode 100644 index 000000000..517de8a3b --- /dev/null +++ b/airflow/runtimes/mocks/RuntimeChecker.go @@ -0,0 +1,14 @@ +package mocks + +import "github.com/stretchr/testify/mock" + +// RuntimeChecker is a mock for GetContainerRuntimeBinary +// This is a manually created mock, not generated by mockery. +type RuntimeChecker struct { + mock.Mock +} + +func (m *RuntimeChecker) GetContainerRuntimeBinary() (string, error) { + args := m.Called() + return args.String(0), args.Error(1) +} diff --git a/airflow/runtimes/os_checker.go b/airflow/runtimes/os_checker.go new file mode 100644 index 000000000..2874c228a --- /dev/null +++ b/airflow/runtimes/os_checker.go @@ -0,0 +1,26 @@ +package runtimes + +import "runtime" + +type OSChecker interface { + IsMac() bool + IsWindows() bool +} + +type osChecker struct{} + +func CreateOSChecker() OSChecker { + return new(osChecker) +} + +// IsWindows is a utility function to determine if the CLI host machine +// is running on Microsoft Windows OS. +func (o osChecker) IsWindows() bool { + return runtime.GOOS == "windows" +} + +// IsMac is a utility function to determine if the CLI host machine +// is running on Apple macOS. +func (o osChecker) IsMac() bool { + return runtime.GOOS == "darwin" +} diff --git a/airflow/runtimes/podman_engine.go b/airflow/runtimes/podman_engine.go new file mode 100644 index 000000000..5952d8649 --- /dev/null +++ b/airflow/runtimes/podman_engine.go @@ -0,0 +1,206 @@ +package runtimes + +import ( + "encoding/json" + "fmt" + "strings" + + "github.com/astronomer/astro-cli/airflow/runtimes/types" + "github.com/astronomer/astro-cli/config" +) + +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." + 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. " + + "Please stop the other machine and try again" +) + +// podmanEngine is the default implementation of PodmanEngine. +type podmanEngine struct{} + +// InitializeMachine initializes our astro Podman machine. +func (e podmanEngine) InitializeMachine(name string) error { + // Grab some optional configurations from the config file. + podmanCmd := Command{ + Command: podman, + Args: []string{ + "machine", + "init", + name, + "--memory", + config.CFG.MachineMemory.GetString(), + "--cpus", + config.CFG.MachineCPU.GetString(), + "--now", + }, + } + output, err := podmanCmd.Execute() + if err != nil { + return ErrorFromOutput("error initializing machine: %s", output) + } + return nil +} + +// StartMachine starts our astro Podman machine. +func (e podmanEngine) StartMachine(name string) error { + podmanCmd := Command{ + Command: podman, + Args: []string{ + "machine", + "start", + name, + }, + } + output, err := podmanCmd.Execute() + if err != nil { + return ErrorFromOutput("error starting machine: %s", output) + } + return nil +} + +// StopMachine stops the given Podman machine. +func (e podmanEngine) StopMachine(name string) error { + podmanCmd := Command{ + Command: podman, + Args: []string{ + "machine", + "stop", + name, + }, + } + output, err := podmanCmd.Execute() + if err != nil { + return ErrorFromOutput("error stopping machine: %s", output) + } + return nil +} + +// RemoveMachine removes the given Podman machine completely, +// such that it can only be started again by re-initializing. +func (e podmanEngine) RemoveMachine(name string) error { + podmanCmd := Command{ + Command: podman, + Args: []string{ + "machine", + "rm", + "-f", + name, + }, + } + output, err := podmanCmd.Execute() + if err != nil { + return ErrorFromOutput("error removing machine: %s", output) + } + return nil +} + +// InspectMachine inspects a given podman machine name. +func (e podmanEngine) InspectMachine(name string) (*types.InspectedMachine, error) { + podmanCmd := Command{ + Command: podman, + Args: []string{ + "machine", + "inspect", + name, + }, + } + output, err := podmanCmd.Execute() + if err != nil { + return nil, ErrorFromOutput("error inspecting machine: %s", output) + } + + var machines []types.InspectedMachine + err = json.Unmarshal([]byte(output), &machines) + if err != nil { + return nil, err + } + if len(machines) == 0 { + return nil, fmt.Errorf("machine not found: %s", name) + } + + return &machines[0], nil +} + +// SetMachineAsDefault sets the given Podman machine as the default. +func (e podmanEngine) SetMachineAsDefault(name string) error { + podmanCmd := Command{ + Command: podman, + Args: []string{ + "system", + "connection", + "default", + name, + }, + } + output, err := podmanCmd.Execute() + if err != nil { + return ErrorFromOutput("error setting default connection: %s", output) + } + return nil +} + +// ListMachines lists all Podman machines. +func (e podmanEngine) ListMachines() ([]types.ListedMachine, error) { + podmanCmd := Command{ + Command: podman, + Args: []string{ + "machine", + "ls", + "--format", + "json", + }, + } + output, err := podmanCmd.Execute() + if err != nil { + return nil, ErrorFromOutput("error listing machines: %s", output) + } + var machines []types.ListedMachine + err = json.Unmarshal([]byte(output), &machines) + if err != nil { + return nil, err + } + return machines, nil +} + +// ListContainers lists all pods in the machine. +func (e podmanEngine) ListContainers() ([]types.ListedContainer, error) { + podmanCmd := Command{ + Command: podman, + Args: []string{ + "ps", + "--format", + "json", + }, + } + output, err := podmanCmd.Execute() + if err != nil { + return nil, ErrorFromOutput("error listing containers: %s", output) + } + var containers []types.ListedContainer + err = json.Unmarshal([]byte(output), &containers) + if err != nil { + return nil, err + } + return containers, nil +} + +// ErrorFromOutput returns an error from the output string. +// This is used to extract the meaningful error message from podman command output. +func ErrorFromOutput(prefix, output string) error { + lines := strings.Split(output, "\n") + for _, line := range lines { + if strings.HasPrefix(line, "Error: ") { + errMsg := strings.Trim(strings.TrimSpace(line), "Error: ") //nolint + return fmt.Errorf(prefix, errMsg) + } + } + // If we didn't find an error message, return the entire output + errMsg := strings.TrimSpace(output) + return fmt.Errorf(prefix, errMsg) +} diff --git a/airflow/runtimes/podman_engine_test.go b/airflow/runtimes/podman_engine_test.go new file mode 100644 index 000000000..be3a2dd6f --- /dev/null +++ b/airflow/runtimes/podman_engine_test.go @@ -0,0 +1,36 @@ +package runtimes + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/suite" +) + +type PodmanEngineSuite struct { + suite.Suite +} + +func TestPodmanEngine(t *testing.T) { + suite.Run(t, new(PodmanEngineSuite)) +} + +func (s *PodmanRuntimeSuite) TestPodmanEngineErrorFromOutput() { + s.Run("returns formatted error when error line is present", func() { + output := "Some output\nError: something went wrong\nMore output" + err := ErrorFromOutput("prefix: %s", output) + assert.EqualError(s.T(), err, "prefix: something went wrong") + }) + + s.Run("returns formatted error when output is empty", func() { + output := "" + err := ErrorFromOutput("prefix: %s", output) + assert.EqualError(s.T(), err, "prefix: ") + }) + + s.Run("returns formatted error when output contains only error line", func() { + output := "Error: something went wrong" + err := ErrorFromOutput("prefix: %s", output) + assert.EqualError(s.T(), err, "prefix: something went wrong") + }) +} diff --git a/airflow/runtimes/podman_runtime.go b/airflow/runtimes/podman_runtime.go index a4547de1b..ae33e5753 100644 --- a/airflow/runtimes/podman_runtime.go +++ b/airflow/runtimes/podman_runtime.go @@ -1,10 +1,325 @@ package runtimes -// PodmanRuntime is a concrete implementation of the ContainerRuntime interface. -// When the podman binary is chosen, this implementation is used. -type PodmanRuntime struct{} +import ( + "errors" + "fmt" + "os" + "strings" + "time" -// Initialize initializes the Podman runtime. -func (p PodmanRuntime) Initialize() error { + "github.com/astronomer/astro-cli/airflow/runtimes/types" + + "github.com/briandowns/spinner" +) + +const ( + podmanMachineName = "astro-machine" + projectNotRunningErrMsg = "this astro project is not running" +) + +type PodmanEngine interface { + InitializeMachine(name string) error + StartMachine(name string) error + StopMachine(name string) error + RemoveMachine(name string) error + InspectMachine(name string) (*types.InspectedMachine, error) + SetMachineAsDefault(name string) error + ListMachines() ([]types.ListedMachine, error) + ListContainers() ([]types.ListedContainer, error) +} + +type PodmanRuntime struct { + Engine PodmanEngine + OSChecker OSChecker +} + +// CreatePodmanRuntime creates a new PodmanRuntime using the provided PodmanEngine. +// The engine allows us to interact with the external podman environment. For unit testing, +// we provide a mock engine that can be used to simulate the podman environment. +func CreatePodmanRuntime(engine PodmanEngine, osChecker OSChecker) PodmanRuntime { + return PodmanRuntime{Engine: engine, OSChecker: osChecker} +} + +func (rt PodmanRuntime) Initialize() error { + // If we're in podman mode, and DOCKER_HOST is not already set + // we need to initialize our astro machine. + // If DOCKER_HOST is already set, we assume the user already has a + // workflow with podman that we don't want to interfere with. + if isDockerHostSet() { + return nil + } + return rt.ensureMachine() +} + +func (rt PodmanRuntime) Configure() error { + // If we're in podman mode, and DOCKER_HOST is not already set + // we need to set things up for our astro machine. + // If DOCKER_HOST is already set, we assume the user already has a + // workflow with podman that we don't want to interfere with. + if isDockerHostSet() { + return nil + } + + // If the astro machine is running, we just configure it + // for usage, so the regular compose commands can carry out. + if rt.astroMachineIsRunning() { + return rt.getAndConfigureMachineForUsage(podmanMachineName) + } + + // Otherwise, we return an error indicating that the project isn't running. + return fmt.Errorf(projectNotRunningErrMsg) +} + +func (rt PodmanRuntime) ConfigureOrKill() error { + // If we're in podman mode, and DOCKER_HOST is not already set + // we need to set things up for our astro machine. + // If DOCKER_HOST is already set, we assume the user already has a + // workflow with podman that we don't want to interfere with. + if isDockerHostSet() { + return nil + } + + // If the astro machine is running, we just configure it + // for usage, so the regular compose kill can carry out. + // We follow up with a machine kill in the post run hook. + if rt.astroMachineIsRunning() { + return rt.getAndConfigureMachineForUsage(podmanMachineName) + } + + // The machine is already not running, + // so we can just ensure its fully killed. + if err := rt.stopAndKillMachine(); err != nil { + return err + } + + // We also return an error indicating that you can't kill + // a project that isn't running. + return fmt.Errorf(projectNotRunningErrMsg) +} + +func (rt PodmanRuntime) Kill() error { + // If we're in podman mode, and DOCKER_HOST is set to the astro machine (in the pre-run hook), + // we'll ensure that the machine is killed. + if isDockerHostSetToAstroMachine() { + return rt.stopAndKillMachine() + } + return nil +} + +func (rt PodmanRuntime) ensureMachine() error { + // Show a spinner message while we're initializing the machine. + s := spinner.New(spinnerCharSet, spinnerRefresh) + s.Suffix = 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. + // Otherwise, we assume the user has some other project running that we don't want to interfere with. + if nonAstroMachineName != "" && rt.OSChecker.IsMac() { + // First, configure the other running machine for usage, so we can check it for containers. + if err := rt.getAndConfigureMachineForUsage(nonAstroMachineName); err != nil { + return err + } + + // Then check the machine for running containers. + containers, err := rt.Engine.ListContainers() + if err != nil { + return err + } + + // There's some other containers running on this machine, so we don't want to stop it. + // We want the user to stop it manually and restart astro. + if len(containers) > 0 { + return errors.New(podmanMachineAlreadyRunningErrMsg) + } + + // If we made it here, we're going to stop the other machine + // and start our own machine, so start the spinner and begin the process. + s.Start() + err = rt.Engine.StopMachine(nonAstroMachineName) + if err != nil { + return err + } + } + + // Check if our astro Podman machine exists. + machine := rt.getAstroMachine() + + // If the machine exists, inspect it and decide what to do. + if machine != nil { + // Inspect the machine and get its details. + iMachine, err := rt.Engine.InspectMachine(podmanMachineName) + if err != nil { + return err + } + + // If the machine is already running, + // just go ahead and configure it for usage. + if iMachine.State == podmanStatusRunning { + return rt.configureMachineForUsage(iMachine) + } + + // If the machine is stopped, + // start it, then configure it for usage. + if iMachine.State == podmanStatusStopped { + s.Start() + if err := rt.Engine.StartMachine(podmanMachineName); err != nil { + return err + } + return rt.configureMachineForUsage(iMachine) + } + } + + // Otherwise, initialize the machine + s.Start() + if err := rt.Engine.InitializeMachine(podmanMachineName); err != nil { + return err + } + + return rt.getAndConfigureMachineForUsage(podmanMachineName) +} + +// stopAndKillMachine attempts to stop and kill the Podman machine. +// If other projects are running, it will leave the machine up. +func (rt PodmanRuntime) stopAndKillMachine() error { + // If the machine doesn't exist, exit early. + if !rt.astroMachineExists() { + return nil + } + + // If the machine exists, and its running, we need to check + // if any other projects are running. If other projects are running, + // we'll leave the machine up, otherwise we stop and kill it. + if rt.astroMachineIsRunning() { + // Get the containers that are running on our machine. + containers, err := rt.Engine.ListContainers() + if err != nil { + return err + } + + // Check the container labels to identify if other projects are running. + projectNames := make(map[string]struct{}) + for _, item := range containers { + // Check if "project.name" exists in the Labels map + if projectName, exists := item.Labels[composeProjectLabel]; exists { + // Add the project name to the map (map keys are unique) + projectNames[projectName] = struct{}{} + } + } + + // At this point in the command hook lifecycle, our project has already been stopped, + // and we are checking to see if any additional projects are running. + if len(projectNames) > 0 { + return nil + } + + // If we made it this far, we can stop the machine, + // as there are no more projects running. + err = rt.Engine.StopMachine(podmanMachineName) + if err != nil { + return err + } + } + + // If we make it here, the machine was already stopped, or was just stopped above. + // We can now remove it. + err := rt.Engine.RemoveMachine(podmanMachineName) + if err != nil { + return err + } + + return nil +} + +// configureMachineForUsage does two things: +// - Sets the DOCKER_HOST environment variable to the machine's socket path +// This allows the docker compose library to function as expected. +// - Sets the podman default connection to the machine +// This allows the podman command to function as expected. +func (rt PodmanRuntime) configureMachineForUsage(machine *types.InspectedMachine) error { + if machine == nil { + return fmt.Errorf("machine does not exist") + } + + // Compute our DOCKER_HOST value depending on the OS. + dockerHost := "unix://" + machine.ConnectionInfo.PodmanSocket.Path + if rt.OSChecker.IsWindows() { + dockerHost = "npipe:////./pipe/podman-" + machine.Name + } + + // Set the DOCKER_HOST environment variable for compose. + err := os.Setenv("DOCKER_HOST", dockerHost) + if err != nil { + return fmt.Errorf("error setting DOCKER_HOST: %s", err) + } + + // Set the podman default connection to our machine. + return rt.Engine.SetMachineAsDefault(machine.Name) +} + +// getAndConfigureMachineForUsage gets our astro machine +// then configures the host machine to use it. +func (rt PodmanRuntime) getAndConfigureMachineForUsage(name string) error { + machine, err := rt.Engine.InspectMachine(name) + if err != nil { + return err + } + return rt.configureMachineForUsage(machine) +} + +// getAstroMachine gets our astro podman machine. +func (rt PodmanRuntime) getAstroMachine() *types.ListedMachine { + machines, _ := rt.Engine.ListMachines() + return findMachineByName(machines, podmanMachineName) +} + +// astroMachineExists checks if our astro podman machine exists. +func (rt PodmanRuntime) astroMachineExists() bool { + machine := rt.getAstroMachine() + return machine != nil +} + +// astroMachineIsRunning checks if our astro podman machine is running. +func (rt PodmanRuntime) astroMachineIsRunning() bool { + machine := rt.getAstroMachine() + return machine != nil && machine.Running +} + +// isAnotherMachineRunning checks if another, non-astro podman machine is running. +func (rt PodmanRuntime) isAnotherMachineRunning() string { + machines, _ := rt.Engine.ListMachines() + for _, machine := range machines { + if machine.Running && machine.Name != podmanMachineName { + return machine.Name + } + } + return "" +} + +// findMachineByName finds a machine by name from a list of machines. +func findMachineByName(items []types.ListedMachine, name string) *types.ListedMachine { + for _, item := range items { + if item.Name == name { + return &item + } + } return nil } + +// isDockerHostSet checks if the DOCKER_HOST environment variable is set. +func isDockerHostSet() bool { + return os.Getenv("DOCKER_HOST") != "" +} + +// isDockerHostSetToAstroMachine checks if the DOCKER_HOST environment variable +// is pointing to the astro machine. +func isDockerHostSetToAstroMachine() bool { + return strings.Contains(os.Getenv("DOCKER_HOST"), podmanMachineName) +} diff --git a/airflow/runtimes/podman_runtime_test.go b/airflow/runtimes/podman_runtime_test.go new file mode 100644 index 000000000..1ce575cf9 --- /dev/null +++ b/airflow/runtimes/podman_runtime_test.go @@ -0,0 +1,466 @@ +package runtimes + +import ( + "os" + "testing" + + "github.com/astronomer/astro-cli/airflow/runtimes/mocks" + "github.com/astronomer/astro-cli/airflow/runtimes/types" + "github.com/stretchr/testify/suite" + + "github.com/stretchr/testify/assert" +) + +var ( + mockListedMachines []types.ListedMachine + mockListedContainers []types.ListedContainer + mockInspectedAstroMachine *types.InspectedMachine + mockInspectedOtherMachine *types.InspectedMachine + mockPodmanEngine *mocks.PodmanEngine + mockPodmanOSChecker *mocks.OSChecker +) + +type PodmanRuntimeSuite struct { + suite.Suite +} + +func TestPodmanRuntime(t *testing.T) { + suite.Run(t, new(PodmanRuntimeSuite)) +} + +// Setenv is a helper function to set an environment variable. +// It panics if an error occurs. +func (s *PodmanRuntimeSuite) Setenv(key, value string) { + if err := os.Setenv(key, value); err != nil { + panic(err) + } +} + +// Unsetenv is a helper function to unset an environment variable. +// It panics if an error occurs. +func (s *PodmanRuntimeSuite) Unsetenv(key string) { + if err := os.Unsetenv(key); err != nil { + panic(err) + } +} + +func (s *PodmanRuntimeSuite) SetupTest() { + // Reset some variables to defaults. + s.Unsetenv("DOCKER_HOST") + mockPodmanEngine = new(mocks.PodmanEngine) + mockPodmanOSChecker = new(mocks.OSChecker) + mockListedMachines = []types.ListedMachine{} + mockListedContainers = []types.ListedContainer{} + mockInspectedAstroMachine = &types.InspectedMachine{ + Name: "astro-machine", + ConnectionInfo: types.ConnectionInfo{ + PodmanSocket: types.PodmanSocket{ + Path: "/path/to/astro-machine.sock", + }, + }, + } + mockInspectedOtherMachine = &types.InspectedMachine{ + Name: "other-machine", + ConnectionInfo: types.ConnectionInfo{ + PodmanSocket: types.PodmanSocket{ + Path: "/path/to/other-machine.sock", + }, + }, + } +} + +func (s *PodmanRuntimeSuite) TestPodmanRuntimeInitializeDockerHostAlreadySet() { + s.Run("DOCKER_HOST is already set, abort initialization", func() { + // Set up mocks. + s.Setenv("DOCKER_HOST", "some_value") + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.Initialize() + assert.Nil(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +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("InspectMachine", podmanMachineName).Return(mockInspectedAstroMachine, nil) + mockPodmanEngine.On("SetMachineAsDefault", podmanMachineName).Return(nil) + mockPodmanOSChecker.On("IsWindows").Return(false) + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.Initialize() + assert.Nil(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +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("InspectMachine", podmanMachineName).Return(mockInspectedAstroMachine, nil) + mockPodmanEngine.On("SetMachineAsDefault", podmanMachineName).Return(nil) + mockPodmanOSChecker.On("IsWindows").Return(true) + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.Initialize() + assert.Nil(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +func (s *PodmanRuntimeSuite) TestPodmanRuntimeInitializeWithAnotherMachineRunningOnMac() { + s.Run("Another machine running on mac, stop it and start the astro machine", func() { + // Set up mocks. + mockListedMachines = []types.ListedMachine{ + { + Name: "other-machine", + Running: true, + }, + } + mockPodmanEngine.On("ListMachines").Return(mockListedMachines, nil).Once() + mockPodmanEngine.On("InspectMachine", mockListedMachines[0].Name).Return(mockInspectedOtherMachine, nil).Once() + mockPodmanEngine.On("SetMachineAsDefault", mockListedMachines[0].Name).Return(nil).Once() + 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("InspectMachine", podmanMachineName).Return(mockInspectedAstroMachine, nil).Once() + mockPodmanEngine.On("SetMachineAsDefault", podmanMachineName).Return(nil).Once() + mockPodmanOSChecker.On("IsMac").Return(true) + mockPodmanOSChecker.On("IsWindows").Return(false) + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.Initialize() + assert.Nil(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +func (s *PodmanRuntimeSuite) TestPodmanRuntimeInitializeWithAnotherMachineRunningWithExistingContainersOnMac() { + s.Run("Another machine running on mac, that has existing containers, return error message", func() { + // Set up mocks. + mockListedMachines = []types.ListedMachine{ + { + Name: "other-machine", + Running: true, + }, + } + mockListedContainers = []types.ListedContainer{ + { + Name: "container1", + Labels: map[string]string{composeProjectLabel: "project1"}, + }, + } + mockPodmanEngine.On("ListMachines").Return(mockListedMachines, nil).Once() + mockPodmanEngine.On("InspectMachine", mockListedMachines[0].Name).Return(mockInspectedOtherMachine, nil).Once() + mockPodmanEngine.On("SetMachineAsDefault", mockListedMachines[0].Name).Return(nil).Once() + mockPodmanEngine.On("ListContainers").Return(mockListedContainers, nil) + mockPodmanOSChecker.On("IsMac").Return(true) + mockPodmanOSChecker.On("IsWindows").Return(false) + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.Initialize() + assert.Error(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +func (s *PodmanRuntimeSuite) TestPodmanRuntimeInitializeAstroMachineAlreadyRunning() { + s.Run("Astro machine is already running, just configure it for usage", func() { + // Set up mocks. + mockListedMachines = []types.ListedMachine{ + { + Name: "astro-machine", + Running: true, + }, + } + mockInspectedAstroMachine.State = podmanStatusRunning + mockPodmanEngine.On("ListMachines").Return(mockListedMachines, nil) + mockPodmanEngine.On("InspectMachine", podmanMachineName).Return(mockInspectedAstroMachine, nil) + mockPodmanEngine.On("SetMachineAsDefault", podmanMachineName).Return(nil) + mockPodmanOSChecker.On("IsWindows").Return(false) + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.Initialize() + assert.Nil(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +func (s *PodmanRuntimeSuite) TestPodmanRuntimeInitializeAstroMachineExistsButStopped() { + s.Run("Astro machine already exists, but is in stopped state, start and configure it for usage", func() { + // Set up mocks. + mockListedMachines = []types.ListedMachine{ + { + Name: "astro-machine", + Running: true, + }, + } + mockInspectedAstroMachine.State = podmanStatusStopped + mockPodmanEngine.On("ListMachines").Return(mockListedMachines, nil) + mockPodmanEngine.On("InspectMachine", podmanMachineName).Return(mockInspectedAstroMachine, nil) + mockPodmanEngine.On("SetMachineAsDefault", podmanMachineName).Return(nil) + mockPodmanEngine.On("StartMachine", podmanMachineName).Return(nil) + mockPodmanOSChecker.On("IsWindows").Return(false) + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.Initialize() + assert.Nil(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +func (s *PodmanRuntimeSuite) TestPodmanRuntimeConfigureDockerHostAlreadySet() { + s.Run("DOCKER_HOST is already set, abort configure", func() { + // Set up mocks. + s.Setenv("DOCKER_HOST", "some_value") + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.Configure() + assert.Nil(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +func (s *PodmanRuntimeSuite) TestPodmanRuntimeConfigureAstroMachineRunning() { + s.Run("Astro machine is already running, so configure it for usage", func() { + // Set up mocks. + mockListedMachines = []types.ListedMachine{ + { + Name: "astro-machine", + Running: true, + }, + } + mockPodmanEngine.On("ListMachines").Return(mockListedMachines, nil) + mockPodmanEngine.On("InspectMachine", podmanMachineName).Return(mockInspectedAstroMachine, nil) + mockPodmanEngine.On("SetMachineAsDefault", podmanMachineName).Return(nil) + mockPodmanOSChecker.On("IsWindows").Return(false) + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.Configure() + assert.Nil(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +func (s *PodmanRuntimeSuite) TestPodmanRuntimeConfigureAstroMachineNotRunning() { + s.Run("Astro machine is not already running, so return error message", func() { + // Set up mocks. + mockPodmanEngine.On("ListMachines").Return(mockListedMachines, nil) + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.Configure() + assert.Error(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +func (s *PodmanRuntimeSuite) TestPodmanRuntimeConfigureOrKillDockerHostAlreadySet() { + s.Run("Astro machine is not already running, so return error message", func() { + // Set up mocks. + s.Setenv("DOCKER_HOST", "some_value") + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.ConfigureOrKill() + assert.Nil(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +func (s *PodmanRuntimeSuite) TestPodmanRuntimeConfigureOrKillAstroMachineRunning() { + s.Run("Astro machine is already running, so configure it for usage", func() { + // Set up mocks. + mockListedMachines = []types.ListedMachine{ + { + Name: "astro-machine", + Running: true, + }, + } + mockPodmanEngine.On("ListMachines").Return(mockListedMachines, nil) + mockPodmanEngine.On("InspectMachine", podmanMachineName).Return(mockInspectedAstroMachine, nil) + mockPodmanEngine.On("SetMachineAsDefault", podmanMachineName).Return(nil) + mockPodmanOSChecker.On("IsWindows").Return(false) + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.ConfigureOrKill() + assert.Nil(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +func (s *PodmanRuntimeSuite) TestPodmanRuntimeConfigureOrKillAstroMachineStopped() { + s.Run("Astro machine is stopped, proceed to kill it", func() { + // Set up mocks. + mockListedMachines = []types.ListedMachine{ + { + Name: "astro-machine", + Running: false, + }, + } + mockPodmanEngine.On("ListMachines").Return(mockListedMachines, nil) + mockPodmanEngine.On("RemoveMachine", podmanMachineName).Return(nil) + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.ConfigureOrKill() + assert.Error(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +func (s *PodmanRuntimeSuite) TestPodmanRuntimeConfigureOrKillAstroMachineNotRunning() { + s.Run("Astro machine is not already running, so return error message", func() { + // Set up mocks. + mockPodmanEngine.On("ListMachines").Return(mockListedMachines, nil) + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.ConfigureOrKill() + assert.Error(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +func (s *PodmanRuntimeSuite) TestPodmanRuntimeKillOtherProjectRunning() { + s.Run("Astro machine running, but another project is still running, so do not stop and kill machine", func() { + // Set up mocks. + s.Setenv("DOCKER_HOST", podmanMachineName) + mockListedMachines = []types.ListedMachine{ + { + Name: "astro-machine", + Running: true, + }, + } + mockListedContainers = []types.ListedContainer{ + { + Name: "container1", + Labels: map[string]string{composeProjectLabel: "project1"}, + }, + } + mockPodmanEngine.On("ListMachines").Return(mockListedMachines, nil) + mockPodmanEngine.On("ListContainers").Return(mockListedContainers, nil) + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.Kill() + assert.Nil(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +func (s *PodmanRuntimeSuite) TestPodmanRuntimeKill() { + s.Run("Astro machine running, no other projects running, so stop and kill the machine", func() { + // Set up mocks. + s.Setenv("DOCKER_HOST", podmanMachineName) + mockListedMachines = []types.ListedMachine{ + { + Name: "astro-machine", + Running: true, + }, + } + mockPodmanEngine.On("ListMachines").Return(mockListedMachines, nil) + mockPodmanEngine.On("ListContainers").Return(mockListedContainers, nil) + mockPodmanEngine.On("StopMachine", podmanMachineName).Return(nil) + mockPodmanEngine.On("RemoveMachine", podmanMachineName).Return(nil) + // Create the runtime with our mock engine and os checker. + rt := CreatePodmanRuntime(mockPodmanEngine, mockPodmanOSChecker) + // Run our test and assert expectations. + err := rt.Kill() + assert.Nil(s.T(), err) + mockPodmanEngine.AssertExpectations(s.T()) + mockPodmanOSChecker.AssertExpectations(s.T()) + }) +} + +func (s *PodmanRuntimeSuite) TestFindMachineByName() { + s.Run("Returns machine when name matches", func() { + machines := []types.ListedMachine{ + {Name: "astro-machine"}, + {Name: "other-machine"}, + } + result := findMachineByName(machines, "astro-machine") + assert.NotNil(s.T(), result) + assert.Equal(s.T(), "astro-machine", result.Name) + }) + + s.Run("Returns nil when no match found", func() { + machines := []types.ListedMachine{ + {Name: "astro-machine"}, + {Name: "other-machine"}, + } + result := findMachineByName(machines, "non-existent-machine") + assert.Nil(s.T(), result) + }) + + s.Run("Returns nil when list is empty", func() { + var machines []types.ListedMachine + result := findMachineByName(machines, "astro-machine") + assert.Nil(s.T(), result) + }) +} + +func (s *PodmanRuntimeSuite) TestIsDockerHostSet() { + s.Run("DOCKER_HOST is set and returns true", func() { + s.Setenv("DOCKER_HOST", "some_value") + result := isDockerHostSet() + assert.True(s.T(), result) + }) + + s.Run("DOCKER_HOST is set and returns true", func() { + s.Unsetenv("DOCKER_HOST") + result := isDockerHostSet() + assert.False(s.T(), result) + }) +} + +func (s *PodmanRuntimeSuite) TestIsDockerHostSetToAstroMachine() { + s.Run("DOCKER_HOST is set to astro-machine and returns true", func() { + s.Setenv("DOCKER_HOST", "unix:///path/to/astro-machine.sock") + result := isDockerHostSetToAstroMachine() + assert.True(s.T(), result) + }) + + s.Run("DOCKER_HOST is set to other-machine and returns false", func() { + s.Setenv("DOCKER_HOST", "unix:///path/to/other-machine.sock") + result := isDockerHostSetToAstroMachine() + assert.False(s.T(), result) + }) + + s.Run("DOCKER_HOST is not set and returns false", func() { + s.Unsetenv("DOCKER_HOST") + result := isDockerHostSetToAstroMachine() + assert.False(s.T(), result) + }) +} diff --git a/airflow/runtimes/types/podman.go b/airflow/runtimes/types/podman.go new file mode 100644 index 000000000..0863e4be6 --- /dev/null +++ b/airflow/runtimes/types/podman.go @@ -0,0 +1,35 @@ +package types + +// ListedMachine contains information about a Podman machine +// as it is provided from the `podman machine ls --format json` command. +type ListedMachine struct { + Name string + Running bool + Starting bool + LastUp string +} + +// PodmanSocket contains the path to the Podman socket. +type PodmanSocket struct { + Path string +} + +// ConnectionInfo contains information about the connection to a Podman machine. +type ConnectionInfo struct { + PodmanSocket PodmanSocket +} + +// InspectedMachine contains information about a Podman machine +// as it is provided from the `podman machine inspect` command. +type InspectedMachine struct { + Name string + ConnectionInfo ConnectionInfo + State string +} + +// ListedContainer contains information about a Podman container +// as it is provided from the `podman ps --format json` command. +type ListedContainer struct { + Name string + Labels map[string]string +} diff --git a/airflow/runtimes/utils.go b/airflow/runtimes/utils.go deleted file mode 100644 index 19f0efd6b..000000000 --- a/airflow/runtimes/utils.go +++ /dev/null @@ -1,15 +0,0 @@ -package runtimes - -import "runtime" - -// isWindows is a utility function to determine if the CLI host machine -// is running on Microsoft Windows OS. -func isWindows() bool { - return runtime.GOOS == "windows" -} - -// isMac is a utility function to determine if the CLI host machine -// is running on Apple macOS. -func isMac() bool { - return runtime.GOOS == "darwin" -} diff --git a/cmd/airflow.go b/cmd/airflow.go index b607a6d77..bc4413adf 100644 --- a/cmd/airflow.go +++ b/cmd/airflow.go @@ -16,7 +16,6 @@ import ( astrocore "github.com/astronomer/astro-cli/astro-client-core" astroplatformcore "github.com/astronomer/astro-cli/astro-client-platform-core" "github.com/astronomer/astro-cli/cloud/environment" - "github.com/astronomer/astro-cli/cmd/utils" "github.com/astronomer/astro-cli/config" "github.com/astronomer/astro-cli/context" "github.com/astronomer/astro-cli/houston" @@ -193,10 +192,10 @@ func newAirflowInitCmd() *cobra.Command { func newAirflowUpgradeTestCmd(platformCoreClient astroplatformcore.CoreClient) *cobra.Command { cmd := &cobra.Command{ - Use: "upgrade-test", - Short: "Run tests to see if your environment and DAGs are compatible with a new version of Airflow or Astro Runtime. This test will produce a series of reports where you can see the test results.", - Long: "Run tests to see if your environment and DAGs are compatible with a new version of Airflow or Astro Runtime. This test will produce a series of reports where you can see the test results.", - PersistentPreRunE: DoNothing, + Use: "upgrade-test", + Short: "Run tests to see if your environment and DAGs are compatible with a new version of Airflow or Astro Runtime. This test will produce a series of reports where you can see the test results.", + Long: "Run tests to see if your environment and DAGs are compatible with a new version of Airflow or Astro Runtime. This test will produce a series of reports where you can see the test results.", + PreRunE: EnsureRuntime, RunE: func(cmd *cobra.Command, args []string) error { return airflowUpgradeTest(cmd, platformCoreClient) }, @@ -260,12 +259,11 @@ func newAirflowStartCmd(astroCoreClient astrocore.CoreClient) *cobra.Command { func newAirflowPSCmd() *cobra.Command { cmd := &cobra.Command{ - Use: "ps", - Short: "List locally running Airflow containers", - Long: "List locally running Airflow containers", - PersistentPreRunE: DoNothing, - PreRunE: utils.EnsureProjectDir, - RunE: airflowPS, + Use: "ps", + Short: "List locally running Airflow containers", + Long: "List locally running Airflow containers", + PreRunE: SetRuntimeIfExists, + RunE: airflowPS, } return cmd } @@ -275,23 +273,21 @@ func newAirflowRunCmd() *cobra.Command { Use: "run", Short: "Run Airflow CLI commands within your local Airflow environment", Long: "Run Airflow CLI commands within your local Airflow environment. These commands run in the webserver container but can interact with your local scheduler, workers, and metadata database.", - PreRunE: utils.EnsureProjectDir, + PreRunE: EnsureRuntime, RunE: airflowRun, Example: RunExample, DisableFlagParsing: true, - PersistentPreRunE: DoNothing, } return cmd } func newAirflowLogsCmd() *cobra.Command { cmd := &cobra.Command{ - Use: "logs", - Short: "Display component logs for your local Airflow environment", - Long: "Display scheduler, worker, and webserver logs for your local Airflow environment", - PreRunE: utils.EnsureProjectDir, - RunE: airflowLogs, - PersistentPreRunE: DoNothing, + Use: "logs", + Short: "Display component logs for your local Airflow environment", + Long: "Display scheduler, worker, and webserver logs for your local Airflow environment", + PreRunE: SetRuntimeIfExists, + RunE: airflowLogs, } cmd.Flags().BoolVarP(&followLogs, "follow", "f", false, "Follow log output") cmd.Flags().BoolVarP(&schedulerLogs, "scheduler", "s", false, "Output scheduler logs") @@ -302,35 +298,33 @@ func newAirflowLogsCmd() *cobra.Command { func newAirflowStopCmd() *cobra.Command { cmd := &cobra.Command{ - Use: "stop", - Short: "Stop all locally running Airflow containers", - Long: "Stop all Airflow containers running on your local machine. This command preserves container data.", - PersistentPreRunE: DoNothing, - PreRunE: utils.EnsureProjectDir, - RunE: airflowStop, + Use: "stop", + Short: "Stop all locally running Airflow containers", + Long: "Stop all Airflow containers running on your local machine. This command preserves container data.", + PreRunE: SetRuntimeIfExists, + RunE: airflowStop, } return cmd } func newAirflowKillCmd() *cobra.Command { cmd := &cobra.Command{ - Use: "kill", - Short: "Kill all locally running Airflow containers", - Long: "Kill all Airflow containers running on your local machine. This command permanently deletes all container data.", - PersistentPreRunE: DoNothing, - PreRunE: utils.EnsureProjectDir, - RunE: airflowKill, + Use: "kill", + Short: "Kill all locally running Airflow containers", + Long: "Kill all Airflow containers running on your local machine. This command permanently deletes all container data.", + PreRunE: KillPreRunHook, + RunE: airflowKill, + PostRunE: KillPostRunHook, } return cmd } func newAirflowRestartCmd(astroCoreClient astrocore.CoreClient) *cobra.Command { cmd := &cobra.Command{ - Use: "restart", - Short: "Restart all locally running Airflow containers", - Long: "Restart all Airflow containers running on your local machine. This command stops and then starts locally running containers to apply changes to your local environment.", - PersistentPreRunE: DoNothing, - PreRunE: utils.EnsureProjectDir, + Use: "restart", + Short: "Restart all locally running Airflow containers", + Long: "Restart all Airflow containers running on your local machine. This command stops and then starts locally running containers to apply changes to your local environment.", + PreRunE: SetRuntimeIfExists, RunE: func(cmd *cobra.Command, args []string) error { return airflowRestart(cmd, args, astroCoreClient) }, @@ -351,12 +345,11 @@ func newAirflowRestartCmd(astroCoreClient astrocore.CoreClient) *cobra.Command { func newAirflowPytestCmd() *cobra.Command { cmd := &cobra.Command{ - Use: "pytest [pytest file/directory]", - Short: "Run pytests in a local Airflow environment", - Long: "This command spins up a local Python environment to run pytests against your DAGs. If a specific pytest file is not specified, all pytests in the tests directory will be run. To run pytests with a different environment file, specify that with the '--env' flag. ", - PersistentPreRunE: DoNothing, - PreRunE: utils.EnsureProjectDir, - RunE: airflowPytest, + Use: "pytest [pytest file/directory]", + Short: "Run pytests in a local Airflow environment", + Long: "This command spins up a local Python environment to run pytests against your DAGs. If a specific pytest file is not specified, all pytests in the tests directory will be run. To run pytests with a different environment file, specify that with the '--env' flag. ", + PreRunE: EnsureRuntime, + RunE: airflowPytest, } cmd.Flags().StringVarP(&pytestArgs, "args", "a", "", "pytest arguments you'd like passed to the pytest command. Surround the args in quotes. For example 'astro dev pytest --args \"--cov-config path\"'") cmd.Flags().StringVarP(&envFile, "env", "e", ".env", "Location of file containing environment variables") @@ -368,13 +361,12 @@ func newAirflowPytestCmd() *cobra.Command { func newAirflowParseCmd() *cobra.Command { cmd := &cobra.Command{ - Use: "parse", - Short: "parse all DAGs in your Astro project for errors", - Long: "This command spins up a local Python environment and checks your DAGs for syntax and import errors.", - Args: cobra.MaximumNArgs(1), - PersistentPreRunE: DoNothing, - PreRunE: utils.EnsureProjectDir, - RunE: airflowParse, + Use: "parse", + Short: "parse all DAGs in your Astro project for errors", + Long: "This command spins up a local Python environment and checks your DAGs for syntax and import errors.", + Args: cobra.MaximumNArgs(1), + PreRunE: EnsureRuntime, + RunE: airflowParse, } cmd.Flags().StringVarP(&envFile, "env", "e", ".env", "Location of file containing environment variables") cmd.Flags().StringVarP(&customImageName, "image-name", "i", "", "Name of a custom built image to run parse with") @@ -388,8 +380,7 @@ func newAirflowUpgradeCheckCmd() *cobra.Command { Use: "upgrade-check", Short: "List DAG and config-level changes required to upgrade to Airflow 2.0", Long: "List DAG and config-level changes required to upgrade to Airflow 2.0", - PersistentPreRunE: DoNothing, - PreRunE: utils.EnsureProjectDir, + PreRunE: EnsureRuntime, RunE: airflowUpgradeCheck, DisableFlagParsing: true, } @@ -398,13 +389,12 @@ func newAirflowUpgradeCheckCmd() *cobra.Command { func newAirflowBashCmd() *cobra.Command { cmd := &cobra.Command{ - Use: "bash", - Short: "Exec into a running an Airflow container", - Long: "Use this command to Exec into either the Webserver, Sechduler, Postgres, or Triggerer Container to run bash commands", - Args: cobra.MaximumNArgs(1), - PersistentPreRunE: DoNothing, - PreRunE: utils.EnsureProjectDir, - RunE: airflowBash, + Use: "bash", + Short: "Exec into a running an Airflow container", + Long: "Use this command to Exec into either the Webserver, Sechduler, Postgres, or Triggerer Container to run bash commands", + Args: cobra.MaximumNArgs(1), + PreRunE: EnsureRuntime, + RunE: airflowBash, } cmd.Flags().BoolVarP(&schedulerExec, "scheduler", "s", false, "Exec into the scheduler container") cmd.Flags().BoolVarP(&webserverExec, "webserver", "w", false, "Exec into the webserver container") @@ -429,12 +419,11 @@ func newAirflowObjectRootCmd() *cobra.Command { func newObjectImportCmd() *cobra.Command { cmd := &cobra.Command{ - Use: "import", - Short: "Create and update local Airflow connections, variables, and pools from a local YAML file", - Long: "This command creates all connections, variables, and pools from a YAML configuration file in your local Airflow environment. Airflow must be running locally for this command to work", - PersistentPreRunE: DoNothing, - PreRunE: utils.EnsureProjectDir, - RunE: airflowSettingsImport, + Use: "import", + Short: "Create and update local Airflow connections, variables, and pools from a local YAML file", + Long: "This command creates all connections, variables, and pools from a YAML configuration file in your local Airflow environment. Airflow must be running locally for this command to work", + PreRunE: EnsureRuntime, + RunE: airflowSettingsImport, } cmd.Flags().BoolVarP(&connections, "connections", "c", false, "Import connections from a settings YAML file") cmd.Flags().BoolVarP(&variables, "variables", "v", false, "Import variables from a settings YAML file") @@ -445,13 +434,12 @@ func newObjectImportCmd() *cobra.Command { func newObjectExportCmd() *cobra.Command { cmd := &cobra.Command{ - Use: "export", - Short: "Export local Airflow connections, variables, pools, and startup configurations as YAML or environment variables.", - Long: "Export local Airflow connections, variables, or pools as YAML or environment variables. Airflow must be running locally to export Airflow objects. Use the '--compose' flag to export the Compose file used to start up Airflow.", - Args: cobra.MaximumNArgs(1), - PersistentPreRunE: DoNothing, - PreRunE: utils.EnsureProjectDir, - RunE: airflowSettingsExport, + Use: "export", + Short: "Export local Airflow connections, variables, pools, and startup configurations as YAML or environment variables.", + Long: "Export local Airflow connections, variables, or pools as YAML or environment variables. Airflow must be running locally to export Airflow objects. Use the '--compose' flag to export the Compose file used to start up Airflow.", + Args: cobra.MaximumNArgs(1), + PreRunE: EnsureRuntime, + RunE: airflowSettingsExport, } cmd.Flags().BoolVarP(&connections, "connections", "c", false, "Export connections to a settings YAML or env file") cmd.Flags().BoolVarP(&variables, "variables", "v", false, "Export variables to a settings YAML or env file") diff --git a/cmd/airflow_hooks.go b/cmd/airflow_hooks.go index c5a8abf0f..743d108f2 100644 --- a/cmd/airflow_hooks.go +++ b/cmd/airflow_hooks.go @@ -1,20 +1,27 @@ package cmd import ( + "fmt" + "os" + "path/filepath" + "github.com/astronomer/astro-cli/airflow/runtimes" "github.com/astronomer/astro-cli/cmd/utils" + "github.com/astronomer/astro-cli/config" "github.com/spf13/cobra" ) +const failedToCreatePluginsDir = "failed to create plugins directory: %w" + // DoNothing is a persistent pre-run hook that does nothing. // Used to clobber the standard astro dev persistent pre-run hook for select commands. func DoNothing(_ *cobra.Command, _ []string) error { return nil } -// ConfigureContainerRuntime sets up the containerRuntime variable. -// The containerRuntime variable is then used in the following pre-run and post-run hooks -// defined here. +// ConfigureContainerRuntime sets up the containerRuntime variable and is defined +// as a PersistentPreRunE hook for all astro dev sub-commands. The containerRuntime +// variable is then used in the following pre-run and post-run hooks defined here. func ConfigureContainerRuntime(_ *cobra.Command, _ []string) error { var err error containerRuntime, err = runtimes.GetContainerRuntime() @@ -30,6 +37,43 @@ func EnsureRuntime(cmd *cobra.Command, args []string) error { if err := utils.EnsureProjectDir(cmd, args); err != nil { return err } + + // Check if the OS is Windows and create the plugins project directory if it doesn't exist. + // In Windows, the compose project will fail if the plugins directory doesn't exist, due + // to the volume mounts we specify. + osChecker := runtimes.CreateOSChecker() + if osChecker.IsWindows() { + pluginsDir := filepath.Join(config.WorkingPath, "plugins") + if err := os.MkdirAll(pluginsDir, os.ModePerm); err != nil && !os.IsExist(err) { + return fmt.Errorf(failedToCreatePluginsDir, err) + } + } + // Initialize the runtime if it's not running. return containerRuntime.Initialize() } + +// SetRuntimeIfExists is a pre-run hook that ensures the project directory exists +// and sets the container runtime if its running, otherwise we bail with an error message. +func SetRuntimeIfExists(cmd *cobra.Command, args []string) error { + if err := utils.EnsureProjectDir(cmd, args); err != nil { + return err + } + return containerRuntime.Configure() +} + +// KillPreRunHook sets the container runtime if its running, +// otherwise we bail with an error message. +func KillPreRunHook(cmd *cobra.Command, args []string) error { + if err := utils.EnsureProjectDir(cmd, args); err != nil { + return err + } + return containerRuntime.ConfigureOrKill() +} + +// KillPostRunHook ensures that we stop and kill the +// podman machine once a project has been killed. +func KillPostRunHook(_ *cobra.Command, _ []string) error { + // Kill the runtime. + return containerRuntime.Kill() +} diff --git a/cmd/airflow_test.go b/cmd/airflow_test.go index 6a1688e58..8dd26ceef 100644 --- a/cmd/airflow_test.go +++ b/cmd/airflow_test.go @@ -35,6 +35,16 @@ func TestAirflow(t *testing.T) { suite.Run(t, new(AirflowSuite)) } +func (s *AirflowSuite) SetupTest() { + testUtil.InitTestConfig(testUtil.LocalPlatform) + dir, err := os.MkdirTemp("", "test_temp_dir_*") + if err != nil { + s.T().Fatalf("failed to create temp dir: %v", err) + } + s.tempDir = dir + config.WorkingPath = s.tempDir +} + func (s *AirflowSuite) SetupSubTest() { testUtil.InitTestConfig(testUtil.LocalPlatform) dir, err := os.MkdirTemp("", "test_temp_dir_*") @@ -140,28 +150,38 @@ func (s *AirflowSuite) TestNewAirflowInitCmd() { } func (s *AirflowSuite) TestNewAirflowRunCmd() { - cmd := newAirflowRunCmd() - s.Nil(cmd.PersistentPreRunE(new(cobra.Command), []string{})) + cmd := newAirflowInitCmd() + cmd.RunE(new(cobra.Command), []string{}) + cmd = newAirflowRunCmd() + s.Nil(cmd.PreRunE(new(cobra.Command), []string{})) } func (s *AirflowSuite) TestNewAirflowPSCmd() { - cmd := newAirflowPSCmd() - s.Nil(cmd.PersistentPreRunE(new(cobra.Command), []string{})) + cmd := newAirflowInitCmd() + cmd.RunE(new(cobra.Command), []string{}) + cmd = newAirflowPSCmd() + s.Nil(cmd.PreRunE(new(cobra.Command), []string{})) } func (s *AirflowSuite) TestNewAirflowLogsCmd() { - cmd := newAirflowLogsCmd() - s.Nil(cmd.PersistentPreRunE(new(cobra.Command), []string{})) + cmd := newAirflowInitCmd() + cmd.RunE(new(cobra.Command), []string{}) + cmd = newAirflowLogsCmd() + s.Nil(cmd.PreRunE(new(cobra.Command), []string{})) } func (s *AirflowSuite) TestNewAirflowKillCmd() { - cmd := newAirflowKillCmd() - s.Nil(cmd.PersistentPreRunE(new(cobra.Command), []string{})) + cmd := newAirflowInitCmd() + cmd.RunE(new(cobra.Command), []string{}) + cmd = newAirflowKillCmd() + s.Nil(cmd.PreRunE(new(cobra.Command), []string{})) } func (s *AirflowSuite) TestNewAirflowUpgradeCheckCmd() { - cmd := newAirflowUpgradeCheckCmd() - s.Nil(cmd.PersistentPreRunE(new(cobra.Command), []string{})) + cmd := newAirflowInitCmd() + cmd.RunE(new(cobra.Command), []string{}) + cmd = newAirflowUpgradeCheckCmd() + s.Nil(cmd.PreRunE(new(cobra.Command), []string{})) } func (s *AirflowSuite) Test_airflowInitNonEmptyDir() { diff --git a/config/config.go b/config/config.go index 9d0efa9e7..95a6b2f75 100644 --- a/config/config.go +++ b/config/config.go @@ -86,6 +86,8 @@ var ( DisableAstroRun: newCfg("disable_astro_run", "false"), DisableEnvObjects: newCfg("disable_env_objects", "false"), AutoSelect: newCfg("auto_select", "false"), + MachineCPU: newCfg("machine.cpu", "2"), + MachineMemory: newCfg("machine.memory", "4096"), } // viperHome is the viper object in the users home directory diff --git a/config/types.go b/config/types.go index 9adca30bd..a747d51a2 100644 --- a/config/types.go +++ b/config/types.go @@ -45,6 +45,8 @@ type cfgs struct { DisableAstroRun cfg DisableEnvObjects cfg AutoSelect cfg + MachineCPU cfg + MachineMemory cfg } // Creates a new cfg struct