add workload identity support

.github/workflows/testacc.yml

@@ -79,24 +79,24 @@ jobs:
          echo "SKIP_CLUSTER_RESOURCE_TESTS=$SKIP_TESTS" >> $GITHUB_ENV
       - env:
           TF_ACC: "1"
-          HYBRID_ORGANIZATION_API_TOKEN: ${{ secrets.DEV_HYBRID_ORGANIZATION_API_TOKEN }}
-          HYBRID_ORGANIZATION_ID: clx44v7op01nf01m5iohqjkk6
-          HOSTED_ORGANIZATION_API_TOKEN: ${{ secrets.DEV_HOSTED_ORGANIZATION_API_TOKEN }}
-          HOSTED_ORGANIZATION_ID: clx42kkcm01fo01o06agtmshg
-          HOSTED_SCIM_ORGANIZATION_API_TOKEN: ${{ secrets.DEV_HOSTED_SCIM_ORGANIZATION_API_TOKEN }}
-          HOSTED_SCIM_ORGANIZATION_ID: clz3bcmd3003m01qemptnfenp
-          HYBRID_CLUSTER_ID: clxkqfzvm001d01ncr9rs80si
-          HYBRID_DRY_RUN_CLUSTER_ID: clxko4djp008601njcuoxt4z5
-          HYBRID_NODE_POOL_ID: clxkqfzvm001c01nc1eosyxzg
-          ASTRO_API_HOST: https://api.astronomer-dev.io
+          HYBRID_ORGANIZATION_API_TOKEN: ${{ secrets.STAGE_HYBRID_ORGANIZATION_API_TOKEN }}
+          HYBRID_ORGANIZATION_ID: clx46ca4y061z01jleyku7sr6
+          HOSTED_ORGANIZATION_API_TOKEN: ${{ secrets.STAGE_HOSTED_ORGANIZATION_API_TOKEN }}
+          HOSTED_ORGANIZATION_ID: clx46acvv060e01ilddqlbsmc
+          HOSTED_SCIM_ORGANIZATION_API_TOKEN: ${{ secrets.STAGE_HOSTED_SCIM_ORGANIZATION_API_TOKEN }}
+          HOSTED_SCIM_ORGANIZATION_ID: clz3blqb500lh01mtkwu9zk5z
+          HYBRID_CLUSTER_ID: clxm3xg9e05bl01ixsrhxje4e
+          HYBRID_DRY_RUN_CLUSTER_ID: clxm3y54805bs01ix5owqhfff
+          HYBRID_NODE_POOL_ID: clxm3xg9e05bk01ixrqk52cob
+          ASTRO_API_HOST: https://api.astronomer-stage.io
+          HOSTED_TEAM_ID: clx486hno068301il306nuhsm
+          HOSTED_USER_ID: clz3a95hw00j301jj5jfmcgwd
+          HOSTED_DUMMY_USER_ID: clzawlsb701vv01ikvsqz5mws
+          HOSTED_DEPLOYMENT_ID: clx4825jb068z01j9931ib5gb
+          HOSTED_STANDARD_DEPLOYMENT_ID: cm077ee2807g301kpjkqdoc15
+          HOSTED_WORKSPACE_ID: clx480rvx068u01j9mp7t7fqh
+          HOSTED_API_TOKEN_ID: clxm46ged05b301neuucdqwox
           SKIP_CLUSTER_RESOURCE_TESTS: ${{ env.SKIP_CLUSTER_RESOURCE_TESTS }}
-          HOSTED_TEAM_ID: clx44rvzr01nc01o06pze6qb7
-          HOSTED_USER_ID: clz3a4ymt004x01on8w5ydq8j
-          HOSTED_DUMMY_USER_ID: clzawipbm00bm01qw98vzzoca
-          HOSTED_DEPLOYMENT_ID: cm1zkps2a0cv301ph39benet6
-          HOSTED_STANDARD_DEPLOYMENT_ID: cm070pg0r00wd01qgnskk0dir
-          HOSTED_WORKSPACE_ID: clx42sxw501gl01o0gjenthnh
-          HOSTED_API_TOKEN_ID: clxm4836f00ql01me3nigmcr6
           TESTARGS: "-failfast"
         run: make testacc
 

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/cluster.md

