Merge branch 'main' into feature/minimize-default-dev-output

cmd/root.go

@@ -1,30 +1,22 @@
 package cmd
 
 import (
-	"errors"
 	"fmt"
-	"net/http"
 	"os"
-	"strings"
-	"time"
-
-	"github.com/astronomer/astro-cli/cmd/registry"
-	"github.com/sirupsen/logrus"
 
 	airflowclient "github.com/astronomer/astro-cli/airflow-client"
 	astrocore "github.com/astronomer/astro-cli/astro-client-core"
 	astroiamcore "github.com/astronomer/astro-cli/astro-client-iam-core"
 	astroplatformcore "github.com/astronomer/astro-cli/astro-client-platform-core"
 	cloudCmd "github.com/astronomer/astro-cli/cmd/cloud"
+	"github.com/astronomer/astro-cli/cmd/registry"
 	softwareCmd "github.com/astronomer/astro-cli/cmd/software"
-	"github.com/astronomer/astro-cli/config"
+	"github.com/astronomer/astro-cli/cmd/utils"
 	"github.com/astronomer/astro-cli/context"
 	"github.com/astronomer/astro-cli/houston"
 	"github.com/astronomer/astro-cli/pkg/ansi"
 	"github.com/astronomer/astro-cli/pkg/httputil"
-	"github.com/astronomer/astro-cli/version"
-
-	"github.com/google/go-github/v48/github"
+	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 )
 
