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

cmd/software/deploy.go

@@ -21,13 +21,18 @@ var (
 
 	ignoreCacheDeploy = false
 
-	EnsureProjectDir               = utils.EnsureProjectDir
-	DeployAirflowImage             = deploy.Airflow
-	DagsOnlyDeploy                 = deploy.DagsOnlyDeploy
-	isDagOnlyDeploy                bool
-	description                    string
-	isImageOnlyDeploy              bool
-	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 = `
@@ -61,6 +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(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")
@@ -119,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)
-	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

@@ -27,7 +27,7 @@ func (s *Suite) TestDeploy() {
 	EnsureProjectDir = func(cmd *cobra.Command, args []string) error {
 		return nil
 	}
-	DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (string, error) {
+	DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) {
 		if description == "" {
 			return deploymentID, fmt.Errorf("description should not be empty")
 		}
@@ -52,7 +52,7 @@ func (s *Suite) TestDeploy() {
 	s.NoError(err)
 
 	// Test when the default description is used
-	DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (string, error) {
+	DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) {
 		expectedDesc := "Deployed via <astro deploy>"
 		if description != expectedDesc {
 			return deploymentID, fmt.Errorf("expected description to be '%s', but got '%s'", expectedDesc, description)
@@ -67,14 +67,14 @@ func (s *Suite) TestDeploy() {
 	DagsOnlyDeploy = deploy.DagsOnlyDeploy
 
 	s.Run("error should be returned for astro deploy, if DeployAirflowImage throws error", func() {
-		DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (string, error) {
+		DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) {
 			return deploymentID, deploy.ErrNoWorkspaceID
 		}
 
 		err := execDeployCmd([]string{"-f"}...)
 		s.ErrorIs(err, deploy.ErrNoWorkspaceID)
 
-		DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (string, error) {
+		DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) {
 			return deploymentID, nil
 		}
 	})
@@ -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() {
@@ -104,15 +103,15 @@ func (s *Suite) TestDeploy() {
 	})
 
 	s.Run("Test for the flag --image for image deployment", func() {
-		DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (string, error) {
+		DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) {
 			return deploymentID, deploy.ErrDeploymentTypeIncorrectForImageOnly
 		}
 		err := execDeployCmd([]string{"test-deployment-id", "--image", "--force"}...)
 		s.ErrorIs(err, deploy.ErrDeploymentTypeIncorrectForImageOnly)
 	})
 
 	s.Run("Test for the flag --image for dags-only deployment", func() {
-		DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (string, error) {
+		DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) {
 			return deploymentID, nil
 		}
 		// This function is not called since --image is passed
@@ -123,6 +122,61 @@ func (s *Suite) TestDeploy() {
 		s.ErrorIs(err, nil)
 	})
 
+	s.Run("Test for the flag --image-name", func() {
+		var capturedImageName string
+		DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) {
+			capturedImageName = imageName // Capture the imageName
+			return deploymentID, nil
+		}
+		DagsOnlyDeploy = func(houstonClient houston.ClientInterface, appConfig *houston.AppConfig, wsID, deploymentID, dagsParentPath string, dagDeployURL *string, cleanUpFiles bool, description string) error {
+			return nil
+		}
+		testImageName := "test-image-name" // Set the expected image name
+		err := execDeployCmd([]string{"test-deployment-id", "--image-name=" + testImageName, "--force", "--workspace-id=" + mockWorkspace.ID}...)
+
+		s.ErrorIs(err, nil)
+		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: "",
@@ -135,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))
 }