@@ -15,7 +15,7 @@ Cluster resource. If creating multiple clusters, add a delay between each cluste
 ```terraform
 resource "astro_cluster" "aws_example" {
   type             = "DEDICATED"
-  name             = "my first aws cluster"
+  name             = "LIOTTA"
   region           = "us-east-1"
   cloud_provider   = "AWS"
   vpc_subnet_range = "172.20.0.0/20"
@@ -27,43 +27,43 @@ resource "astro_cluster" "aws_example" {
   }
 }
 
-resource "astro_cluster" "azure_example" {
-  type             = "DEDICATED"
-  name             = "my first azure cluster"
-  region           = "westus2"
-  cloud_provider   = "AZURE"
-  vpc_subnet_range = "172.20.0.0/19"
-  workspace_ids    = ["clv4wcf6f003u01m3zp7gsvzg"]
-}
-
-resource "astro_cluster" "gcp_example" {
-  type                  = "DEDICATED"
-  name                  = "my first gcp cluster"
-  region                = "us-central1"
-  cloud_provider        = "GCP"
-  pod_subnet_range      = "172.21.0.0/19"
-  service_peering_range = "172.23.0.0/20"
-  service_subnet_range  = "172.22.0.0/22"
-  vpc_subnet_range      = "172.20.0.0/22"
-  workspace_ids         = []
-}
-
-// Import an existing cluster
-import {
-  id = "clozc036j01to01jrlgvuf98d" // ID of the existing cluster
-  to = astro_cluster.imported_cluster
-}
-resource "astro_cluster" "imported_cluster" {
-  type                  = "DEDICATED"
-  name                  = "an existing cluster to import"
-  region                = "us-central1"
-  cloud_provider        = "GCP"
-  pod_subnet_range      = "172.21.0.0/19"
-  service_peering_range = "172.23.0.0/20"
-  service_subnet_range  = "172.22.0.0/22"
-  vpc_subnet_range      = "172.20.0.0/22"
-  workspace_ids         = []
-}
+# resource "astro_cluster" "azure_example" {
+#   type             = "DEDICATED"
+#   name             = "my first azure cluster"
+#   region           = "westus2"
+#   cloud_provider   = "AZURE"
+#   vpc_subnet_range = "172.20.0.0/19"
+#   workspace_ids    = ["clv4wcf6f003u01m3zp7gsvzg"]
+# }
+#
+# resource "astro_cluster" "gcp_example" {
+#   type                  = "DEDICATED"
+#   name                  = "my first gcp cluster"
+#   region                = "us-central1"
+#   cloud_provider        = "GCP"
+#   pod_subnet_range      = "172.21.0.0/19"
+#   service_peering_range = "172.23.0.0/20"
+#   service_subnet_range  = "172.22.0.0/22"
+#   vpc_subnet_range      = "172.20.0.0/22"
+#   workspace_ids         = []
+# }
+#
+# // Import an existing cluster
+# import {
+#   id = "clozc036j01to01jrlgvuf98d" // ID of the existing cluster
+#   to = astro_cluster.imported_cluster
+# }
+# resource "astro_cluster" "imported_cluster" {
+#   type                  = "DEDICATED"
+#   name                  = "an existing cluster to import"
+#   region                = "us-central1"
+#   cloud_provider        = "GCP"
+#   pod_subnet_range      = "172.21.0.0/19"
+#   service_peering_range = "172.23.0.0/20"
+#   service_subnet_range  = "172.22.0.0/22"
+#   vpc_subnet_range      = "172.20.0.0/22"
+#   workspace_ids         = []
+# }
 ```
 
 <!-- schema generated by tfplugindocs -->

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_cluster/resource.tf

