From bee657a04ad9e454b8211dbf593659db1ed0a863 Mon Sep 17 00:00:00 2001
From: Neel Dalsania <neel.dalsania@astronomer.io>
Date: Mon, 1 Aug 2022 19:04:51 +0530
Subject: [PATCH] added support for multi arch runtime images (#655)

---
 airflow/docker_image.go      | 4 ++++
 airflow/docker_image_test.go | 5 +++--
 airflow/docker_test.go       | 8 ++++++++
 airflow/types/types.go       | 7 ++++---
 cloud/deploy/deploy.go       | 4 +++-
 cloud/deploy/deploy_test.go  | 9 +++++++++
 software/deploy/deploy.go    | 8 ++++++--
 7 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/airflow/docker_image.go b/airflow/docker_image.go
index caa8bc037..0422f109c 100644
--- a/airflow/docker_image.go
+++ b/airflow/docker_image.go
@@ -55,6 +55,10 @@ func (d *DockerImage) Build(config airflowTypes.ImageBuildConfig) error {
 	if config.NoCache {
 		args = append(args, "--no-cache")
 	}
+
+	if len(config.TargetPlatforms) > 0 {
+		args = append(args, fmt.Sprintf("--platform=%s", strings.Join(config.TargetPlatforms, ",")))
+	}
 	// Build image
 	var stdout, stderr io.Writer
 	if config.Output {
diff --git a/airflow/docker_image_test.go b/airflow/docker_image_test.go
index 0c02f76a8..7e5e9a685 100644
--- a/airflow/docker_image_test.go
+++ b/airflow/docker_image_test.go
@@ -24,8 +24,9 @@ func TestDockerImageBuild(t *testing.T) {
 	assert.NoError(t, err)
 
 	options := airflowTypes.ImageBuildConfig{
-		Path:    cwd,
-		NoCache: false,
+		Path:            cwd,
+		TargetPlatforms: []string{"linux/amd64"},
+		NoCache:         false,
 	}
 
 	previousCmdExec := cmdExec
diff --git a/airflow/docker_test.go b/airflow/docker_test.go
index f61f5f1e7..32777a4cc 100644
--- a/airflow/docker_test.go
+++ b/airflow/docker_test.go
@@ -875,6 +875,10 @@ func TestCheckWebserverHealth(t *testing.T) {
 			mockEventsCall.ReturnArguments = mock.Arguments{err}
 		}
 
+		openURL = func(url string) error {
+			return nil
+		}
+
 		orgInitSetting := initSettings
 		initSettings = func(id string, version uint64) error {
 			return nil
@@ -912,6 +916,10 @@ func TestCheckWebserverHealth(t *testing.T) {
 			mockEventsCall.ReturnArguments = mock.Arguments{err}
 		}
 
+		openURL = func(url string) error {
+			return nil
+		}
+
 		err := checkWebserverHealth(&types.Project{Name: "test"}, composeMock, 2)
 		assert.ErrorIs(t, err, errMockDocker)
 	})
diff --git a/airflow/types/types.go b/airflow/types/types.go
index e97f1fb79..4a3c85cbc 100644
--- a/airflow/types/types.go
+++ b/airflow/types/types.go
@@ -2,7 +2,8 @@ package types
 
 // ImageBuildConfig defines options when building a container image
 type ImageBuildConfig struct {
-	Path    string
-	NoCache bool
-	Output  bool
+	Path            string
+	TargetPlatforms []string
+	NoCache         bool
+	Output          bool
 }
diff --git a/cloud/deploy/deploy.go b/cloud/deploy/deploy.go
index 61a736fd9..9c7e1d9ba 100644
--- a/cloud/deploy/deploy.go
+++ b/cloud/deploy/deploy.go
@@ -43,6 +43,8 @@ var (
 	pytestFile string
 	dockerfile = "Dockerfile"
 
+	deployImagePlatformSupport = []string{"linux/amd64"}
+
 	// Monkey patched to write unit tests
 	airflowImageHandler  = airflow.ImageHandlerInit
 	containerHandlerInit = airflow.ContainerHandlerInit
@@ -261,7 +263,7 @@ func buildImage(c *config.Context, path, currentVersion, deployImage, imageName
 	imageHandler := airflowImageHandler(deployImage)
 
 	if imageName == "" {
-		err := imageHandler.Build(types.ImageBuildConfig{Path: path, Output: true})
+		err := imageHandler.Build(types.ImageBuildConfig{Path: path, Output: true, TargetPlatforms: deployImagePlatformSupport})
 		if err != nil {
 			return "", err
 		}
diff --git a/cloud/deploy/deploy_test.go b/cloud/deploy/deploy_test.go
index 1461f501e..74b213148 100644
--- a/cloud/deploy/deploy_test.go
+++ b/cloud/deploy/deploy_test.go
@@ -171,6 +171,15 @@ func TestBuildImageFailure(t *testing.T) {
 	ctx.SetSystemAdmin(true)
 
 	mockImageHandler := new(mocks.ImageHandler)
+
+	// image build failure
+	airflowImageHandler = func(image string) airflow.ImageHandler {
+		mockImageHandler.On("Build", mock.Anything).Return(errMock).Once()
+		return mockImageHandler
+	}
+	_, err = buildImage(&ctx, "./testfiles/", "4.2.5", "", "", nil)
+	assert.ErrorIs(t, err, errMock)
+
 	airflowImageHandler = func(image string) airflow.ImageHandler {
 		mockImageHandler.On("Build", mock.Anything).Return(nil)
 		mockImageHandler.On("GetLabel", runtimeImageLabel).Return("4.2.5", nil)
diff --git a/software/deploy/deploy.go b/software/deploy/deploy.go
index c9d85f77c..2abb2b619 100644
--- a/software/deploy/deploy.go
+++ b/software/deploy/deploy.go
@@ -23,6 +23,8 @@ var (
 	imageHandlerInit = airflow.ImageHandlerInit
 
 	dockerfile = "Dockerfile"
+
+	deployImagePlatformSupport = []string{"linux/amd64"}
 )
 
 var (
@@ -216,8 +218,10 @@ func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Conte
 	imageHandler := imageHandlerInit(imageName)
 
 	buildConfig := types.ImageBuildConfig{
-		Path:    config.WorkingPath,
-		NoCache: ignoreCacheDeploy,
+		Path:            config.WorkingPath,
+		NoCache:         ignoreCacheDeploy,
+		TargetPlatforms: deployImagePlatformSupport,
+		Output:          true,
 	}
 	err = imageHandler.Build(buildConfig)
 	if err != nil {