@@ -73,37 +65,10 @@ func NewRootCmd() *cobra.Command {
     \__\/\__\/ \_____\/   \__\/    \_\/ \_\/ \_____\/           \_____\/ \_____\/\________\/
 
 Welcome to the Astro CLI, the modern command line interface for data orchestration. You can use it for Astro, Astronomer Software, or Local Development.`,
-		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
-			// Check for latest version
-			if config.CFG.UpgradeMessage.GetBool() {
-				// create github client with 3 second timeout, setting an aggressive timeout since its not mandatory to get a response in each command execution
-				githubClient := github.NewClient(&http.Client{Timeout: 3 * time.Second})
-				// compare current version to latest
-				err = version.CompareVersions(githubClient, "astronomer", "astro-cli")
-				if err != nil {
-					softwareCmd.InitDebugLogs = append(softwareCmd.InitDebugLogs, "Error comparing CLI versions: "+err.Error())
-				}
-			}
-			if isCloudCtx {
-				err = cloudCmd.Setup(cmd, platformCoreClient, astroCoreClient)
-				if err != nil {
-					if strings.Contains(err.Error(), "token is invalid or malformed") {
-						return errors.New("API Token is invalid or malformed") //nolint
-					}
-					if strings.Contains(err.Error(), "the API token given has expired") {
-						return errors.New("API Token is expired") //nolint
-					}
-					softwareCmd.InitDebugLogs = append(softwareCmd.InitDebugLogs, "Error during cmd setup: "+err.Error())
-				}
-			}
-			// common PersistentPreRunE component between software & cloud
-			// setting up log verbosity and dumping debug logs collected during CLI-initialization
-			if err := softwareCmd.SetUpLogs(os.Stdout, verboseLevel); err != nil {
-				return err
-			}
-			softwareCmd.PrintDebugLogs()
-			return nil
-		},
+		PersistentPreRunE: utils.ChainRunEs(
+			SetupLoggingPersistentPreRunE,
+			CreateRootPersistentPreRunE(astroCoreClient, platformCoreClient),
+		),
 	}
 
 	rootCmd.AddCommand(

cmd/root_hooks.go

@@ -0,0 +1,56 @@
+package cmd
+
+import (
+	"errors"
+	"net/http"
+	"os"
+	"strings"
+	"time"
+
+	astrocore "github.com/astronomer/astro-cli/astro-client-core"
+	astroplatformcore "github.com/astronomer/astro-cli/astro-client-platform-core"
+	cloudCmd "github.com/astronomer/astro-cli/cmd/cloud"
+	softwareCmd "github.com/astronomer/astro-cli/cmd/software"
+	"github.com/astronomer/astro-cli/config"
+	"github.com/astronomer/astro-cli/context"
+	"github.com/astronomer/astro-cli/version"
+	"github.com/google/go-github/v48/github"
+	"github.com/spf13/cobra"
+)
+
+// SetupLoggingPersistentPreRunE is a pre-run hook shared between software & cloud
+// setting up log verbosity.
+func SetupLoggingPersistentPreRunE(_ *cobra.Command, _ []string) error {
+	return softwareCmd.SetUpLogs(os.Stdout, verboseLevel)
+}
+
+// CreateRootPersistentPreRunE takes clients as arguments and returns a cobra
+// pre-run hook that sets up the context and checks for the latest version.
+func CreateRootPersistentPreRunE(astroCoreClient astrocore.CoreClient, platformCoreClient astroplatformcore.CoreClient) func(cmd *cobra.Command, args []string) error {
+	return func(cmd *cobra.Command, args []string) error {
+		// Check for latest version
+		if config.CFG.UpgradeMessage.GetBool() {
+			// create github client with 3 second timeout, setting an aggressive timeout since its not mandatory to get a response in each command execution
+			githubClient := github.NewClient(&http.Client{Timeout: 3 * time.Second})
+			// compare current version to latest
+			err := version.CompareVersions(githubClient, "astronomer", "astro-cli")
+			if err != nil {
+				softwareCmd.InitDebugLogs = append(softwareCmd.InitDebugLogs, "Error comparing CLI versions: "+err.Error())
+			}
+		}
+		if context.IsCloudContext() {
+			err := cloudCmd.Setup(cmd, platformCoreClient, astroCoreClient)
+			if err != nil {
+				if strings.Contains(err.Error(), "token is invalid or malformed") {
+					return errors.New("API Token is invalid or malformed") //nolint
+				}
+				if strings.Contains(err.Error(), "the API token given has expired") {
+					return errors.New("API Token is expired") //nolint
+				}
+				softwareCmd.InitDebugLogs = append(softwareCmd.InitDebugLogs, "Error during cmd setup: "+err.Error())
+			}
+		}
+		softwareCmd.PrintDebugLogs()
+		return nil
+	}
+}

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,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: "",

cmd/utils/utils.go

@@ -7,6 +7,20 @@ import (
 	"github.com/spf13/cobra"
 )
 
+type RunE func(cmd *cobra.Command, args []string) error
+
+// ChainRunEs chains multiple RunE functions together for cleaner composition.
+func ChainRunEs(runEs ...RunE) RunE {
+	return func(cmd *cobra.Command, args []string) error {
+		for _, runE := range runEs {
+			if err := runE(cmd, args); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+}
+
 func EnsureProjectDir(cmd *cobra.Command, args []string) error {
 	isProjectDir, err := config.IsProjectDir(config.WorkingPath)
 	if err != nil {

cmd/utils/utils_test.go

@@ -1,6 +1,7 @@
 package utils
 
 import (
+	"errors"
 	"testing"
 
 	"github.com/astronomer/astro-cli/config"
@@ -46,3 +47,42 @@ func TestGetDefaultDeployDescription(t *testing.T) {
 	descriptionWithDags := GetDefaultDeployDescription(true)
 	assert.Equal(t, "Deployed via <astro deploy --dags>", descriptionWithDags)
 }
+
+func TestChainRunEsExecutesAllFunctionsSuccessfully(t *testing.T) {
+	runE1 := func(cmd *cobra.Command, args []string) error {
+		return nil
+	}
+	runE2 := func(cmd *cobra.Command, args []string) error {
+		return nil
+	}
+	chain := ChainRunEs(runE1, runE2)
+	err := chain(&cobra.Command{}, []string{})
+	assert.NoError(t, err)
+}
+
+func TestChainRunEsReturnsErrorIfAnyFunctionFails(t *testing.T) {
+	runE1 := func(cmd *cobra.Command, args []string) error {
+		return nil
+	}
+	runE2 := func(cmd *cobra.Command, args []string) error {
+		return errors.New("error in runE2")
+	}
+	chain := ChainRunEs(runE1, runE2)
+	err := chain(&cobra.Command{}, []string{})
+	assert.Error(t, err)
+	assert.Equal(t, "error in runE2", err.Error())
+}
+
+func TestChainRunEsStopsExecutionAfterError(t *testing.T) {
+	runE1 := func(cmd *cobra.Command, args []string) error {
+		return errors.New("error in runE1")
+	}
+	runE2 := func(cmd *cobra.Command, args []string) error {
+		t.FailNow() // This should not be called
+		return nil
+	}
+	chain := ChainRunEs(runE1, runE2)
+	err := chain(&cobra.Command{}, []string{})
+	assert.Error(t, err)
+	assert.Equal(t, "error in runE1", err.Error())
+}