@@ -1,6 +1,6 @@
 resource "astro_cluster" "aws_example" {
   type             = "DEDICATED"
-  name             = "my first aws cluster"
+  name             = "LIOTTA"
   region           = "us-east-1"
   cloud_provider   = "AWS"
   vpc_subnet_range = "172.20.0.0/20"
@@ -12,40 +12,40 @@ resource "astro_cluster" "aws_example" {
   }
 }
 
-resource "astro_cluster" "azure_example" {
-  type             = "DEDICATED"
-  name             = "my first azure cluster"
-  region           = "westus2"
-  cloud_provider   = "AZURE"
-  vpc_subnet_range = "172.20.0.0/19"
-  workspace_ids    = ["clv4wcf6f003u01m3zp7gsvzg"]
-}
-
-resource "astro_cluster" "gcp_example" {
-  type                  = "DEDICATED"
-  name                  = "my first gcp cluster"
-  region                = "us-central1"
-  cloud_provider        = "GCP"
-  pod_subnet_range      = "172.21.0.0/19"
-  service_peering_range = "172.23.0.0/20"
-  service_subnet_range  = "172.22.0.0/22"
-  vpc_subnet_range      = "172.20.0.0/22"
-  workspace_ids         = []
-}
-
-// Import an existing cluster
-import {
-  id = "clozc036j01to01jrlgvuf98d" // ID of the existing cluster
-  to = astro_cluster.imported_cluster
-}
-resource "astro_cluster" "imported_cluster" {
-  type                  = "DEDICATED"
-  name                  = "an existing cluster to import"
-  region                = "us-central1"
-  cloud_provider        = "GCP"
-  pod_subnet_range      = "172.21.0.0/19"
-  service_peering_range = "172.23.0.0/20"
-  service_subnet_range  = "172.22.0.0/22"
-  vpc_subnet_range      = "172.20.0.0/22"
-  workspace_ids         = []
-}
+# resource "astro_cluster" "azure_example" {
+#   type             = "DEDICATED"
+#   name             = "my first azure cluster"
+#   region           = "westus2"
+#   cloud_provider   = "AZURE"
+#   vpc_subnet_range = "172.20.0.0/19"
+#   workspace_ids    = ["clv4wcf6f003u01m3zp7gsvzg"]
+# }
+#
+# resource "astro_cluster" "gcp_example" {
+#   type                  = "DEDICATED"
+#   name                  = "my first gcp cluster"
+#   region                = "us-central1"
+#   cloud_provider        = "GCP"
+#   pod_subnet_range      = "172.21.0.0/19"
+#   service_peering_range = "172.23.0.0/20"
+#   service_subnet_range  = "172.22.0.0/22"
+#   vpc_subnet_range      = "172.20.0.0/22"
+#   workspace_ids         = []
+# }
+#
+# // Import an existing cluster
+# import {
+#   id = "clozc036j01to01jrlgvuf98d" // ID of the existing cluster
+#   to = astro_cluster.imported_cluster
+# }
+# resource "astro_cluster" "imported_cluster" {
+#   type                  = "DEDICATED"
+#   name                  = "an existing cluster to import"
+#   region                = "us-central1"
+#   cloud_provider        = "GCP"
+#   pod_subnet_range      = "172.21.0.0/19"
+#   service_peering_range = "172.23.0.0/20"
+#   service_subnet_range  = "172.22.0.0/22"
+#   vpc_subnet_range      = "172.20.0.0/22"
+#   workspace_ids         = []
+# }

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/datasources/data_source_deployments_test.go

@@ -150,7 +150,7 @@ resource "astro_deployment" "test_deployment_celery" {
 	name = "%v-2"
 	description = "%v"
 	type = "STANDARD"
-	region = "us-east-1"
+	region = "us-west-2"
 	cloud_provider = "AWS"
 	contact_emails = []
 	default_task_pod_cpu = "0.25"

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"`