AWS Batch Compute Environment add update_policy parameter

.changelog/34353.txt

@@ -0,0 +1,7 @@
+```release-note:enhancement
+resource/aws_batch_compute_environment: Add `update_policy` parameter
+```
+
+```release-note:enhancement
+data-source/aws_batch_compute_environment: Add `update_policy` attribute
+```

internal/service/batch/compute_environment.go

@@ -259,6 +259,23 @@ func ResourceComputeEnvironment() *schema.Resource {
 				},
 				ValidateFunc: validation.StringInSlice(batch.CEType_Values(), true),
 			},
+			"update_policy": {
+				Type:     schema.TypeList,
+				Optional: true,
+				MaxItems: 1,
+				Elem: &schema.Resource{
+					Schema: map[string]*schema.Schema{
+						"job_execution_timeout_minutes": {
+							Type:     schema.TypeInt,
+							Required: true,
+						},
+						"terminate_jobs_on_update": {
+							Type:     schema.TypeBool,
+							Required: true,
+						},
+					},
+				},
+			},
 		},
 	}
 }
@@ -300,6 +317,23 @@ func resourceComputeEnvironmentCreate(ctx context.Context, d *schema.ResourceDat
 		return sdkdiag.AppendErrorf(diags, "waiting for Batch Compute Environment (%s) create: %s", d.Id(), err)
 	}
 
+	// UpdatePolicy is not possible to set with CreateComputeEnvironment
+	if v, ok := d.GetOk("update_policy"); ok && len(v.([]interface{})) > 0 && v.([]interface{})[0] != nil {
+		inputeUpdateOnCreate := &batch.UpdateComputeEnvironmentInput{
+			ComputeEnvironment: aws.String(d.Id()),
+			UpdatePolicy:       expandComputeEnvironmentUpdatePolicy(v.([]interface{})),
+		}
+		log.Printf("[DEBUG] Creating Batch Compute Environment extra arguments: %s", input)
+
+		if _, err := conn.UpdateComputeEnvironmentWithContext(ctx, inputeUpdateOnCreate); err != nil {
+			return sdkdiag.AppendErrorf(diags, "Create Batch Compute Environment extra arguments through UpdateComputeEnvironment (%s): %s", d.Id(), err)
+		}
+
+		if err := waitComputeEnvironmentUpdated(ctx, conn, d.Id(), d.Timeout(schema.TimeoutUpdate)); err != nil {
+			return sdkdiag.AppendErrorf(diags, "Create waiting for Batch Compute Environment (%s) extra arguments through UpdateComputeEnvironment: %s", d.Id(), err)
+		}
+	}
+
 	return append(diags, resourceComputeEnvironmentRead(ctx, d, meta)...)
 }
 
