Added --image-name flag for software, similar to astro

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
 	}

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
 		}
 	})
@@ -104,15 +104,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 +123,17 @@ func (s *Suite) TestDeploy() {
 		s.ErrorIs(err, nil)
 	})
 
+	s.Run("Test for the flag --image-name", func() {
+		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
+		}
+		DagsOnlyDeploy = func(houstonClient houston.ClientInterface, appConfig *houston.AppConfig, wsID, deploymentID, dagsParentPath string, dagDeployURL *string, cleanUpFiles bool, description string) error {
+			return nil
+		}
+		err := execDeployCmd([]string{"test-deployment-id", "--image-name", "--force", "--workspace-id=" + mockWorkspace.ID}...)
+		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

@@ -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,72 @@ func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Conte
 			os.Exit(1)
 		}
 	}
-	imageName := airflow.ImageName(name, "latest")
+	return nil
+}
 
+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)
 	}
 
-	buildConfig := types.ImageBuildConfig{
-		Path:            config.WorkingPath,
-		NoCache:         ignoreCacheDeploy,
-		TargetPlatforms: deployImagePlatformSupport,
-		Output:          true,
-		Labels:          deployLabels,
-	}
+	var err error
+	if customImageName == "" {
+		// 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)
+		}
 
-	err = imageHandler.Build("", "", buildConfig)
-	if err != nil {
-		return 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)
+		if err != nil {
+			return err
+		}
+	} else {
+		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)
+		if err != nil {
+			return err
+		}
 	}
+
 	var registry, remoteImage, token string
 	if byoRegistryEnabled {
 		registry = byoRegistryDomain