From c13b5f8383ae55c36a65fc890675b6304acafc51 Mon Sep 17 00:00:00 2001 From: Rujhan Arora <95576577+rujhan-arora-astronomer@users.noreply.github.com> Date: Wed, 11 Dec 2024 17:51:12 +0530 Subject: [PATCH] Introduction of sha_as_tag=True config in astro cli which will be used during Astro deploy (#1752) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- airflow/container.go | 1 + airflow/docker_image.go | 23 +++++++++++++++++++++++ airflow/mocks/ImageHandler.go | 24 ++++++++++++++++++++++++ config/config.go | 1 + config/types.go | 1 + software/deploy/deploy.go | 12 +++++++++--- software/deploy/deploy_test.go | 32 ++++++++++++++++++++++++++++++++ 7 files changed, 91 insertions(+), 3 deletions(-) diff --git a/airflow/container.go b/airflow/container.go index 0264f64ca..2f671d24d 100644 --- a/airflow/container.go +++ b/airflow/container.go @@ -56,6 +56,7 @@ type ImageHandler interface { Pytest(pytestFile, airflowHome, envFile, testHomeDirectory string, pytestArgs []string, htmlReport bool, config types.ImageBuildConfig) (string, error) ConflictTest(workingDirectory, testHomeDirectory string, buildConfig types.ImageBuildConfig) (string, error) CreatePipFreeze(altImageName, pipFreezeFile string) error + GetImageSha() (string, error) } type DockerComposeAPI interface { diff --git a/airflow/docker_image.go b/airflow/docker_image.go index 7e1807043..4d58a5c58 100644 --- a/airflow/docker_image.go +++ b/airflow/docker_image.go @@ -410,6 +410,29 @@ func (d *DockerImage) Push(remoteImage, username, token string) error { return nil } +func (d *DockerImage) GetImageSha() (string, error) { + containerRuntime, err := runtimes.GetContainerRuntimeBinary() + if err != nil { + return "", err + } + // Get the digest of the pushed image + remoteDigest := "" + out := &bytes.Buffer{} + err = cmdExec(containerRuntime, out, nil, "inspect", "--format={{index .RepoDigests 0}}", d.imageName) + if err != nil { + return remoteDigest, fmt.Errorf("failed to get digest for image %s: %w", d.imageName, err) + } + // Parse and clean the output + digestOutput := strings.TrimSpace(out.String()) + if digestOutput != "" { + parts := strings.Split(digestOutput, "@") + if len(parts) == 2 { + remoteDigest = parts[1] // Extract the digest part (after '@') + } + } + return remoteDigest, nil +} + func (d *DockerImage) pushWithClient(authConfig *cliTypes.AuthConfig, remoteImage string) error { ctx := context.Background() diff --git a/airflow/mocks/ImageHandler.go b/airflow/mocks/ImageHandler.go index 52e6eea4f..eedba6aac 100644 --- a/airflow/mocks/ImageHandler.go +++ b/airflow/mocks/ImageHandler.go @@ -78,6 +78,30 @@ func (_m *ImageHandler) DoesImageExist(image string) error { return r0 } +// GetImageSha provides a mock function with given fields: +func (_m *ImageHandler) GetImageSha() (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 +} + // GetLabel provides a mock function with given fields: altImageName, labelName func (_m *ImageHandler) GetLabel(altImageName string, labelName string) (string, error) { ret := _m.Called(altImageName, labelName) diff --git a/config/config.go b/config/config.go index 9d0efa9e7..db42dbabf 100644 --- a/config/config.go +++ b/config/config.go @@ -86,6 +86,7 @@ var ( DisableAstroRun: newCfg("disable_astro_run", "false"), DisableEnvObjects: newCfg("disable_env_objects", "false"), AutoSelect: newCfg("auto_select", "false"), + ShaAsTag: newCfg("sha_as_tag", "false"), } // viperHome is the viper object in the users home directory diff --git a/config/types.go b/config/types.go index 9adca30bd..ec1a03f9f 100644 --- a/config/types.go +++ b/config/types.go @@ -45,6 +45,7 @@ type cfgs struct { DisableAstroRun cfg DisableEnvObjects cfg AutoSelect cfg + ShaAsTag cfg } // Creates a new cfg struct diff --git a/software/deploy/deploy.go b/software/deploy/deploy.go index 1e091f12c..8a4000dac 100644 --- a/software/deploy/deploy.go +++ b/software/deploy/deploy.go @@ -181,11 +181,11 @@ func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Conte Output: true, Labels: deployLabels, } + err = imageHandler.Build("", "", buildConfig) if err != nil { return err } - var registry, remoteImage, token string if byoRegistryEnabled { registry = byoRegistryDomain @@ -195,13 +195,19 @@ func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Conte remoteImage = fmt.Sprintf("%s/%s", registry, airflow.ImageName(name, nextTag)) token = c.Token } - err = imageHandler.Push(remoteImage, "", token) if err != nil { return err } - if byoRegistryEnabled { + useShaAsTag := config.CFG.ShaAsTag.GetBool() + if useShaAsTag { + sha, err := imageHandler.GetImageSha() + if (sha == "") || (err != nil) { + return fmt.Errorf("failed to get image sha: %w", err) + } + remoteImage = fmt.Sprintf("%s@%s", registry, sha) + } runtimeVersion, _ := imageHandler.GetLabel("", runtimeImageLabel) airflowVersion, _ := imageHandler.GetLabel("", airflowImageLabel) req := houston.UpdateDeploymentImageRequest{ReleaseName: name, Image: remoteImage, AirflowVersion: airflowVersion, RuntimeVersion: runtimeVersion} diff --git a/software/deploy/deploy_test.go b/software/deploy/deploy_test.go index 6262dab84..982b5e24b 100644 --- a/software/deploy/deploy_test.go +++ b/software/deploy/deploy_test.go @@ -209,6 +209,38 @@ func (s *Suite) TestBuildPushDockerImageSuccessWithBYORegistry() { assert.Contains(s.T(), capturedBuildConfig.Labels, expectedLabel) mockImageHandler.AssertExpectations(s.T()) houstonMock.AssertExpectations(s.T()) + + // Case when SHA is used as tag + houstonMock.On("UpdateDeploymentImage", houston.UpdateDeploymentImageRequest{ReleaseName: "test", Image: "test.registry.io@image_sha", AirflowVersion: "1.10.12", RuntimeVersion: ""}).Return(nil, nil) + imageHandlerInit = func(image string) airflow.ImageHandler { + // Mock the Build function, capturing the buildConfig + mockImageHandler.On("Build", mock.Anything, mock.Anything, mock.MatchedBy(func(buildConfig types.ImageBuildConfig) bool { + // Capture buildConfig for later assertions + capturedBuildConfig = buildConfig + // Check if the deploy label contains the correct description + for _, label := range buildConfig.Labels { + if label == deployRevisionDescriptionLabel+"="+description { + return true + } + } + return false + })).Return(nil).Once() + + mockImageHandler.On("Push", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil) + mockImageHandler.On("GetLabel", "", runtimeImageLabel).Return("", nil).Once() + mockImageHandler.On("GetLabel", "", airflowImageLabel).Return("1.10.12", nil).Once() + mockImageHandler.On("GetImageSha", mock.Anything, mock.Anything).Return("image_sha", nil).Once() + + return mockImageHandler + } + config.CFG.ShaAsTag.SetHomeString("true") + defer config.CFG.ShaAsTag.SetHomeString("false") + err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "test.registry.io", false, true, description) + s.NoError(err) + expectedLabel = deployRevisionDescriptionLabel + "=" + description + assert.Contains(s.T(), capturedBuildConfig.Labels, expectedLabel) + mockImageHandler.AssertExpectations(s.T()) + houstonMock.AssertExpectations(s.T()) } func (s *Suite) TestBuildPushDockerImageFailure() {