From 23bd1720806fec15f9b8021d1f4b3ce313c48293 Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Wed, 11 Dec 2024 19:17:19 +0530 Subject: [PATCH 1/2] Added --image flag for astro deploy command for software --- cmd/software/deploy.go | 23 ++++++--- cmd/software/deploy_test.go | 33 +++++++++++-- software/deploy/deploy.go | 11 ++++- software/deploy/deploy_test.go | 88 ++++++++++++++++++++++++++++++---- 4 files changed, 134 insertions(+), 21 deletions(-) diff --git a/cmd/software/deploy.go b/cmd/software/deploy.go index 4bfb82698..878f31612 100644 --- a/cmd/software/deploy.go +++ b/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,23 @@ 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 { + return nil + } err = DagsOnlyDeploy(houstonClient, appConfig, ws, deploymentID, config.WorkingPath, nil, true, description) // Don't throw the error if dag-deploy itself is disabled diff --git a/cmd/software/deploy_test.go b/cmd/software/deploy_test.go index 43d592f2e..816ba88ec 100644 --- a/cmd/software/deploy_test.go +++ b/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 " 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: "", diff --git a/software/deploy/deploy.go b/software/deploy/deploy.go index 8a4000dac..501ed98e8 100644 --- a/software/deploy/deploy.go +++ b/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 diff --git a/software/deploy/deploy_test.go b/software/deploy/deploy_test.go index 982b5e24b..3f1777f2e 100644 --- a/software/deploy/deploy_test.go +++ b/software/deploy/deploy_test.go @@ -323,14 +323,14 @@ 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()) @@ -338,7 +338,7 @@ func (s *Suite) TestAirflowFailure() { 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" From aeb716b1c47eae32f181763fe8800478071f9ca4 Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Thu, 12 Dec 2024 11:20:09 +0530 Subject: [PATCH 2/2] Added print statement - Dags in the project will not be deployed since --image is passed. --- cmd/software/deploy.go | 1 + 1 file changed, 1 insertion(+) diff --git a/cmd/software/deploy.go b/cmd/software/deploy.go index 878f31612..b305f2979 100644 --- a/cmd/software/deploy.go +++ b/cmd/software/deploy.go @@ -126,6 +126,7 @@ func deployAirflow(cmd *cobra.Command, args []string) error { } // 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 }