diff --git a/cmd/software/deploy.go b/cmd/software/deploy.go index b305f2979..ed98bd413 100644 --- a/cmd/software/deploy.go +++ b/cmd/software/deploy.go @@ -27,6 +27,7 @@ var ( 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") ) @@ -61,6 +62,7 @@ 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") 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") @@ -120,7 +122,7 @@ func deployAirflow(cmd *cobra.Command, args []string) error { } // 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) + deploymentID, err = DeployAirflowImage(houstonClient, config.WorkingPath, deploymentID, ws, byoRegistryDomain, ignoreCacheDeploy, byoRegistryEnabled, forcePrompt, description, isImageOnlyDeploy, imageName) if err != nil { return err } diff --git a/cmd/software/deploy_test.go b/cmd/software/deploy_test.go index 816ba88ec..2ae169904 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, 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 " 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 } }) @@ -104,7 +104,7 @@ 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"}...) @@ -112,7 +112,7 @@ func (s *Suite) TestDeploy() { }) 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 +123,22 @@ 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("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 501ed98e8..c46ab8fbd 100644 --- a/software/deploy/deploy.go +++ b/software/deploy/deploy.go @@ -45,6 +45,8 @@ var ( 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") ) const ( @@ -58,9 +60,10 @@ const ( warningInvalidNameTag = "WARNING! You are about to push an image using the '%s' tag. This is not recommended.\nPlease use one of the following tags: %s.\nAre you sure you want to continue?" warningInvalidNameTagEmptyRecommendations = "WARNING! You are about to push an image using the '%s' tag. This is not recommended.\nAre you sure you want to continue?" - registryDomainPrefix = "registry." - runtimeImageLabel = "io.astronomer.docker.runtime.version" - airflowImageLabel = "io.astronomer.docker.airflow.version" + registryDomainPrefix = "registry." + runtimeImageLabel = "io.astronomer.docker.runtime.version" + airflowImageLabel = "io.astronomer.docker.airflow.version" + composeSkipImageBuildingPromptMsg = "Skipping building image since --image-name flag is used..." ) var tab = printutil.Table{ @@ -69,7 +72,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, isImageOnlyDeploy bool) (string, error) { +func Airflow(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { deploymentID, deployments, err := getDeploymentIDForCurrentCommand(houstonClient, wsID, deploymentID, prompt) if err != nil { return deploymentID, err @@ -107,7 +110,7 @@ func Airflow(houstonClient houston.ClientInterface, path, deploymentID, wsID, by fmt.Printf(houstonDeploymentPrompt, releaseName) // Build the image to deploy - err = buildPushDockerImage(houstonClient, &c, deploymentInfo, releaseName, path, nextTag, cloudDomain, byoRegistryDomain, ignoreCacheDeploy, byoRegistryEnabled, description) + err = buildPushDockerImage(houstonClient, &c, deploymentInfo, releaseName, path, nextTag, cloudDomain, byoRegistryDomain, ignoreCacheDeploy, byoRegistryEnabled, description, imageName) if err != nil { return deploymentID, err } @@ -129,24 +132,7 @@ func deploymentExists(deploymentID string, deployments []houston.Deployment) boo return false } -func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Context, deploymentInfo *houston.Deployment, name, path, nextTag, cloudDomain, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled bool, description string) error { - // Build our image - fmt.Println(imageBuildingPrompt) - - // parse dockerfile - cmds, err := docker.ParseFile(filepath.Join(path, dockerfile)) - if err != nil { - return fmt.Errorf("failed to parse dockerfile: %s: %w", filepath.Join(path, dockerfile), err) - } - - image, tag := docker.GetImageTagFromParsedFile(cmds) - if config.CFG.ShowWarnings.GetBool() && !validAirflowImageRepo(image) && !validRuntimeImageRepo(image) { - i, _ := input.Confirm(fmt.Sprintf(warningInvalidImageName, image)) - if !i { - fmt.Println("Canceling deploy...") - os.Exit(1) - } - } +func validateRuntimeVersion(houstonClient houston.ClientInterface, tag string, deploymentInfo *houston.Deployment) error { // Get valid image tags for platform using Deployment Info request deploymentConfig, err := houston.Call(houstonClient.GetDeploymentConfig)(nil) if err != nil { @@ -175,26 +161,10 @@ func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Conte os.Exit(1) } } - imageName := airflow.ImageName(name, "latest") - - imageHandler := imageHandlerInit(imageName) - - if description != "" { - deployLabels = append(deployLabels, "io.astronomer.deploy.revision.description="+description) - } - - buildConfig := types.ImageBuildConfig{ - Path: config.WorkingPath, - NoCache: ignoreCacheDeploy, - TargetPlatforms: deployImagePlatformSupport, - Output: true, - Labels: deployLabels, - } + return nil +} - err = imageHandler.Build("", "", buildConfig) - if err != nil { - return 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 { registry = byoRegistryDomain @@ -204,7 +174,7 @@ func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Conte remoteImage = fmt.Sprintf("%s/%s", registry, airflow.ImageName(name, nextTag)) token = c.Token } - err = imageHandler.Push(remoteImage, "", token) + err := imageHandler.Push(remoteImage, "", token) if err != nil { return err } @@ -223,10 +193,82 @@ func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Conte _, err := houston.Call(houstonClient.UpdateDeploymentImage)(req) return err } - return nil } +func buildDockerImageForCustomImage(imageHandler airflow.ImageHandler, customImageName string, deploymentInfo *houston.Deployment, houstonClient houston.ClientInterface) error { + fmt.Println(composeSkipImageBuildingPromptMsg) + err := imageHandler.TagLocalImage(customImageName) + if err != nil { + return err + } + runtimeLabel, err := imageHandler.GetLabel("", airflow.RuntimeImageLabel) + if err != nil { + fmt.Println("unable get runtime version from image") + return err + } + if runtimeLabel == "" { + return ErrNoRuntimeLabelOnCustomImage + } + err = validateRuntimeVersion(houstonClient, runtimeLabel, deploymentInfo) + return err +} + +func buildDockerImageForWorkingDir(path string, imageHandler airflow.ImageHandler, houstonClient houston.ClientInterface, deploymentInfo *houston.Deployment, ignoreCacheDeploy bool) error { + // all these checks inside Dockerfile should happen only when no image-name is provided + // parse dockerfile + cmds, err := docker.ParseFile(filepath.Join(path, dockerfile)) + if err != nil { + return fmt.Errorf("failed to parse dockerfile: %s: %w", filepath.Join(path, dockerfile), err) + } + + image, tag := docker.GetImageTagFromParsedFile(cmds) + if config.CFG.ShowWarnings.GetBool() && !validAirflowImageRepo(image) && !validRuntimeImageRepo(image) { + i, _ := input.Confirm(fmt.Sprintf(warningInvalidImageName, image)) + if !i { + fmt.Println("Canceling deploy...") + os.Exit(1) + } + } + // Get valid image tags for platform using Deployment Info request + err = validateRuntimeVersion(houstonClient, tag, deploymentInfo) + if err != nil { + return err + } + // Build our image + fmt.Println(imageBuildingPrompt) + buildConfig := types.ImageBuildConfig{ + Path: config.WorkingPath, + NoCache: ignoreCacheDeploy, + TargetPlatforms: deployImagePlatformSupport, + Output: true, + Labels: deployLabels, + } + + err = imageHandler.Build("", "", buildConfig) + return err +} + +func buildDockerImage(ignoreCacheDeploy bool, deploymentInfo *houston.Deployment, customImageName, path string, imageHandler airflow.ImageHandler, houstonClient houston.ClientInterface) error { + if customImageName == "" { + return buildDockerImageForWorkingDir(path, imageHandler, houstonClient, deploymentInfo, ignoreCacheDeploy) + } + return buildDockerImageForCustomImage(imageHandler, customImageName, deploymentInfo, houstonClient) +} + +func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Context, deploymentInfo *houston.Deployment, name, path, nextTag, cloudDomain, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled bool, description, customImageName string) error { + imageName := airflow.ImageName(name, "latest") + imageHandler := imageHandlerInit(imageName) + if description != "" { + deployLabels = append(deployLabels, "io.astronomer.deploy.revision.description="+description) + } + err := buildDockerImage(ignoreCacheDeploy, deploymentInfo, customImageName, path, imageHandler, houstonClient) + if err != nil { + return err + } + return pushDockerImage(byoRegistryEnabled, byoRegistryDomain, name, nextTag, cloudDomain, imageHandler, houstonClient, c) +} + func validAirflowImageRepo(image string) bool { validDockerfileBaseImages := map[string]bool{ "quay.io/astronomer/ap-airflow": true, diff --git a/software/deploy/deploy_test.go b/software/deploy/deploy_test.go index 3f1777f2e..c299b8843 100644 --- a/software/deploy/deploy_test.go +++ b/software/deploy/deploy_test.go @@ -124,7 +124,7 @@ func (s *Suite) TestBuildPushDockerImageSuccessWithTagWarning() { houstonMock.On("GetRuntimeReleases", "").Return(houston.RuntimeReleases{}, nil) houstonMock.On("GetDeploymentConfig", nil).Return(mockedDeploymentConfig, nil) - err := buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description) + err := buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description, "") s.NoError(err) mockImageHandler.AssertExpectations(s.T()) houstonMock.AssertExpectations(s.T()) @@ -155,7 +155,7 @@ func (s *Suite) TestBuildPushDockerImageSuccessWithImageRepoWarning() { houstonMock.On("GetDeploymentConfig", nil).Return(mockedDeploymentConfig, nil) houstonMock.On("GetRuntimeReleases", "").Return(houston.RuntimeReleases{}, nil) - err := buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description) + err := buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description, "") s.NoError(err) mockImageHandler.AssertExpectations(s.T()) houstonMock.AssertExpectations(s.T()) @@ -202,7 +202,7 @@ func (s *Suite) TestBuildPushDockerImageSuccessWithBYORegistry() { houstonMock.On("GetRuntimeReleases", "").Return(houston.RuntimeReleases{}, nil) houstonMock.On("UpdateDeploymentImage", houston.UpdateDeploymentImageRequest{ReleaseName: "test", Image: "test.registry.io:test-test", AirflowVersion: "1.10.12", RuntimeVersion: ""}).Return(nil, nil) - err := buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "test.registry.io", false, true, description) + err := buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "test.registry.io", false, true, description, "") s.NoError(err) expectedLabel := deployRevisionDescriptionLabel + "=" + description @@ -235,7 +235,7 @@ func (s *Suite) TestBuildPushDockerImageSuccessWithBYORegistry() { } config.CFG.ShaAsTag.SetHomeString("true") defer config.CFG.ShaAsTag.SetHomeString("false") - err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "test.registry.io", false, true, description) + err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "test.registry.io", false, true, description, "") s.NoError(err) expectedLabel = deployRevisionDescriptionLabel + "=" + description assert.Contains(s.T(), capturedBuildConfig.Labels, expectedLabel) @@ -246,7 +246,7 @@ func (s *Suite) TestBuildPushDockerImageSuccessWithBYORegistry() { func (s *Suite) TestBuildPushDockerImageFailure() { // invalid dockerfile test dockerfile = "Dockerfile.invalid" - err := buildPushDockerImage(nil, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description) + err := buildPushDockerImage(nil, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description, "") s.EqualError(err, "failed to parse dockerfile: testfiles/Dockerfile.invalid: when using JSON array syntax, arrays must be comprised of strings only") dockerfile = "Dockerfile" @@ -262,7 +262,7 @@ func (s *Suite) TestBuildPushDockerImageFailure() { houstonMock.On("GetDeploymentConfig", nil).Return(nil, errMockHouston).Once() houstonMock.On("GetRuntimeReleases", "").Return(houston.RuntimeReleases{}, nil) // houston GetDeploymentConfig call failure - err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description) + err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description, "") s.Error(err, errMockHouston) houstonMock.On("GetDeploymentConfig", nil).Return(mockedDeploymentConfig, nil).Twice() @@ -274,7 +274,7 @@ func (s *Suite) TestBuildPushDockerImageFailure() { } // build error test case - err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description) + err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description, "") s.Error(err, errSomeContainerIssue.Error()) mockImageHandler.AssertExpectations(s.T()) @@ -286,7 +286,7 @@ func (s *Suite) TestBuildPushDockerImageFailure() { } // push error test case - err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description) + err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description, "") s.Error(err, errSomeContainerIssue.Error()) mockImageHandler.AssertExpectations(s.T()) houstonMock.AssertExpectations(s.T()) @@ -323,14 +323,14 @@ func (s *Suite) TestGetAirflowUILinkFailure() { func (s *Suite) TestAirflowFailure() { // No workspace ID test case - _, err := Airflow(nil, "", "", "", "", false, false, false, description, false) + _, 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, false) + _, 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, false) + _, 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, false) + _, 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, false) + _, 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, false) + _, 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, false) + _, 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, false) + _, 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, false) + _, 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,7 +413,7 @@ 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, false) + _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, false, "") s.NoError(err) houstonMock.AssertExpectations(s.T()) } @@ -452,9 +452,84 @@ func (s *Suite) TestAirflowSuccessForImageOnly() { 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) + _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, "") s.NoError(err) houstonMock.AssertExpectations(s.T()) + mockImageHandler.AssertExpectations(s.T()) +} + +func (s *Suite) TestAirflowSuccessForImageName() { + fs := afero.NewMemMapFs() + configYaml := testUtil.NewTestConfig("localhost") + afero.WriteFile(fs, config.HomeConfigFile, configYaml, 0o777) + config.InitConfig(fs) + + customImageName := "test-image-name" + mockImageHandler := new(mocks.ImageHandler) + imageHandlerInit = func(image string) airflow.ImageHandler { + mockImageHandler.On("Push", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil) + mockImageHandler.On("TagLocalImage", customImageName).Return(nil) + mockImageHandler.On("GetLabel", "", airflow.RuntimeImageLabel).Return("test", 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, customImageName) + s.NoError(err) + houstonMock.AssertExpectations(s.T()) + mockImageHandler.AssertExpectations(s.T()) +} + +func (s *Suite) TestAirflowFailForImageNameWhenImageHasNoRuntimeLabel() { + fs := afero.NewMemMapFs() + configYaml := testUtil.NewTestConfig("localhost") + afero.WriteFile(fs, config.HomeConfigFile, configYaml, 0o777) + config.InitConfig(fs) + + customImageName := "test-image-name" + mockImageHandler := new(mocks.ImageHandler) + imageHandlerInit = func(image string) airflow.ImageHandler { + mockImageHandler.On("TagLocalImage", customImageName).Return(nil) + mockImageHandler.On("GetLabel", "", airflow.RuntimeImageLabel).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: "dag-only", + } + 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, customImageName) + s.Error(err, ErrNoRuntimeLabelOnCustomImage) + houstonMock.AssertExpectations(s.T()) + mockImageHandler.AssertExpectations(s.T()) } func (s *Suite) TestAirflowFailureForImageOnly() { @@ -481,9 +556,10 @@ func (s *Suite) TestAirflowFailureForImageOnly() { houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() - _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true) + _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, "") s.Error(err, ErrDeploymentTypeIncorrectForImageOnly) houstonMock.AssertExpectations(s.T()) + mockImageHandler.AssertExpectations(s.T()) } func (s *Suite) TestDeployDagsOnlyFailure() {