Added --image flag for astro deploy command for software

cmd/software/deploy.go

@@ -21,11 +21,13 @@ 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
+	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")
 )
 
 var deployExample = `
@@ -58,6 +60,7 @@ 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().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.")
 
 	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")
@@ -108,15 +111,24 @@ func deployAirflow(cmd *cobra.Command, args []string) error {
 		description = utils.GetDefaultDeployDescription(isDagOnlyDeploy)
 	}
 
+	if isImageOnlyDeploy && isDagOnlyDeploy {
+		return ErrBothDagsOnlyAndImageOnlySet
+	}
+
 	if isDagOnlyDeploy {
 		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)
+	deploymentID, err = DeployAirflowImage(houstonClient, config.WorkingPath, deploymentID, ws, byoRegistryDomain, ignoreCacheDeploy, byoRegistryEnabled, forcePrompt, description, isImageOnlyDeploy)
 	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.")
+		return nil
+	}
 
 	err = DagsOnlyDeploy(houstonClient, appConfig, ws, deploymentID, config.WorkingPath, nil, true, description)
 	// Don't throw the error if dag-deploy itself is disabled

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) (string, error) {
+	DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (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) (string, error) {
+	DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (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) (string, error) {
+		DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (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) (string, error) {
+		DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (string, error) {
 			return deploymentID, nil
 		}
 	})
@@ -98,6 +98,31 @@ func (s *Suite) TestDeploy() {
 		s.ErrorIs(err, deploy.ErrDagOnlyDeployDisabledInConfig)
 	})
 
+	s.Run("Test when both the flags --dags and --image are passed", func() {
+		err := execDeployCmd([]string{"test-deployment-id", "--dags", "--image", "--force"}...)
+		s.ErrorIs(err, ErrBothDagsOnlyAndImageOnlySet)
+	})
+
+	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) {
+			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) {
+			return deploymentID, nil
+		}
+		// This function is not called since --image is passed
+		DagsOnlyDeploy = func(houstonClient houston.ClientInterface, appConfig *houston.AppConfig, wsID, deploymentID, dagsParentPath string, dagDeployURL *string, cleanUpFiles bool, description string) error {
+			return deploy.ErrNoWorkspaceID
+		}
+		err := execDeployCmd([]string{"test-deployment-id", "--image", "--force"}...)
+		s.ErrorIs(err, nil)
+	})
+
 	s.Run("error should be returned if BYORegistryEnabled is true but BYORegistryDomain is empty", func() {
 		appConfig = &houston.AppConfig{
 			BYORegistryDomain: "",

software/deploy/deploy.go

@@ -44,6 +44,7 @@ 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
+	ErrDeploymentTypeIncorrectForImageOnly  = errors.New("--image only works for Dag-only, Git-sync-based and NFS-based deployments")
 )
 
 const (
@@ -68,7 +69,7 @@ var tab = printutil.Table{
 	Header:         []string{"#", "LABEL", "DEPLOYMENT NAME", "WORKSPACE", "DEPLOYMENT ID"},
 }
 
-func Airflow(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string) (string, error) {
+func Airflow(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (string, error) {
 	deploymentID, deployments, err := getDeploymentIDForCurrentCommand(houstonClient, wsID, deploymentID, prompt)
 	if err != nil {
 		return deploymentID, err
@@ -95,6 +96,14 @@ func Airflow(houstonClient houston.ClientInterface, path, deploymentID, wsID, by
 		return deploymentID, fmt.Errorf("failed to get deployment info: %w", err)
 	}
 
+	// isImageOnlyDeploy is not valid for image-based deployments since image-based deployments inherently mean that the image itself contains dags.
+	// If we deploy only the image, the deployment will not have any dags for image-based deployments.
+	// Even on astro, image-based deployments are not allowed to be deployed with --image flag.
+	if isImageOnlyDeploy && deploymentInfo.DagDeployment.Type == houston.ImageDeploymentType {
+		return "", ErrDeploymentTypeIncorrectForImageOnly
+	}
+	// We don't need to exclude the dags from the image because the dags present in the image are not respected anyways for non-image based deployments
+
 	fmt.Printf(houstonDeploymentPrompt, releaseName)
 
 	// Build the image to deploy

software/deploy/deploy_test.go

@@ -323,22 +323,22 @@ func (s *Suite) TestGetAirflowUILinkFailure() {
 
 func (s *Suite) TestAirflowFailure() {
 	// No workspace ID test case
-	_, err := Airflow(nil, "", "", "", "", false, false, false, description)
+	_, err := Airflow(nil, "", "", "", "", false, false, false, description, false)
 	s.ErrorIs(err, ErrNoWorkspaceID)
 
 	// houston GetWorkspace failure case
 	houstonMock := new(houston_mocks.ClientInterface)
 	houstonMock.On("GetWorkspace", mock.Anything).Return(nil, errMockHouston).Once()
 
-	_, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description)
+	_, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description, false)
 	s.ErrorIs(err, errMockHouston)
 	houstonMock.AssertExpectations(s.T())
 
 	// houston ListDeployments failure case
 	houstonMock.On("GetWorkspace", mock.Anything).Return(&houston.Workspace{}, nil)
 	houstonMock.On("ListDeployments", mock.Anything).Return(nil, errMockHouston).Once()
 
-	_, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description)
+	_, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description, false)
 	s.ErrorIs(err, errMockHouston)
 	houstonMock.AssertExpectations(s.T())
 
@@ -352,35 +352,35 @@ func (s *Suite) TestAirflowFailure() {
 	// config GetCurrentContext failure case
 	config.ResetCurrentContext()
 
-	_, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description)
+	_, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description, false)
 	s.EqualError(err, "no context set, have you authenticated to Astro or Astronomer Software? Run astro login and try again")
 
 	context.Switch("localhost")
 
 	// Invalid deployment name case
-	_, err = Airflow(houstonMock, "", "test-deployment-id", "test-workspace-id", "", false, false, false, description)
+	_, err = Airflow(houstonMock, "", "test-deployment-id", "test-workspace-id", "", false, false, false, description, false)
 	s.ErrorIs(err, errInvalidDeploymentID)
 
 	// No deployment in the current workspace case
-	_, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description)
+	_, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description, false)
 	s.ErrorIs(err, errDeploymentNotFound)
 	houstonMock.AssertExpectations(s.T())
 
 	// Invalid deployment selection case
 	houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil)
-	_, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description)
+	_, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description, false)
 	s.ErrorIs(err, errInvalidDeploymentSelected)
 
 	// return error When houston get deployment throws an error
 	houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil)
 	houstonMock.On("GetDeployment", mock.Anything).Return(nil, errMockHouston).Once()
-	_, err = Airflow(houstonMock, "", "test-deployment-id", "test-workspace-id", "", false, false, false, description)
+	_, err = Airflow(houstonMock, "", "test-deployment-id", "test-workspace-id", "", false, false, false, description, false)
 	s.Equal(err.Error(), "failed to get deployment info: "+errMockHouston.Error())
 
 	// buildPushDockerImage failure case
 	houstonMock.On("GetDeployment", "test-deployment-id").Return(&houston.Deployment{}, nil)
 	dockerfile = "Dockerfile.invalid"
-	_, err = Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description)
+	_, err = Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, false)
 	dockerfile = "Dockerfile"
 	s.Error(err)
 	s.Contains(err.Error(), "failed to parse dockerfile")
@@ -413,11 +413,79 @@ func (s *Suite) TestAirflowSuccess() {
 	houstonMock.On("GetDeployment", mock.Anything).Return(&houston.Deployment{}, nil).Once()
 	houstonMock.On("GetRuntimeReleases", "").Return(mockRuntimeReleases, nil)
 
-	_, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description)
+	_, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, false)
 	s.NoError(err)
 	houstonMock.AssertExpectations(s.T())
 }
 
+func (s *Suite) TestAirflowSuccessForImageOnly() {
+	fs := afero.NewMemMapFs()
+	configYaml := testUtil.NewTestConfig("localhost")
+	afero.WriteFile(fs, config.HomeConfigFile, configYaml, 0o777)
+	config.InitConfig(fs)
+
+	mockImageHandler := new(mocks.ImageHandler)
+	imageHandlerInit = func(image string) airflow.ImageHandler {
+		mockImageHandler.On("Build", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+		mockImageHandler.On("Push", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+		return mockImageHandler
+	}
+
+	mockedDeploymentConfig := &houston.DeploymentConfig{
+		AirflowImages: mockAirflowImageList,
+	}
+	mockRuntimeReleases := houston.RuntimeReleases{
+		houston.RuntimeRelease{Version: "4.2.4", AirflowVersion: "2.2.5"},
+		houston.RuntimeRelease{Version: "4.2.5", AirflowVersion: "2.2.5"},
+	}
+	houstonMock := new(houston_mocks.ClientInterface)
+	houstonMock.On("GetWorkspace", mock.Anything).Return(&houston.Workspace{}, nil).Once()
+	houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil).Once()
+	houstonMock.On("GetDeploymentConfig", nil).Return(mockedDeploymentConfig, nil).Once()
+	dagDeployment := &houston.DagDeploymentConfig{
+		Type: "dag-only",
+	}
+	deployment := &houston.Deployment{
+		DagDeployment: *dagDeployment,
+	}
+
+	houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once()
+	houstonMock.On("GetRuntimeReleases", "").Return(mockRuntimeReleases, nil)
+
+	_, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true)
+	s.NoError(err)
+	houstonMock.AssertExpectations(s.T())
+}
+
+func (s *Suite) TestAirflowFailureForImageOnly() {
+	fs := afero.NewMemMapFs()
+	configYaml := testUtil.NewTestConfig("localhost")
+	afero.WriteFile(fs, config.HomeConfigFile, configYaml, 0o777)
+	config.InitConfig(fs)
+
+	mockImageHandler := new(mocks.ImageHandler)
+	imageHandlerInit = func(image string) airflow.ImageHandler {
+		mockImageHandler.On("Build", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+		mockImageHandler.On("Push", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+		return mockImageHandler
+	}
+	houstonMock := new(houston_mocks.ClientInterface)
+	houstonMock.On("GetWorkspace", mock.Anything).Return(&houston.Workspace{}, nil).Once()
+	houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil).Once()
+	dagDeployment := &houston.DagDeploymentConfig{
+		Type: "image",
+	}
+	deployment := &houston.Deployment{
+		DagDeployment: *dagDeployment,
+	}
+
+	houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once()
+
+	_, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true)
+	s.Error(err, ErrDeploymentTypeIncorrectForImageOnly)
+	houstonMock.AssertExpectations(s.T())
+}
+
 func (s *Suite) TestDeployDagsOnlyFailure() {
 	testUtil.InitTestConfig(testUtil.SoftwarePlatform)
 	deploymentID := "test-deployment-id"