Added the support for --image-name flag in the astro deploy command

cmd/software/deploy.go

@@ -21,11 +21,14 @@ var (
 
 	ignoreCacheDeploy = false
 
-	EnsureProjectDir   = utils.EnsureProjectDir
-	DeployAirflowImage = deploy.Airflow
-	DagsOnlyDeploy     = deploy.DagsOnlyDeploy
-	isDagOnlyDeploy    bool
-	description        string
+	EnsureProjectDir           = utils.EnsureProjectDir
+	DeployAirflowImage         = deploy.Airflow
+	DagsOnlyDeploy             = deploy.DagsOnlyDeploy
+	UpdateDeploymentImage      = deploy.UpdateDeploymentImage
+	isDagOnlyDeploy            bool
+	description                string
+	imageName                  string
+	runtimeVersionForImageName string
 )
 
 var deployExample = `
@@ -58,6 +61,8 @@ func NewDeployCmd() *cobra.Command {
 	cmd.Flags().BoolVarP(&ignoreCacheDeploy, "no-cache", "", false, "Do not use cache when building container image")
 	cmd.Flags().StringVar(&workspaceID, "workspace-id", "", "workspace assigned to deployment")
 	cmd.Flags().StringVar(&description, "description", "", "Improve traceability by attaching a description to a code deploy. If you don't provide a description, the system automatically assigns a default description based on the deploy type.")
+	cmd.Flags().StringVar(&imageName, "image-name", "", "Name of the image to deploy. Example - quay.io/astronomer/astro-runtime:12.1.1 or quay.io/astronomer/astro-runtime@sha256:77d85a8b6cd3bdfbcb25cbdb6ffbdb4ac2c1bb390b5951b836451dbdefb6c325")
+	cmd.Flags().StringVar(&runtimeVersionForImageName, "runtime-version", "", "Runtime version of the image to deploy. Example - 12.1.1. Mandatory if image-name is provided")
 
 	if !context.IsCloudContext() && houston.VerifyVersionMatch(houstonVersion, houston.VersionRestrictions{GTE: "0.34.0"}) {
 		cmd.Flags().BoolVarP(&isDagOnlyDeploy, "dags", "d", false, "Push only DAGs to your Deployment")
@@ -112,6 +117,9 @@ func deployAirflow(cmd *cobra.Command, args []string) error {
 		return DagsOnlyDeploy(houstonClient, appConfig, ws, deploymentID, config.WorkingPath, nil, true, description)
 	}
 
+	if imageName != "" || runtimeVersionForImageName != "" {
+		return UpdateDeploymentImage(houstonClient, deploymentID, ws, runtimeVersionForImageName, imageName)
+	}
 	// Since we prompt the user to enter the deploymentID in come cases for DeployAirflowImage, reusing the same  deploymentID for DagsOnlyDeploy
 	deploymentID, err = DeployAirflowImage(houstonClient, config.WorkingPath, deploymentID, ws, byoRegistryDomain, ignoreCacheDeploy, byoRegistryEnabled, forcePrompt, description)
 	if err != nil {

software/deploy/deploy.go

@@ -44,6 +44,8 @@ var (
 	ErrDagOnlyDeployNotEnabledForDeployment = errors.New("to perform this operation, first set the Deployment type to 'dag_deploy' via the UI or the API or the CLI")
 	ErrEmptyDagFolderUserCancelledOperation = errors.New("no DAGs found in the dags folder. User canceled the operation")
 	ErrBYORegistryDomainNotSet              = errors.New("Custom registry host is not set in config. It can be set at astronomer.houston.config.deployments.registry.protectedCustomRegistry.updateRegistry.host") //nolint
+	ErrNoRuntimeVersionPassed               = errors.New("no runtime version was provided")
+	ErrNoImageNamePassed                    = errors.New("no image name was provided")
 )
 
 const (
@@ -204,14 +206,18 @@ func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Conte
 	if byoRegistryEnabled {
 		runtimeVersion, _ := imageHandler.GetLabel("", runtimeImageLabel)
 		airflowVersion, _ := imageHandler.GetLabel("", airflowImageLabel)
-		req := houston.UpdateDeploymentImageRequest{ReleaseName: name, Image: remoteImage, AirflowVersion: airflowVersion, RuntimeVersion: runtimeVersion}
-		_, err := houston.Call(houstonClient.UpdateDeploymentImage)(req)
+		_, err := updateDeploymentImageAPICall(houstonClient, remoteImage, name, airflowVersion, runtimeVersion)
 		return err
 	}
 
 	return nil
 }
 
+func updateDeploymentImageAPICall(houstonClient houston.ClientInterface, imageName, releaseName, airflowVersion, runtimeVersion string) (interface{}, error) {
+	req := houston.UpdateDeploymentImageRequest{ReleaseName: releaseName, Image: imageName, AirflowVersion: airflowVersion, RuntimeVersion: runtimeVersion}
+	return houston.Call(houstonClient.UpdateDeploymentImage)(req)
+}
+
 func validAirflowImageRepo(image string) bool {
 	validDockerfileBaseImages := map[string]bool{
 		"quay.io/astronomer/ap-airflow": true,
@@ -427,3 +433,27 @@ func DagsOnlyDeploy(houstonClient houston.ClientInterface, appConfig *houston.Ap
 	}
 	return fileutil.UploadFile(&uploadFileArgs)
 }
+
+func UpdateDeploymentImage(houstonClient houston.ClientInterface, deploymentID, wsID, runtimeVersion, imageName string) error {
+	if runtimeVersion == "" {
+		return ErrNoRuntimeVersionPassed
+	}
+	if imageName == "" {
+		return ErrNoImageNamePassed
+	}
+	deploymentIDForCurrentCmd, _, err := getDeploymentIDForCurrentCommandVar(houstonClient, wsID, deploymentID, deploymentID == "")
+	if err != nil {
+		return err
+	}
+	deploymentID = deploymentIDForCurrentCmd
+	if deploymentID == "" {
+		return errInvalidDeploymentID
+	}
+	deploymentInfo, err := houston.Call(houstonClient.GetDeployment)(deploymentID)
+	if err != nil {
+		return fmt.Errorf("failed to get deployment info: %w", err)
+	}
+	_, err = updateDeploymentImageAPICall(houstonClient, imageName, deploymentInfo.ReleaseName, "", runtimeVersion)
+	fmt.Println("Image successfully updated")
+	return err
+}

software/deploy/deploy_test.go

@@ -841,3 +841,83 @@ func (s *Suite) TestDeployDagsOnlyFailure() {
 		s.True(os.IsNotExist(err))
 	})
 }
+
+func (s *Suite) TestUpdateDeploymentImage() {
+	testUtil.InitTestConfig(testUtil.SoftwarePlatform)
+	deploymentID := "test-deployment-id"
+	wsID := "test-workspace-id"
+	runtimeVersion := "12.1.1"
+	imageName := "imageName"
+	releaseName := "releaseName"
+
+	s.Run("When runtimeVersion is empty", func() {
+		houstonMock := new(houston_mocks.ClientInterface)
+		err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, "", imageName)
+		s.ErrorIs(err, ErrNoRuntimeVersionPassed)
+		houstonMock.AssertExpectations(s.T())
+	})
+
+	s.Run("When imageName is empty", func() {
+		houstonMock := new(houston_mocks.ClientInterface)
+		err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, "")
+		s.ErrorIs(err, ErrNoImageNamePassed)
+		houstonMock.AssertExpectations(s.T())
+	})
+
+	s.Run("When getDeploymentIDForCurrentCommandVar gives an error", func() {
+		getDeploymentIDForCurrentCommandVar = func(houstonClient houston.ClientInterface, wsID, deploymentID string, prompt bool) (string, []houston.Deployment, error) {
+			return deploymentID, nil, errDeploymentNotFound
+		}
+		houstonMock := new(houston_mocks.ClientInterface)
+		err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName)
+		s.ErrorIs(err, errDeploymentNotFound)
+		houstonMock.AssertExpectations(s.T())
+	})
+
+	s.Run("When an error occurs in the GetDeployment api call", func() {
+		getDeploymentIDForCurrentCommandVar = func(houstonClient houston.ClientInterface, wsID, deploymentID string, prompt bool) (string, []houston.Deployment, error) {
+			return deploymentID, nil, nil
+		}
+		houstonMock := new(houston_mocks.ClientInterface)
+		houstonMock.On("GetDeployment", mock.Anything).Return(nil, errMockHouston).Once()
+
+		err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName)
+		s.ErrorContains(err, "failed to get deployment info: some houston error")
+		houstonMock.AssertExpectations(s.T())
+	})
+
+	s.Run("Houston API call throws error", func() {
+		getDeploymentIDForCurrentCommandVar = func(houstonClient houston.ClientInterface, wsID, deploymentID string, prompt bool) (string, []houston.Deployment, error) {
+			return deploymentID, nil, nil
+		}
+		houstonMock := new(houston_mocks.ClientInterface)
+		deployment := &houston.Deployment{
+			ReleaseName: releaseName,
+		}
+		houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once()
+		houstonMock.On("UpdateDeploymentImage", mock.Anything).Return(nil, errMockHouston).Once()
+		err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName)
+		s.ErrorContains(err, "some houston error")
+		houstonMock.AssertExpectations(s.T())
+	})
+
+	s.Run("Successful API call", func() {
+		getDeploymentIDForCurrentCommandVar = func(houstonClient houston.ClientInterface, wsID, deploymentID string, prompt bool) (string, []houston.Deployment, error) {
+			return deploymentID, nil, nil
+		}
+		houstonMock := new(houston_mocks.ClientInterface)
+		updateDeploymentImageResp := &houston.UpdateDeploymentImageResp{
+			ReleaseName:    releaseName,
+			AirflowVersion: "",
+			RuntimeVersion: runtimeVersion,
+		}
+		deployment := &houston.Deployment{
+			ReleaseName: releaseName,
+		}
+		houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once()
+		houstonMock.On("UpdateDeploymentImage", mock.Anything).Return(updateDeploymentImageResp, nil).Once()
+		err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName)
+		houstonMock.AssertExpectations(s.T())
+		s.ErrorIs(err, nil)
+	})
+}