@@ -345,6 +379,10 @@ func resourceComputeEnvironmentRead(ctx context.Context, d *schema.ResourceData,
 	d.Set("status_reason", computeEnvironment.StatusReason)
 	d.Set("type", computeEnvironmentType)
 
+	if err := d.Set("update_policy", flattenComputeEnvironmentUpdatePolicy(computeEnvironment.UpdatePolicy)); err != nil {
+		return sdkdiag.AppendErrorf(diags, "setting update_policy: %s", err)
+	}
+
 	setTagsOut(ctx, computeEnvironment.Tags)
 
 	return diags
@@ -367,6 +405,10 @@ func resourceComputeEnvironmentUpdate(ctx context.Context, d *schema.ResourceDat
 			input.State = aws.String(d.Get("state").(string))
 		}
 
+		if d.HasChange("update_policy") {
+			input.UpdatePolicy = expandComputeEnvironmentUpdatePolicy(d.Get("update_policy").([]interface{}))
+		}
+
 		if computeEnvironmentType := strings.ToUpper(d.Get("type").(string)); computeEnvironmentType == batch.CETypeManaged {
 			// "At least one compute-resources attribute must be specified"
 			computeResourceUpdate := &batch.ComputeResourceUpdate{
@@ -479,7 +521,7 @@ func resourceComputeEnvironmentUpdate(ctx context.Context, d *schema.ResourceDat
 			return sdkdiag.AppendErrorf(diags, "updating Batch Compute Environment (%s): %s", d.Id(), err)
 		}
 
-		if _, err := waitComputeEnvironmentUpdated(ctx, conn, d.Id(), d.Timeout(schema.TimeoutUpdate)); err != nil {
+		if err := waitComputeEnvironmentUpdated(ctx, conn, d.Id(), d.Timeout(schema.TimeoutUpdate)); err != nil {
 			return sdkdiag.AppendErrorf(diags, "waiting for Batch Compute Environment (%s) update: %s", d.Id(), err)
 		}
 	}
@@ -754,7 +796,7 @@ func waitComputeEnvironmentDisabled(ctx context.Context, conn *batch.Batch, name
 	return nil, err
 }
 
-func waitComputeEnvironmentUpdated(ctx context.Context, conn *batch.Batch, name string, timeout time.Duration) (*batch.ComputeEnvironmentDetail, error) {
+func waitComputeEnvironmentUpdated(ctx context.Context, conn *batch.Batch, name string, timeout time.Duration) error {
 	stateConf := &retry.StateChangeConf{
 		Pending: []string{batch.CEStatusUpdating},
 		Target:  []string{batch.CEStatusValid},
@@ -764,11 +806,11 @@ func waitComputeEnvironmentUpdated(ctx context.Context, conn *batch.Batch, name
 
 	outputRaw, err := stateConf.WaitForStateContext(ctx)
 
-	if v, ok := outputRaw.(*batch.ComputeEnvironmentDetail); ok {
-		return v, err
+	if _, ok := outputRaw.(*batch.ComputeEnvironmentDetail); ok {
+		return err
 	}
 
-	return nil, err
+	return err
 }
 
 func isFargateType(computeResourceType string) bool {
@@ -1213,3 +1255,31 @@ func flattenLaunchTemplateSpecification(apiObject *batch.LaunchTemplateSpecifica
 
 	return tfMap
 }
+
+func expandComputeEnvironmentUpdatePolicy(l []interface{}) *batch.UpdatePolicy {
+	if len(l) == 0 || l[0] == nil {
+		return nil
+	}
+
+	m := l[0].(map[string]interface{})
+
+	up := &batch.UpdatePolicy{
+		JobExecutionTimeoutMinutes: aws.Int64(int64(m["job_execution_timeout_minutes"].(int))),
+		TerminateJobsOnUpdate:      aws.Bool(m["terminate_jobs_on_update"].(bool)),
+	}
+
+	return up
+}
+
+func flattenComputeEnvironmentUpdatePolicy(up *batch.UpdatePolicy) []interface{} {
+	if up == nil {
+		return []interface{}{}
+	}
+
+	m := map[string]interface{}{
+		"job_execution_timeout_minutes": aws.Int64Value(up.JobExecutionTimeoutMinutes),
+		"terminate_jobs_on_update":      aws.BoolValue(up.TerminateJobsOnUpdate),
+	}
+
+	return []interface{}{m}
+}

internal/service/batch/compute_environment_data_source.go

@@ -62,6 +62,22 @@ func DataSourceComputeEnvironment() *schema.Resource {
 				Type:     schema.TypeString,
 				Computed: true,
 			},
+			"update_policy": {
+				Type:     schema.TypeList,
+				Computed: true,
+				Elem: &schema.Resource{
+					Schema: map[string]*schema.Schema{
+						"job_execution_timeout_minutes": {
+							Type:     schema.TypeInt,
+							Computed: true,
+						},
+						"terminate_jobs_on_update": {
+							Type:     schema.TypeBool,
+							Computed: true,
+						},
+					},
+				},
+			},
 		},
 	}
 }
@@ -97,6 +113,10 @@ func dataSourceComputeEnvironmentRead(ctx context.Context, d *schema.ResourceDat
 	d.Set("status_reason", computeEnvironment.StatusReason)
 	d.Set("state", computeEnvironment.State)
 
+	if err := d.Set("update_policy", flattenComputeEnvironmentUpdatePolicy(computeEnvironment.UpdatePolicy)); err != nil {
+		return sdkdiag.AppendErrorf(diags, "setting update_policy: %s", err)
+	}
+
 	if err := d.Set("tags", KeyValueTags(ctx, computeEnvironment.Tags).IgnoreAWS().IgnoreConfig(ignoreTagsConfig).Map()); err != nil {
 		return sdkdiag.AppendErrorf(diags, "setting tags: %s", err)
 	}

internal/service/batch/compute_environment_data_source_test.go

@@ -34,6 +34,32 @@ func TestAccBatchComputeEnvironmentDataSource_basic(t *testing.T) {
 					resource.TestCheckResourceAttrPair(datasourceName, "state", resourceName, "state"),
 					resource.TestCheckResourceAttrPair(datasourceName, "tags.%", resourceName, "tags.%"),
 					resource.TestCheckResourceAttrPair(datasourceName, "type", resourceName, "type"),
+					resource.TestCheckResourceAttr(datasourceName, "update_policy.#", "0"),
+				),
+			},
+		},
+	})
+}
+
+func TestAccBatchComputeEnvironmentDataSource_basicUpdatePolicy(t *testing.T) {
+	ctx := acctest.Context(t)
+	rName := sdkacctest.RandomWithPrefix("tf_acc_test_")
+	resourceName := "aws_batch_compute_environment.test"
+	datasourceName := "data.aws_batch_compute_environment.by_name"
+
+	resource.ParallelTest(t, resource.TestCase{
+		PreCheck:                 func() { acctest.PreCheck(ctx, t); testAccPreCheck(ctx, t) },
+		ErrorCheck:               acctest.ErrorCheck(t, batch.EndpointsID),
+		ProtoV5ProviderFactories: acctest.ProtoV5ProviderFactories,
+		Steps: []resource.TestStep{
+			{
+				Config: testAccComputeEnvironmentDataSourceConfig_updatePolicy(rName, 30, false),
+				Check: resource.ComposeTestCheckFunc(
+					resource.TestCheckResourceAttrPair(datasourceName, "arn", resourceName, "arn"),
+					resource.TestCheckResourceAttr(datasourceName, "update_policy.#", "1"),
+					resource.TestCheckResourceAttr(datasourceName, "update_policy.0.%", "2"),
+					resource.TestCheckResourceAttr(datasourceName, "update_policy.0.terminate_jobs_on_update", "false"),
+					resource.TestCheckResourceAttr(datasourceName, "update_policy.0.job_execution_timeout_minutes", "30"),
 				),
 			},
 		},
@@ -73,30 +99,6 @@ resource "aws_iam_instance_profile" "ecs_instance_role" {
   role = aws_iam_role.ecs_instance_role.name
 }
 
-resource "aws_iam_role" "aws_batch_service_role" {
-  name = "batch_%[1]s"
-
-  assume_role_policy = <<EOF
-{
-  "Version": "2012-10-17",
-  "Statement": [
-    {
-      "Action": "sts:AssumeRole",
-      "Effect": "Allow",
-      "Principal": {
-        "Service": "batch.${data.aws_partition.current.dns_suffix}"
-      }
-    }
-  ]
-}
-EOF
-}
-
-resource "aws_iam_role_policy_attachment" "aws_batch_service_role" {
-  role       = aws_iam_role.aws_batch_service_role.name
-  policy_arn = "arn:${data.aws_partition.current.partition}:iam::aws:policy/service-role/AWSBatchServiceRole"
-}
-
 resource "aws_security_group" "sample" {
   name = "%[1]s"
 }
@@ -134,42 +136,44 @@ resource "aws_batch_compute_environment" "test" {
     type = "EC2"
   }
 
-  service_role = aws_iam_role.aws_batch_service_role.arn
-  type         = "MANAGED"
-  depends_on   = [aws_iam_role_policy_attachment.aws_batch_service_role]
+  type = "MANAGED"
 }
 
-resource "aws_batch_compute_environment" "wrong" {
-  compute_environment_name = "%[1]s_wrong"
-
-  compute_resources {
-    instance_role = aws_iam_instance_profile.ecs_instance_role.arn
-
-    instance_type = [
-      "c4.large",
-    ]
+data "aws_batch_compute_environment" "by_name" {
+  compute_environment_name = aws_batch_compute_environment.test.compute_environment_name
+}
+`, rName)
+}
 
-    max_vcpus = 16
-    min_vcpus = 0
+func testAccComputeEnvironmentDataSourceConfig_updatePolicy(rName string, timeout int, terminate bool) string {
+	return acctest.ConfigCompose(testAccComputeEnvironmentConfig_baseDefaultSLR(rName), fmt.Sprintf(`
+resource "aws_batch_compute_environment" "test" {
+  compute_environment_name = %[1]q
 
+  compute_resources {
+    allocation_strategy = "BEST_FIT_PROGRESSIVE"
+    instance_role       = aws_iam_instance_profile.ecs_instance.arn
+    instance_type       = ["optimal"]
+    max_vcpus           = 4
+    min_vcpus           = 0
     security_group_ids = [
-      aws_security_group.sample.id,
+      aws_security_group.test.id
     ]
-
     subnets = [
-      aws_subnet.sample.id,
+      aws_subnet.test.id
     ]
-
     type = "EC2"
   }
+  update_policy {
+    job_execution_timeout_minutes = %[2]d
+    terminate_jobs_on_update      = %[3]t
+  }
 
-  service_role = aws_iam_role.aws_batch_service_role.arn
-  type         = "MANAGED"
-  depends_on   = [aws_iam_role_policy_attachment.aws_batch_service_role]
+  type = "MANAGED"
 }
 
 data "aws_batch_compute_environment" "by_name" {
   compute_environment_name = aws_batch_compute_environment.test.compute_environment_name
 }
-`, rName)
+`, rName, timeout, terminate))
 }