Introduction of astro deploy --image-name=image_name --remote --runtime-version=rt_version

cmd/software/deploy.go

@@ -21,14 +21,18 @@ var (
 
 	ignoreCacheDeploy = false
 
-	EnsureProjectDir               = utils.EnsureProjectDir
-	DeployAirflowImage             = deploy.Airflow
-	DagsOnlyDeploy                 = deploy.DagsOnlyDeploy
-	isDagOnlyDeploy                bool
-	description                    string
-	isImageOnlyDeploy              bool
-	imageName                      string
-	ErrBothDagsOnlyAndImageOnlySet = errors.New("cannot use both --dags and --image together. Run 'astro deploy' to update both your image and dags")
+	EnsureProjectDir                   = utils.EnsureProjectDir
+	DeployAirflowImage                 = deploy.Airflow
+	DagsOnlyDeploy                     = deploy.DagsOnlyDeploy
+	UpdateDeploymentImage              = deploy.UpdateDeploymentImage
+	isDagOnlyDeploy                    bool
+	description                        string
+	isImageOnlyDeploy                  bool
+	imageName                          string
+	runtimeVersionForImageName         string
+	imagePresentOnRemote               bool
+	ErrBothDagsOnlyAndImageOnlySet     = errors.New("cannot use both --dags and --image together. Run 'astro deploy' to update both your image and dags")
+	ErrImageNameNotPassedForRemoteFlag = errors.New("--image-name is mandatory when --remote flag is passed")
 )
 
 var deployExample = `
@@ -62,7 +66,9 @@ func NewDeployCmd() *cobra.Command {
 	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().BoolVarP(&isImageOnlyDeploy, "image", "", false, "Push only an image to your Astro Deployment. This only works for Dag-only, Git-sync-based and NFS-based deployments.")
-	cmd.Flags().StringVarP(&imageName, "image-name", "i", "", "Name of the custom image(present locally) to deploy")
+	cmd.Flags().StringVarP(&imageName, "image-name", "i", "", "Name of the custom image(should be present locally unless --remote is specified) to deploy")
+	cmd.Flags().StringVar(&runtimeVersionForImageName, "runtime-version", "", "Runtime version of the image to deploy. Example - 12.1.1. Mandatory if --image-name --remote is provided")
+	cmd.Flags().BoolVarP(&imagePresentOnRemote, "remote", "", false, "Custom image which is present on the remote registry. Can only be used with --image-name flag")
 
 	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")
@@ -121,11 +127,22 @@ func deployAirflow(cmd *cobra.Command, args []string) error {
 		return DagsOnlyDeploy(houstonClient, appConfig, ws, deploymentID, config.WorkingPath, nil, true, description)
 	}
 
-	// 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, isImageOnlyDeploy, imageName)
-	if err != nil {
-		return err
+	if imagePresentOnRemote {
+		if imageName == "" {
+			return ErrImageNameNotPassedForRemoteFlag
+		}
+		deploymentID, err = UpdateDeploymentImage(houstonClient, deploymentID, ws, runtimeVersionForImageName, imageName)
+		if err != nil {
+			return err
+		}
+	} else {
+		// 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, isImageOnlyDeploy, imageName)
+		if err != nil {
+			return err
+		}
 	}
+
 	// Don't deploy dags even for dags-only deployments --image is passed
 	if isImageOnlyDeploy {
 		fmt.Println("Dags in the project will not be deployed since --image is passed.")

cmd/software/deploy_test.go

@@ -85,7 +85,6 @@ func (s *Suite) TestDeploy() {
 		}
 		err := execDeployCmd([]string{"-f"}...)
 		s.ErrorIs(err, deploy.ErrNoWorkspaceID)
-		DagsOnlyDeploy = deploy.DagsOnlyDeploy
 	})
 
 	s.Run("No error should be returned for astro deploy, if dags deploy throws error but the feature itself is disabled", func() {
@@ -139,6 +138,45 @@ func (s *Suite) TestDeploy() {
 		s.Equal(testImageName, capturedImageName, "The imageName passed to DeployAirflowImage is incorrect")
 	})
 
+	s.Run("Test for the flag --image-name with --remote. Dags should be deployed but DeployAirflowImage shouldn't be called", func() {
+		DagsOnlyDeploy = func(houstonClient houston.ClientInterface, appConfig *houston.AppConfig, wsID, deploymentID, dagsParentPath string, dagDeployURL *string, cleanUpFiles bool, description string) error {
+			return nil
+		}
+		// Create a flag to track if DeployAirflowImage is called
+		deployAirflowImageCalled := false
+
+		// Mock function for DeployAirflowImage
+		DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) {
+			deployAirflowImageCalled = true // Set the flag if this function is called
+			return deploymentID, nil
+		}
+		UpdateDeploymentImage = func(houstonClient houston.ClientInterface, deploymentID, wsID, runtimeVersion, imageName string) (string, error) {
+			return "", nil
+		}
+		testImageName := "test-image-name" // Set the expected image name
+		err := execDeployCmd([]string{"test-deployment-id", "--image-name=" + testImageName, "--force", "--remote", "--workspace-id=" + mockWorkspace.ID}...)
+		s.ErrorIs(err, nil)
+		// Assert that DeployAirflowImage was NOT called
+		s.False(deployAirflowImageCalled, "DeployAirflowImage should not be called when --remote is specified")
+	})
+
+	s.Run("Test for the flag --image-name with --remote. Dags should not be deployed if UpdateDeploymentImage throws an error", func() {
+		UpdateDeploymentImage = func(houstonClient houston.ClientInterface, deploymentID, wsID, runtimeVersion, imageName string) (string, error) {
+			return "", errNoWorkspaceFound
+		}
+		testImageName := "test-image-name" // Set the expected image name
+		err := execDeployCmd([]string{"test-deployment-id", "--image-name=" + testImageName, "--force", "--remote", "--workspace-id=" + mockWorkspace.ID}...)
+		s.ErrorIs(err, errNoWorkspaceFound)
+	})
+
+	s.Run("Test for the flag --remote without --image-name. It should throw an error", func() {
+		UpdateDeploymentImage = func(houstonClient houston.ClientInterface, deploymentID, wsID, runtimeVersion, imageName string) (string, error) {
+			return "", errNoWorkspaceFound
+		}
+		err := execDeployCmd([]string{"test-deployment-id", "--force", "--remote", "--workspace-id=" + mockWorkspace.ID}...)
+		s.ErrorIs(err, ErrImageNameNotPassedForRemoteFlag)
+	})
+
 	s.Run("error should be returned if BYORegistryEnabled is true but BYORegistryDomain is empty", func() {
 		appConfig = &houston.AppConfig{
 			BYORegistryDomain: "",
@@ -151,6 +189,5 @@ func (s *Suite) TestDeploy() {
 		}
 		err := execDeployCmd([]string{"-f"}...)
 		s.ErrorIs(err, deploy.ErrBYORegistryDomainNotSet)
-		DagsOnlyDeploy = deploy.DagsOnlyDeploy
 	})
 }

cmd/software/suite_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	testUtil "github.com/astronomer/astro-cli/pkg/testing"
+	"github.com/astronomer/astro-cli/software/deploy"
 	"github.com/stretchr/testify/suite"
 )
 
@@ -22,9 +23,20 @@ func (s *Suite) SetupSuite() {
 func (s *Suite) SetupTest() {
 	// Reset the version once this is torn down
 	houstonVersion = "0.34.0"
+	DagsOnlyDeploy = deploy.DagsOnlyDeploy
+}
+
+func (s *Suite) SetupSubTest() {
+	DagsOnlyDeploy = deploy.DagsOnlyDeploy
+}
+
+func (s *Suite) TearDownSuite() {
+	DagsOnlyDeploy = deploy.DagsOnlyDeploy
+	UpdateDeploymentImage = deploy.UpdateDeploymentImage
 }
 
 var (
-	_ suite.SetupAllSuite  = (*Suite)(nil)
-	_ suite.SetupTestSuite = (*Suite)(nil)
+	_ suite.SetupAllSuite    = (*Suite)(nil)
+	_ suite.SetupTestSuite   = (*Suite)(nil)
+	_ suite.TearDownAllSuite = (*Suite)(nil)
 )

houston/suite_test.go

@@ -3,13 +3,18 @@ package houston
 import (
 	"testing"
 
+	testUtil "github.com/astronomer/astro-cli/pkg/testing"
 	"github.com/stretchr/testify/suite"
 )
 
 type Suite struct {
 	suite.Suite
 }
 
+func SetupSuite() {
+	testUtil.InitTestConfig(testUtil.SoftwarePlatform)
+}
+
 func TestHouston(t *testing.T) {
 	suite.Run(t, new(Suite))
 }

software/deploy/deploy.go

@@ -33,18 +33,19 @@ var (
 )
 
 var (
-	ErrNoWorkspaceID                        = errors.New("no workspace id provided")
-	errNoDomainSet                          = errors.New("no domain set, re-authenticate")
-	errInvalidDeploymentID                  = errors.New("please specify a valid deployment ID")
-	errDeploymentNotFound                   = errors.New("no airflow deployments found")
-	errInvalidDeploymentSelected            = errors.New("invalid deployment selection\n") //nolint
-	ErrDagOnlyDeployDisabledInConfig        = errors.New("to perform this operation, set both deployments.dagOnlyDeployment and deployments.configureDagDeployment to true in your Astronomer cluster")
-	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
-	ErrDeploymentTypeIncorrectForImageOnly  = errors.New("--image only works for Dag-only, Git-sync-based and NFS-based deployments")
-	WarningInvalidImageNameMsg              = "WARNING! The image in your Dockerfile '%s' is not based on Astro Runtime and is not supported. Change your Dockerfile with an image that pulls from 'quay.io/astronomer/astro-runtime' to proceed.\n"
-	ErrNoRuntimeLabelOnCustomImage          = errors.New("the image should have label io.astronomer.docker.runtime.version")
+	ErrNoWorkspaceID                         = errors.New("no workspace id provided")
+	errNoDomainSet                           = errors.New("no domain set, re-authenticate")
+	errInvalidDeploymentID                   = errors.New("please specify a valid deployment ID")
+	errDeploymentNotFound                    = errors.New("no airflow deployments found")
+	errInvalidDeploymentSelected             = errors.New("invalid deployment selection\n") //nolint
+	ErrDagOnlyDeployDisabledInConfig         = errors.New("to perform this operation, set both deployments.dagOnlyDeployment and deployments.configureDagDeployment to true in your Astronomer cluster")
+	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
+	ErrDeploymentTypeIncorrectForImageOnly   = errors.New("--image only works for Dag-only, Git-sync-based and NFS-based deployments")
+	WarningInvalidImageNameMsg               = "WARNING! The image in your Dockerfile '%s' is not based on Astro Runtime and is not supported. Change your Dockerfile with an image that pulls from 'quay.io/astronomer/astro-runtime' to proceed.\n"
+	ErrNoRuntimeLabelOnCustomImage           = errors.New("the image should have label io.astronomer.docker.runtime.version")
+	ErrRuntimeVersionNotPassedForRemoteImage = errors.New("if --image-name and --remote is passed, it's mandatory to pass --runtime-version")
 )
 
 const (
@@ -162,6 +163,28 @@ func validateRuntimeVersion(houstonClient houston.ClientInterface, tag string, d
 	return nil
 }
 
+func UpdateDeploymentImage(houstonClient houston.ClientInterface, deploymentID, wsID, runtimeVersion, imageName string) (string, error) {
+	if runtimeVersion == "" {
+		return "", ErrRuntimeVersionNotPassedForRemoteImage
+	}
+	deploymentID, _, err := getDeploymentIDForCurrentCommandVar(houstonClient, wsID, deploymentID, deploymentID == "")
+	if err != nil {
+		return "", err
+	}
+	if deploymentID == "" {
+		return "", errInvalidDeploymentID
+	}
+	deploymentInfo, err := houston.Call(houstonClient.GetDeployment)(deploymentID)
+	if err != nil {
+		return "", fmt.Errorf("failed to get deployment info: %w", err)
+	}
+	fmt.Println("Skipping building the image since --image-name flag is used...")
+	req := houston.UpdateDeploymentImageRequest{ReleaseName: deploymentInfo.ReleaseName, Image: imageName, AirflowVersion: "", RuntimeVersion: runtimeVersion}
+	_, err = houston.Call(houstonClient.UpdateDeploymentImage)(req)
+	fmt.Println("Image successfully updated")
+	return deploymentID, err
+}
+
 func pushDockerImage(byoRegistryEnabled bool, byoRegistryDomain, name, nextTag, cloudDomain string, imageHandler airflow.ImageHandler, houstonClient houston.ClientInterface, c *config.Context) error {
 	var registry, remoteImage, token string
 	if byoRegistryEnabled {
@@ -188,7 +211,7 @@ func pushDockerImage(byoRegistryEnabled bool, byoRegistryDomain, name, nextTag,
 		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 = houston.Call(houstonClient.UpdateDeploymentImage)(req)
 		return err
 	}
 	return nil
@@ -402,11 +425,10 @@ func DagsOnlyDeploy(houstonClient houston.ClientInterface, appConfig *houston.Ap
 		return ErrDagOnlyDeployDisabledInConfig
 	}
 
-	deploymentIDForCurrentCmd, _, err := getDeploymentIDForCurrentCommandVar(houstonClient, wsID, deploymentID, deploymentID == "")
+	deploymentID, _, err := getDeploymentIDForCurrentCommandVar(houstonClient, wsID, deploymentID, deploymentID == "")
 	if err != nil {
 		return err
 	}
-	deploymentID = deploymentIDForCurrentCmd
 
 	if deploymentID == "" {
 		return errInvalidDeploymentID