Fix issues from bug bash (#1460)

* Fix issues from bug bash * refactor code * Update cloud/deploy/deploy.go Co-authored-by: kushalmalani <[email protected]> * fix test * fix test --------- Co-authored-by: kushalmalani <[email protected]>
astronomer · Nov 29, 2023 · 4764527 · 4764527
1 parent 1823da2
commit 4764527
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 29 deletions.
diff --git a/cloud/deploy/deploy.go b/cloud/deploy/deploy.go
@@ -70,9 +70,10 @@ var (
 )
 
 var (
-	errDagsParseFailed       = errors.New("your local DAGs did not parse. Fix the listed errors or use `astro deploy [deployment-id] -f` to force deploy") //nolint:revive
-	envFileMissing           = errors.New("Env file path is incorrect: ")                                                                                  //nolint:revive
-	errCiCdEnforcementUpdate = errors.New("cannot update dag deploy since ci/cd enforcement is enabled for this deployment. Please use API Tokens or API Keys instead")
+	errDagsParseFailed        = errors.New("your local DAGs did not parse. Fix the listed errors or use `astro deploy [deployment-id] -f` to force deploy") //nolint:revive
+	envFileMissing            = errors.New("Env file path is incorrect: ")                                                                                  //nolint:revive
+	errCiCdEnforcementUpdate  = errors.New("cannot update dag deploy since ci/cd enforcement is enabled for this deployment. Please use API Tokens or API Keys instead")
+	errImageDeployNoPriorDags = errors.New("cannot do image only deploy with no prior DAGs deployed. Please deploy DAGs to your deployment first")
 )
 
 var (
@@ -83,16 +84,17 @@ var (
 )
 
 type deploymentInfo struct {
-	deploymentID     string
-	namespace        string
-	deployImage      string
-	currentVersion   string
-	organizationID   string
-	workspaceID      string
-	webserverURL     string
-	dagDeployEnabled bool
-	deploymentType   string
-	cicdEnforcement  bool
+	deploymentID             string
+	namespace                string
+	deployImage              string
+	currentVersion           string
+	organizationID           string
+	workspaceID              string
+	webserverURL             string
+	deploymentType           string
+	desiredDagTarballVersion string
+	dagDeployEnabled         bool
+	cicdEnforcement          bool
 }
 
 type InputDeploy struct {
@@ -232,6 +234,15 @@ func Deploy(deployInput InputDeploy, client astro.Client, coreClient astrocore.C
 		return nil
 	}
 
+	if deployInput.Image {
+		if !deployInfo.dagDeployEnabled {
+			return fmt.Errorf(enableDagDeployMsg, deployInfo.deploymentID) //nolint
+		}
+		if deployInfo.desiredDagTarballVersion == "" {
+			return errImageDeployNoPriorDags
+		}
+	}
+
 	deploymentURL, err := deployment.GetDeploymentURL(deployInfo.deploymentID, deployInfo.workspaceID)
 	if err != nil {
 		return err
@@ -372,23 +383,12 @@ func Deploy(deployInput InputDeploy, client astro.Client, coreClient astrocore.C
 					return err
 				}
 			} else {
-				if !deployInfo.dagDeployEnabled {
-					return fmt.Errorf(enableDagDeployMsg, deployInfo.deploymentID) //nolint
-				}
 				fmt.Println("Image Deploy only. Skipping deploying DAG...")
 			}
 		}
 		// finish deploy
 		if deployInput.Image {
-			coreDeployment, err := deployment.CoreGetDeployment(deployInfo.workspaceID, deployInfo.organizationID, deployInfo.deploymentID, coreClient)
-			if err != nil {
-				return err
-			}
-			if coreDeployment.CurrentDagTarballVersion != nil {
-				dagTarballVersion = *coreDeployment.CurrentDagTarballVersion
-			} else {
-				dagTarballVersion = ""
-			}
+			dagTarballVersion = deployInfo.desiredDagTarballVersion
 		}
 		err = updateDeploy(deployID, deployInfo.deploymentID, deployInfo.organizationID, dagTarballVersion, deployInfo.dagDeployEnabled, coreClient)
 		if err != nil {
@@ -430,6 +430,17 @@ func getDeploymentInfo(deploymentID, wsID, deploymentName string, prompt bool, c
 		if err != nil {
 			return deploymentInfo{}, err
 		}
+		coreDeployment, err := deployment.CoreGetDeployment(currentDeployment.Workspace.ID, currentDeployment.Workspace.OrganizationID, currentDeployment.ID, coreClient)
+		if err != nil {
+			return deploymentInfo{}, err
+		}
+		var desiredDagTarballVersion string
+		if coreDeployment.DesiredDagTarballVersion != nil {
+			desiredDagTarballVersion = *coreDeployment.DesiredDagTarballVersion
+		} else {
+			desiredDagTarballVersion = ""
+		}
+
 		return deploymentInfo{
 			currentDeployment.ID,
 			currentDeployment.ReleaseName,
@@ -438,8 +449,9 @@ func getDeploymentInfo(deploymentID, wsID, deploymentName string, prompt bool, c
 			currentDeployment.Workspace.OrganizationID,
 			currentDeployment.Workspace.ID,
 			currentDeployment.DeploymentSpec.Webserver.URL,
-			currentDeployment.DagDeployEnabled,
 			currentDeployment.Type,
+			desiredDagTarballVersion,
+			currentDeployment.DagDeployEnabled,
 			currentDeployment.APIKeyOnlyDeployments,
 		}, nil
 	}

diff --git a/cloud/deploy/deploy_test.go b/cloud/deploy/deploy_test.go
@@ -30,6 +30,7 @@ var (
 	ws                    = "test-ws-id"
 	dagTarballVersionTest = "test-version"
 	dagsUploadTestURL     = "test-url"
+	deploymentID          = "test-id"
 	createDeployResponse  = astrocore.CreateDeployResponse{
 		HTTPResponse: &http.Response{
 			StatusCode: 200,
@@ -74,6 +75,20 @@ var (
 			IsDagDeployEnabled: false,
 		},
 	}
+	mockCoreDeploymentResponse = []astrocore.Deployment{
+		{
+			Id:     deploymentID,
+			Status: "HEALTHY",
+		},
+	}
+	mockListDeploymentsResponse = astrocore.ListDeploymentsResponse{
+		HTTPResponse: &http.Response{
+			StatusCode: 200,
+		},
+		JSON200: &astrocore.DeploymentsPaginated{
+			Deployments: mockCoreDeploymentResponse,
+		},
+	}
 )
 
 func TestDeployWithoutDagsDeploySuccess(t *testing.T) {
@@ -95,6 +110,7 @@ func TestDeployWithoutDagsDeploySuccess(t *testing.T) {
 	mockClient := new(astro_mocks.Client)
 
 	mockCoreClient.On("GetDeploymentWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&deploymentResponse, nil).Times(4)
+	mockCoreClient.On("ListDeploymentsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&mockListDeploymentsResponse, nil).Times(1)
 	mockClient.On("ListDeployments", org, ws).Return([]astro.Deployment{{ID: "test-id", Workspace: astro.Workspace{ID: ws}}}, nil).Once()
 	mockCoreClient.On("GetDeploymentOptionsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&getDeploymentOptionsResponse, nil).Times(5)
 	mockCoreClient.On("CreateDeployWithResponse", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&createDeployResponse, nil).Times(5)
@@ -201,7 +217,7 @@ func TestDeployOnCiCdEnforcedDeployment(t *testing.T) {
 	canCiCdDeploy = func(astroAPIToken string) bool {
 		return false
 	}
-
+	mockCoreClient.On("ListDeploymentsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&mockListDeploymentsResponse, nil).Times(1)
 	mockClient.On("ListDeployments", org, ws).Return([]astro.Deployment{{ID: "test-id", Workspace: astro.Workspace{ID: ws}, DagDeployEnabled: true, APIKeyOnlyDeployments: true}}, nil).Once()
 
 	err := Deploy(deployInput, mockClient, mockCoreClient)
@@ -237,6 +253,7 @@ func TestDeployWithDagsDeploySuccess(t *testing.T) {
 
 	mockCoreClient.On("GetDeploymentWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&deploymentResponse, nil).Times(5)
 	mockClient.On("ListDeployments", org, ws).Return([]astro.Deployment{{ID: "test-id", Workspace: astro.Workspace{ID: ws}, DagDeployEnabled: true}}, nil).Times(2)
+	mockCoreClient.On("ListDeploymentsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&mockListDeploymentsResponse, nil).Times(2)
 	mockCoreClient.On("GetDeploymentOptionsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&getDeploymentOptionsResponse, nil).Times(7)
 	mockCoreClient.On("CreateDeployWithResponse", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&createDeployResponse, nil).Times(7)
 	mockCoreClient.On("UpdateDeployWithResponse", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&updateDeployResponse, nil).Times(7)
@@ -393,6 +410,7 @@ func TestDagsDeploySuccess(t *testing.T) {
 
 	mockCoreClient.On("GetDeploymentOptionsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&getDeploymentOptionsResponse, nil).Times(3)
 	mockClient.On("ListDeployments", mock.Anything, mock.Anything).Return(mockDeplyResp, nil).Times(5)
+	mockCoreClient.On("ListDeploymentsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&mockListDeploymentsResponse, nil).Times(5)
 	mockCoreClient.On("CreateDeployWithResponse", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&createDeployResponse, nil).Times(5)
 	mockCoreClient.On("UpdateDeployWithResponse", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&updateDeployResponse, nil).Times(5)
 
@@ -488,6 +506,7 @@ func TestNoDagsDeploy(t *testing.T) {
 	}
 
 	mockClient.On("ListDeployments", mock.Anything, mock.Anything).Return(mockDeplyResp, nil).Times(1)
+	mockCoreClient.On("ListDeploymentsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&mockListDeploymentsResponse, nil).Times(1)
 	mockCoreClient.On("CreateDeployWithResponse", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&createDeployResponse, nil).Times(1)
 
 	deployInput := InputDeploy{
@@ -553,6 +572,7 @@ func TestDagsDeployFailed(t *testing.T) {
 		Dags:           true,
 	}
 	mockClient.On("ListDeployments", mock.Anything, mock.Anything).Return(mockDeplyResp, nil).Times(3)
+	mockCoreClient.On("ListDeploymentsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&mockListDeploymentsResponse, nil).Times(3)
 	mockCoreClient.On("GetDeploymentOptionsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&getDeploymentOptionsResponse, nil).Times(2)
 	mockCoreClient.On("CreateDeployWithResponse", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&createDeployResponse, nil).Times(3)
 
@@ -633,6 +653,7 @@ func TestDeployFailure(t *testing.T) {
 	testUtil.InitTestConfig(testUtil.CloudPlatform)
 	mockClient := new(astro_mocks.Client)
 	mockClient.On("ListDeployments", org, ws).Return(mockDeplyResp, nil).Times(2)
+	mockCoreClient.On("ListDeploymentsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&mockListDeploymentsResponse, nil).Times(3)
 	mockCoreClient.On("CreateDeployWithResponse", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&createDeployResponse, nil).Times(2)
 	mockCoreClient.On("GetDeploymentOptionsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&getDeploymentOptionsResponse, nil).Once()
 
@@ -740,6 +761,7 @@ func TestDeployMonitoringDAGNonHosted(t *testing.T) {
 
 	mockCoreClient.On("GetDeploymentOptionsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&getDeploymentOptionsResponse, nil).Times(3)
 	mockClient.On("ListDeployments", mock.Anything, mock.Anything).Return(mockDeplyResp, nil).Times(4)
+	mockCoreClient.On("ListDeploymentsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&mockListDeploymentsResponse, nil).Times(4)
 	mockCoreClient.On("CreateDeployWithResponse", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&createDeployResponse, nil).Times(4)
 	mockCoreClient.On("UpdateDeployWithResponse", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&updateDeployResponse, nil).Times(4)
 
@@ -844,6 +866,7 @@ func TestDeployNoMonitoringDAGHosted(t *testing.T) {
 
 	mockCoreClient.On("GetDeploymentOptionsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&getDeploymentOptionsResponse, nil).Times(3)
 	mockClient.On("ListDeployments", mock.Anything, mock.Anything).Return(mockDeplyResp, nil).Times(4)
+	mockCoreClient.On("ListDeploymentsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&mockListDeploymentsResponse, nil).Times(4)
 	mockCoreClient.On("CreateDeployWithResponse", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&createDeployResponse, nil).Times(4)
 	mockCoreClient.On("UpdateDeployWithResponse", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&updateDeployResponse, nil).Times(4)
 

diff --git a/cloud/deployment/deployment.go b/cloud/deployment/deployment.go
@@ -696,6 +696,10 @@ func Update(deploymentID, label, ws, description, deploymentName, dagDeploy, exe
 		}
 	}
 
+	// determine dagDeploy enabled/disabled
+	if dagDeploy == "" {
+		deploymentUpdate.DagDeployEnabled = currentDeployment.DagDeployEnabled
+	}
 	if dagDeploy == "enable" {
 		if currentDeployment.DagDeployEnabled {
 			fmt.Println("\nDAG deploys are already enabled for this Deployment. Your DAGs will continue to run as scheduled.")

diff --git a/cloud/deployment/deployment_test.go b/cloud/deployment/deployment_test.go
@@ -2605,7 +2605,7 @@ func TestUpdate(t *testing.T) { //nolint
 			ClusterID:        "",
 			Label:            "",
 			Description:      "",
-			DagDeployEnabled: false,
+			DagDeployEnabled: true,
 			DeploymentSpec: astro.DeploymentCreateSpec{
 				Executor:  KubeExecutor,
 				Scheduler: astro.Scheduler{AU: 5, Replicas: 3},
@@ -2678,7 +2678,7 @@ func TestUpdate(t *testing.T) { //nolint
 			ClusterID:        "",
 			Label:            "",
 			Description:      "",
-			DagDeployEnabled: false,
+			DagDeployEnabled: true,
 			DeploymentSpec: astro.DeploymentCreateSpec{
 				Executor:  CeleryExecutor,
 				Scheduler: astro.Scheduler{AU: 5, Replicas: 3},