Add workload identity to hosted deployment (#1708)

Co-authored-by: kushalmalani <[email protected]>
astronomer · Aug 30, 2024 · 64b6c42 · 64b6c42
1 parent 1008c5c
commit 64b6c42
Show file tree

Hide file tree

Showing 8 changed files with 289 additions and 96 deletions.
diff --git a/astro-client-core/api.gen.go b/astro-client-core/api.gen.go
diff --git a/cloud/deployment/deployment.go b/cloud/deployment/deployment.go
@@ -63,12 +63,11 @@ const (
 )
 
 var (
-	sleepTime                  = 180
-	tickNum                    = 10
-	timeoutNum                 = 180
-	listLimit                  = 1000
-	dedicatedDeploymentRequest = astroplatformcore.UpdateDedicatedDeploymentRequest{}
-	dagDeployEnabled           bool
+	sleepTime        = 180
+	tickNum          = 10
+	timeoutNum       = 180
+	listLimit        = 1000
+	dagDeployEnabled bool
 )
 
 func newTableOut() *printutil.Table {
@@ -135,6 +134,7 @@ func List(ws string, fromAllWorkspaces bool, platformCoreClient astroplatformcor
 	return nil
 }
 
+// TODO (https://github.com/astronomer/astro-cli/issues/1709): move these input arguments to a struct, and drop the nolint
 func Logs(deploymentID, ws, deploymentName, keyword string, logWebserver, logScheduler, logTriggerer, logWorkers, warnLogs, errorLogs, infoLogs bool, logCount int, platformCoreClient astroplatformcore.CoreClient, coreClient astrocore.CoreClient) error {
 	var logLevel string
 	var i int
@@ -212,6 +212,7 @@ func Logs(deploymentID, ws, deploymentName, keyword string, logWebserver, logSch
 	return nil
 }
 
+// TODO (https://github.com/astronomer/astro-cli/issues/1709): move these input arguments to a struct, and drop the nolint
 func Create(name, workspaceID, description, clusterID, runtimeVersion, dagDeploy, executor, cloudProvider, region, schedulerSize, highAvailability, developmentMode, cicdEnforcement, defaultTaskPodCpu, defaultTaskPodMemory, resourceQuotaCpu, resourceQuotaMemory, workloadIdentity string, deploymentType astroplatformcore.DeploymentType, schedulerAU, schedulerReplicas int, platformCoreClient astroplatformcore.CoreClient, coreClient astrocore.CoreClient, waitForStatus bool) error { //nolint
 	var organizationID string
 	var currentWorkspace astrocore.Workspace
@@ -325,6 +326,10 @@ func Create(name, workspaceID, description, clusterID, runtimeVersion, dagDeploy
 		if resourceQuotaMemory == "" {
 			resourceQuotaMemory = configOption.ResourceQuotas.ResourceQuota.Memory.Default
 		}
+		var deplWorkloadIdentity *string
+		if workloadIdentity != "" {
+			deplWorkloadIdentity = &workloadIdentity
+		}
 		// build standard input
 		if IsDeploymentStandard(deploymentType) {
 			var requestedCloudProvider astroplatformcore.CreateStandardDeploymentRequestCloudProvider
@@ -360,6 +365,7 @@ func Create(name, workspaceID, description, clusterID, runtimeVersion, dagDeploy
 				DefaultTaskPodMemory: defaultTaskPodMemory,
 				ResourceQuotaCpu:     resourceQuotaCpu,
 				ResourceQuotaMemory:  resourceQuotaMemory,
+				WorkloadIdentity:     deplWorkloadIdentity,
 			}
 			if strings.EqualFold(executor, CeleryExecutor) || strings.EqualFold(executor, CELERY) {
 				standardDeploymentRequest.WorkerQueues = &defautWorkerQueue
@@ -409,6 +415,7 @@ func Create(name, workspaceID, description, clusterID, runtimeVersion, dagDeploy
 				DefaultTaskPodMemory: defaultTaskPodMemory,
 				ResourceQuotaCpu:     resourceQuotaCpu,
 				ResourceQuotaMemory:  resourceQuotaMemory,
+				WorkloadIdentity:     deplWorkloadIdentity,
 			}
 			if strings.EqualFold(executor, CeleryExecutor) || strings.EqualFold(executor, CELERY) {
 				dedicatedDeploymentRequest.WorkerQueues = &defautWorkerQueue
@@ -738,6 +745,7 @@ func HealthPoll(deploymentID, ws string, sleepTime, tickNum, timeoutNum int, pla
 	}
 }
 
+// TODO (https://github.com/astronomer/astro-cli/issues/1709): move these input arguments to a struct, and drop the nolint
 func Update(deploymentID, name, ws, description, deploymentName, dagDeploy, executor, schedulerSize, highAvailability, developmentMode, cicdEnforcement, defaultTaskPodCpu, defaultTaskPodMemory, resourceQuotaCpu, resourceQuotaMemory, workloadIdentity string, schedulerAU, schedulerReplicas int, wQueueList []astroplatformcore.WorkerQueueRequest, hybridQueueList []astroplatformcore.HybridWorkerQueueRequest, newEnvironmentVariables []astroplatformcore.DeploymentEnvironmentVariableRequest, force bool, coreClient astrocore.CoreClient, platformCoreClient astroplatformcore.CoreClient) error { //nolint
 	var queueCreateUpdate, confirmWithUser bool
 	// get deployment
@@ -911,6 +919,10 @@ func Update(deploymentID, name, ws, description, deploymentName, dagDeploy, exec
 		if resourceQuotaMemory == "" {
 			resourceQuotaMemory = *currentDeployment.ResourceQuotaMemory
 		}
+		var deplWorkloadIdentity *string
+		if workloadIdentity != "" {
+			deplWorkloadIdentity = &workloadIdentity
+		}
 		if IsDeploymentStandard(*currentDeployment.Type) {
 			var requestedExecutor astroplatformcore.UpdateStandardDeploymentRequestExecutor
 			switch strings.ToUpper(executor) {
@@ -941,6 +953,7 @@ func Update(deploymentID, name, ws, description, deploymentName, dagDeploy, exec
 				EnvironmentVariables: deploymentEnvironmentVariablesRequest,
 				DefaultTaskPodCpu:    defaultTaskPodCpu,
 				DefaultTaskPodMemory: defaultTaskPodMemory,
+				WorkloadIdentity:     deplWorkloadIdentity,
 			}
 			switch schedulerSize {
 			case strings.ToLower(string(astrocore.CreateStandardDeploymentRequestSchedulerSizeSMALL)):
@@ -988,7 +1001,7 @@ func Update(deploymentID, name, ws, description, deploymentName, dagDeploy, exec
 			case strings.ToUpper(KUBERNETES):
 				requestedExecutor = astroplatformcore.UpdateDedicatedDeploymentRequestExecutorKUBERNETES
 			}
-			dedicatedDeploymentRequest = astroplatformcore.UpdateDedicatedDeploymentRequest{
+			dedicatedDeploymentRequest := astroplatformcore.UpdateDedicatedDeploymentRequest{
 				Description:          &description,
 				Name:                 name,
 				Executor:             requestedExecutor,
@@ -1004,6 +1017,7 @@ func Update(deploymentID, name, ws, description, deploymentName, dagDeploy, exec
 				ResourceQuotaMemory:  resourceQuotaMemory,
 				EnvironmentVariables: deploymentEnvironmentVariablesRequest,
 				WorkerQueues:         &workerQueuesRequest,
+				WorkloadIdentity:     deplWorkloadIdentity,
 			}
 			switch schedulerSize {
 			case strings.ToLower(string(astrocore.CreateStandardDeploymentRequestSchedulerSizeSMALL)):