Skip to content

Commit

Permalink
Added --image flag for astro deploy command for software (#1753)
Browse files Browse the repository at this point in the history
  • Loading branch information
@rujhan-arora-astronomer
rujhan-arora-astronomer authored Dec 12, 2024
1 parent 19672e5 commit dce037b
Show file tree
Hide file tree
Showing 4 changed files with 135 additions and 21 deletions.
24 changes: 18 additions & 6 deletions cmd/software/deploy.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 = `
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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
Expand Down
33 changes: 29 additions & 4 deletions cmd/software/deploy_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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")
}
Expand All @@ -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)
Expand All @@ -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
}
})
Expand All @@ -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: "",
Expand Down
11 changes: 10 additions & 1 deletion software/deploy/deploy.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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
Expand All @@ -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
Expand Down
88 changes: 78 additions & 10 deletions software/deploy/deploy_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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())

Expand All @@ -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")
Expand Down Expand Up @@ -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"
Expand Down

0 comments on commit dce037b

Please sign in to comment.