Introduction of sha_as_tag=True config in astro cli which will be use…

…d during Astro deploy (#1752)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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 {

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()
 

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

config/types.go

@@ -45,6 +45,7 @@ type cfgs struct {
 	DisableAstroRun       cfg
 	DisableEnvObjects     cfg
 	AutoSelect            cfg
+	ShaAsTag              cfg
 }
 
 // Creates a new cfg struct

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}

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() {