add workload identity support

CONTRIBUTING.md

@@ -66,6 +66,17 @@ Ensure you have the following installed:
    export ASTRO_API_TOKEN=<your-api-token>
    ```
 
+### Setting up the Import script for Local Development
+
+1. Build the import script from the import directory
+   ```
+   go build import_script.go
+   ```
+2. Run the import script
+   ```
+   ./import_script -resources deployment -organizationId <your-org-id> -host dev -token YOU_API_TOKEN
+   ```
+
 ## Making Changes
 
 1. Create a new branch for your changes:

docs/resources/deployment.md

@@ -31,6 +31,7 @@ resource "astro_deployment" "dedicated" {
   resource_quota_memory          = "20Gi"
   scheduler_size                 = "SMALL"
   workspace_id                   = "clnp86ly5000401ndaga21g81"
+  desired_workload_identity      = "arn:aws:iam::123456789:role/AirflowS3Logs-clmk2qqia000008mhff3ndjr0"
   environment_variables = [{
     key       = "key1"
     value     = "value1"
@@ -164,6 +165,7 @@ resource "astro_deployment" "imported_deployment" {
 - `cluster_id` (String) Deployment cluster identifier - required for 'HYBRID' and 'DEDICATED' deployments. If changing this value, the deployment will be recreated in the new cluster
 - `default_task_pod_cpu` (String) Deployment default task pod CPU - required for 'STANDARD' and 'DEDICATED' deployments
 - `default_task_pod_memory` (String) Deployment default task pod memory - required for 'STANDARD' and 'DEDICATED' deployments
+- `desired_workload_identity` (String) Deployment's desired workload identity. The Terraform provider will use this provided workload identity to create the Deployment. If it is not provided the workload identity will be assigned automatically.
 - `is_development_mode` (Boolean) Deployment development mode - required for 'STANDARD' and 'DEDICATED' deployments. If changing from 'False' to 'True', the deployment will be recreated
 - `is_high_availability` (Boolean) Deployment high availability - required for 'STANDARD' and 'DEDICATED' deployments
 - `original_astro_runtime_version` (String) Deployment's original Astro Runtime version. The Terraform provider will use this provided Astro runtime version to create the Deployment. The Astro runtime version can be updated with your Astro project Dockerfile, but if this value is changed, the Deployment will be recreated with this new Astro runtime version.

examples/resources/astro_deployment/resource.tf

@@ -16,6 +16,7 @@ resource "astro_deployment" "dedicated" {
   resource_quota_memory          = "20Gi"
   scheduler_size                 = "SMALL"
   workspace_id                   = "clnp86ly5000401ndaga21g81"
+  desired_workload_identity      = "arn:aws:iam::123456789:role/AirflowS3Logs-clmk2qqia000008mhff3ndjr0"
   environment_variables = [{
     key       = "key1"
     value     = "value1"

import/import_script.go

@@ -824,6 +824,12 @@ func generateDeploymentHCL(ctx context.Context, platformClient *platform.ClientW
 
 		deploymentType := deployment.Type
 
+		workloadIdentity := deployment.WorkloadIdentity
+		workloadIdentityString := ""
+		if workloadIdentity != nil {
+			workloadIdentityString = fmt.Sprintf(`desired_workload_identity = "%s"`, *workloadIdentity)
+		}
+
 		if *deploymentType == platform.DeploymentTypeDEDICATED {
 			deploymentHCL = fmt.Sprintf(`
 resource "astro_deployment" "deployment_%s" {
@@ -845,6 +851,7 @@ resource "astro_deployment" "deployment_%s" {
 	type = "%s"
 	workspace_id = "%s"
 	%s
+    %s
 }
 `,
 				deployment.Id,
@@ -866,6 +873,7 @@ resource "astro_deployment" "deployment_%s" {
 				stringValue((*string)(deploymentType)),
 				deployment.WorkspaceId,
 				workerQueuesString,
+				workloadIdentityString,
 			)
 		} else if *deploymentType == platform.DeploymentTypeSTANDARD {
 			deploymentHCL = fmt.Sprintf(`
@@ -889,6 +897,7 @@ resource "astro_deployment" "deployment_%s" {
 	type = "%s"
 	workspace_id = "%s"
 	%s
+    %s
 }
 `,
 				deployment.Id,
@@ -911,6 +920,7 @@ resource "astro_deployment" "deployment_%s" {
 				stringValue((*string)(deploymentType)),
 				deployment.WorkspaceId,
 				workerQueuesString,
+				workloadIdentityString,
 			)
 		} else {
 			log.Printf("Skipping deployment %s: unsupported deployment type %s", deployment.Id, stringValue((*string)(deploymentType)))

internal/provider/models/deployment.go

@@ -49,6 +49,7 @@ type DeploymentResource struct {
 	DesiredDagTarballVersion    types.String `tfsdk:"desired_dag_tarball_version"`
 	IsCicdEnforced              types.Bool   `tfsdk:"is_cicd_enforced"`
 	IsDagDeployEnabled          types.Bool   `tfsdk:"is_dag_deploy_enabled"`
+	DesiredWorkloadIdentity     types.String `tfsdk:"desired_workload_identity"`
 	WorkloadIdentity            types.String `tfsdk:"workload_identity"`
 	ExternalIps                 types.Set    `tfsdk:"external_ips"`
 	OidcIssuerUrl               types.String `tfsdk:"oidc_issuer_url"`

internal/provider/resources/resource_deployment.go

@@ -106,6 +106,8 @@ func (r *DeploymentResource) Create(
 		}
 	}
 
+	desiredWorkloadIdentity := data.DesiredWorkloadIdentity.ValueString()
+
 	switch data.Type.ValueString() {
 	case string(platform.DeploymentTypeSTANDARD):
 		createStandardDeploymentRequest := platform.CreateStandardDeploymentRequest{
@@ -127,6 +129,9 @@ func (r *DeploymentResource) Create(
 			Type:                 platform.CreateStandardDeploymentRequestTypeSTANDARD,
 			WorkspaceId:          data.WorkspaceId.ValueString(),
 		}
+		if desiredWorkloadIdentity != "" {
+			createStandardDeploymentRequest.WorkloadIdentity = &desiredWorkloadIdentity
+		}
 
 		// contact emails
 		contactEmails, diags := utils.TypesSetToStringSlice(ctx, data.ContactEmails)
@@ -187,6 +192,9 @@ func (r *DeploymentResource) Create(
 			Type:                 platform.CreateDedicatedDeploymentRequestTypeDEDICATED,
 			WorkspaceId:          data.WorkspaceId.ValueString(),
 		}
+		if desiredWorkloadIdentity != "" {
+			createDedicatedDeploymentRequest.WorkloadIdentity = &desiredWorkloadIdentity
+		}
 
 		// contact emails
 		contactEmails, diags := utils.TypesSetToStringSlice(ctx, data.ContactEmails)
@@ -246,6 +254,10 @@ func (r *DeploymentResource) Create(
 			WorkspaceId:       data.WorkspaceId.ValueString(),
 		}
 
+		if desiredWorkloadIdentity != "" {
+			createHybridDeploymentRequest.WorkloadIdentity = &desiredWorkloadIdentity
+		}
+
 		// contact emails
 		contactEmails, diags := utils.TypesSetToStringSlice(ctx, data.ContactEmails)
 		createHybridDeploymentRequest.ContactEmails = &contactEmails
@@ -386,6 +398,8 @@ func (r *DeploymentResource) Update(
 	var updateDeploymentRequest platform.UpdateDeploymentRequest
 	var envVars []platform.DeploymentEnvironmentVariableRequest
 
+	desiredWorkloadIdentity := data.DesiredWorkloadIdentity.ValueString()
+
 	switch data.Type.ValueString() {
 	case string(platform.DeploymentTypeSTANDARD):
 		updateStandardDeploymentRequest := platform.UpdateStandardDeploymentRequest{
@@ -405,6 +419,10 @@ func (r *DeploymentResource) Update(
 			WorkspaceId:          data.WorkspaceId.ValueString(),
 		}
 
+		if desiredWorkloadIdentity != "" {
+			updateStandardDeploymentRequest.WorkloadIdentity = &desiredWorkloadIdentity
+		}
+
 		// contact emails
 		contactEmails, diags := utils.TypesSetToStringSlice(ctx, data.ContactEmails)
 		updateStandardDeploymentRequest.ContactEmails = &contactEmails
@@ -463,6 +481,10 @@ func (r *DeploymentResource) Update(
 			WorkspaceId:          data.WorkspaceId.ValueString(),
 		}
 
+		if desiredWorkloadIdentity != "" {
+			updateDedicatedDeploymentRequest.WorkloadIdentity = &desiredWorkloadIdentity
+		}
+
 		// contact emails
 		contactEmails, diags := utils.TypesSetToStringSlice(ctx, data.ContactEmails)
 		updateDedicatedDeploymentRequest.ContactEmails = &contactEmails
@@ -519,6 +541,10 @@ func (r *DeploymentResource) Update(
 			WorkspaceId:       data.WorkspaceId.ValueString(),
 		}
 
+		if desiredWorkloadIdentity != "" {
+			updateHybridDeploymentRequest.WorkloadIdentity = &desiredWorkloadIdentity
+		}
+
 		// contact emails
 		contactEmails, diags := utils.TypesSetToStringSlice(ctx, data.ContactEmails)
 		updateHybridDeploymentRequest.ContactEmails = &contactEmails

internal/provider/resources/resource_deployment_test.go

@@ -84,13 +84,15 @@ func TestAcc_ResourceDeploymentHybrid(t *testing.T) {
 					IncludeEnvironmentVariables: false,
 					SchedulerAu:                 6,
 					NodePoolId:                  nodePoolId,
+					DesiredWorkloadIdentity:     "arn:aws:iam::123456789:role/AirflowS3Logs-clmk2qqia000008mhff3ndjr0",
 				}),
 				Check: resource.ComposeTestCheckFunc(
 					resource.TestCheckResourceAttr(resourceVar, "description", utils.TestResourceDescription),
 					resource.TestCheckResourceAttr(resourceVar, "worker_queues.0.name", "default"),
 					resource.TestCheckResourceAttr(resourceVar, "environment_variables.#", "0"),
 					resource.TestCheckResourceAttr(resourceVar, "executor", "CELERY"),
 					resource.TestCheckResourceAttr(resourceVar, "scheduler_au", "6"),
+					resource.TestCheckResourceAttr(resourceVar, "workload_identity", "arn:aws:iam::123456789:role/AirflowS3Logs-clmk2qqia000008mhff3ndjr0"),
 					// Check via API that deployment exists
 					testAccCheckDeploymentExistence(t, deploymentName, false, true),
 				),
@@ -194,13 +196,15 @@ func TestAcc_ResourceDeploymentStandard(t *testing.T) {
 					SchedulerSize:               string(platform.SchedulerMachineNameEXTRALARGE),
 					IncludeEnvironmentVariables: false,
 					WorkerQueuesStr:             workerQueuesStr(""),
+					DesiredWorkloadIdentity:     "arn:aws:iam::123456789:role/AirflowS3Logs-clmk2qqia000008mhff3ndjr0",
 				}),
 				Check: resource.ComposeTestCheckFunc(
 					resource.TestCheckResourceAttr(awsResourceVar, "description", utils.TestResourceDescription),
 					resource.TestCheckResourceAttr(awsResourceVar, "scheduler_size", string(platform.SchedulerMachineNameEXTRALARGE)),
 					resource.TestCheckResourceAttr(awsResourceVar, "worker_queues.0.name", "default"),
 					resource.TestCheckNoResourceAttr(awsResourceVar, "environment_variables.0.key"),
 					resource.TestCheckResourceAttr(awsResourceVar, "executor", "CELERY"),
+					resource.TestCheckResourceAttr(awsResourceVar, "workload_identity", "arn:aws:iam::123456789:role/AirflowS3Logs-clmk2qqia000008mhff3ndjr0"),
 					// Check via API that deployment exists
 					testAccCheckDeploymentExistence(t, awsDeploymentName, true, true),
 				),
@@ -722,6 +726,7 @@ type hybridDeploymentInput struct {
 	SchedulerAu                 int
 	NodePoolId                  string
 	DuplicateWorkerQueues       bool
+	DesiredWorkloadIdentity     string
 }
 
 func hybridDeployment(input hybridDeploymentInput) string {
@@ -736,6 +741,10 @@ func hybridDeployment(input hybridDeploymentInput) string {
 	} else {
 		taskPodNodePoolIdStr = fmt.Sprintf(`task_pod_node_pool_id = "%v"`, input.NodePoolId)
 	}
+	desiredWorkloadIdentityStr := ""
+	if input.DesiredWorkloadIdentity != "" {
+		desiredWorkloadIdentityStr = input.DesiredWorkloadIdentity
+	}
 
 	return fmt.Sprintf(`
 resource "astro_workspace" "%v_workspace" {
@@ -758,12 +767,13 @@ resource "astro_deployment" "%v" {
 	%v
 	%v
 	%v
+    %v
   }
 `,
 		input.Name, input.Name, utils.TestResourceDescription,
 		input.Name, input.Name, utils.TestResourceDescription,
 		input.ClusterId, input.Executor, input.SchedulerAu, input.Name,
-		envVarsStr(input.IncludeEnvironmentVariables), wqStr, taskPodNodePoolIdStr)
+		envVarsStr(input.IncludeEnvironmentVariables), wqStr, taskPodNodePoolIdStr, desiredWorkloadIdentityStr)
 }
 
 func developmentDeployment(scalingSpecDeploymentName, scalingSpec string) string {
@@ -791,6 +801,7 @@ type standardDeploymentInput struct {
 	IsDevelopmentMode           bool
 	ScalingSpec                 string
 	WorkerQueuesStr             string
+	DesiredWorkloadIdentity     string
 }
 
 func standardDeployment(input standardDeploymentInput) string {
@@ -816,6 +827,10 @@ func standardDeployment(input standardDeploymentInput) string {
 			scalingSpecStr = input.ScalingSpec
 		}
 	}
+	desiredWorkloadIdentityStr := ""
+	if input.DesiredWorkloadIdentity != "" {
+		desiredWorkloadIdentityStr = input.DesiredWorkloadIdentity
+	}
 	return fmt.Sprintf(`
 resource "astro_workspace" "%v_workspace" {
 	name = "%s"
@@ -844,10 +859,11 @@ resource "astro_deployment" "%v" {
 	%v
 	%v
     %v
+    %v
 }
 `,
 		input.Name, input.Name, utils.TestResourceDescription, input.Name, input.Name, input.Description, input.Region, input.CloudProvider, input.Executor, input.IsDevelopmentMode, input.SchedulerSize, input.Name,
-		envVarsStr(input.IncludeEnvironmentVariables), input.WorkerQueuesStr, scalingSpecStr)
+		envVarsStr(input.IncludeEnvironmentVariables), input.WorkerQueuesStr, scalingSpecStr, desiredWorkloadIdentityStr)
 }
 
 func standardDeploymentWithVariableName(input standardDeploymentInput) string {

internal/provider/schemas/deployment.go

@@ -195,6 +195,10 @@ func DeploymentResourceSchemaAttributes() map[string]resourceSchema.Attribute {
 				stringplanmodifier.UseStateForUnknown(),
 			},
 		},
+		"desired_workload_identity": resourceSchema.StringAttribute{
+			MarkdownDescription: "Deployment's desired workload identity. The Terraform provider will use this provided workload identity to create the Deployment. If it is not provided the workload identity will be assigned automatically.",
+			Optional:            true,
+		},
 		"workload_identity": resourceSchema.StringAttribute{
 			MarkdownDescription: "Deployment workload identity. This value can be changed via the Astro API if applicable.",
 			Computed:            true,