Add support for database retention in Cloud Composer (#12565) (#20769)

[upstream:7c7f625529fc0cfac7e8ac687428172f52a019d2]

Signed-off-by: Modular Magician <[email protected]>

Author: modular-magician

.changelog/12565.txt

@@ -0,0 +1,3 @@
+```release-note:enhancement
+composer: added `airflow_metadata_retention_config` field to `composer_environment`
+```

google/services/composer/resource_composer_environment.go

@@ -706,7 +706,7 @@ func ResourceComposerEnvironment() *schema.Resource {
 									"task_logs_retention_config": {
 										Type:        schema.TypeList,
 										Description: `Optional. The configuration setting for Task Logs.`,
-										Required:    true,
+										Optional:    true,
 										Elem: &schema.Resource{
 											Schema: map[string]*schema.Schema{
 												"storage_mode": {
@@ -718,6 +718,30 @@ func ResourceComposerEnvironment() *schema.Resource {
 											},
 										},
 									},
+									"airflow_metadata_retention_config": {
+										Type:        schema.TypeList,
+										Description: `Optional. The configuration setting for database retention.`,
+										Optional:    true,
+										Computed:    true,
+										Elem: &schema.Resource{
+											Schema: map[string]*schema.Schema{
+												"retention_mode": {
+													Type:         schema.TypeString,
+													Optional:     true,
+													Computed:     true,
+													ValidateFunc: validation.StringInSlice([]string{"RETENTION_MODE_ENABLED", "RETENTION_MODE_DISABLED"}, false),
+													Description:  `Whether database retention is enabled or not. This field is supported for Cloud Composer environments in composer 3 and newer.`,
+												},
+												"retention_days": {
+													Type:         schema.TypeInt,
+													Optional:     true,
+													Computed:     true,
+													ValidateFunc: validation.IntBetween(30, 730),
+													Description:  `How many days data should be retained for. This field is supported for Cloud Composer environments in composer 3 and newer.`,
+												},
+											},
+										},
+									},
 								},
 							},
 						},
@@ -1470,6 +1494,22 @@ func resourceComposerEnvironmentUpdate(d *schema.ResourceData, meta interface{})
 				return err
 			}
 		}
+		if d.HasChange("config.0.data_retention_config.0.airflow_metadata_retention_config") {
+			patchObj := &composer.Environment{
+				Config: &composer.EnvironmentConfig{
+					DataRetentionConfig: &composer.DataRetentionConfig{
+						AirflowMetadataRetentionConfig: &composer.AirflowMetadataRetentionPolicyConfig{},
+					},
+				},
+			}
+			if config != nil && config.DataRetentionConfig != nil && config.DataRetentionConfig.AirflowMetadataRetentionConfig != nil {
+				patchObj.Config.DataRetentionConfig.AirflowMetadataRetentionConfig = config.DataRetentionConfig.AirflowMetadataRetentionConfig
+			}
+			err = resourceComposerEnvironmentPatchField("config.DataRetentionConfig.AirflowMetadataRetentionConfig", userAgent, patchObj, d, tfConfig)
+			if err != nil {
+				return err
+			}
+		}
 		if d.HasChange("config.0.recovery_config.0.scheduled_snapshots_config") {
 			patchObj := &composer.Environment{Config: &composer.EnvironmentConfig{}}
 			if config != nil {
@@ -1766,6 +1806,7 @@ func flattenComposerEnvironmentConfigDataRetentionConfig(dataRetentionConfig *co
 
 	transformed := make(map[string]interface{})
 	transformed["task_logs_retention_config"] = flattenComposerEnvironmentConfigDataRetentionConfigTaskLogsRetentionConfig(dataRetentionConfig.TaskLogsRetentionConfig)
+	transformed["airflow_metadata_retention_config"] = flattenComposerEnvironmentConfigDataRetentionConfigAirflowMetadataRetentionConfig(dataRetentionConfig.AirflowMetadataRetentionConfig)
 
 	return []interface{}{transformed}
 }
@@ -1781,6 +1822,18 @@ func flattenComposerEnvironmentConfigDataRetentionConfigTaskLogsRetentionConfig(
 	return []interface{}{transformed}
 }
 
+func flattenComposerEnvironmentConfigDataRetentionConfigAirflowMetadataRetentionConfig(airflowMetadataRetentionConfig *composer.AirflowMetadataRetentionPolicyConfig) interface{} {
+	if airflowMetadataRetentionConfig == nil {
+		return nil
+	}
+
+	transformed := make(map[string]interface{})
+	transformed["retention_mode"] = airflowMetadataRetentionConfig.RetentionMode
+	transformed["retention_days"] = airflowMetadataRetentionConfig.RetentionDays
+
+	return []interface{}{transformed}
+}
+
 func flattenComposerEnvironmentConfigWorkloadsConfig(workloadsConfig *composer.WorkloadsConfig) interface{} {
 	if workloadsConfig == nil {
 		return nil
@@ -2249,6 +2302,13 @@ func expandComposerEnvironmentConfigDataRetentionConfig(v interface{}, d *schema
 		}
 		transformed.TaskLogsRetentionConfig = transformedTaskLogsRetentionConfig
 	}
+	if airflowMetadataRetentionConfig, ok := original["airflow_metadata_retention_config"]; ok {
+		transformedAirflowMetadataRetentionConfig, err := expandComposerEnvironmentConfigDataRetentionConfigAirflowMetadataRetentionConfig(airflowMetadataRetentionConfig, d, config)
+		if err != nil {
+			return nil, err
+		}
+		transformed.AirflowMetadataRetentionConfig = transformedAirflowMetadataRetentionConfig
+	}
 
 	return transformed, nil
 }
@@ -2269,6 +2329,25 @@ func expandComposerEnvironmentConfigDataRetentionConfigTaskLogsRetentionConfig(v
 	return transformed, nil
 }
 
+func expandComposerEnvironmentConfigDataRetentionConfigAirflowMetadataRetentionConfig(v interface{}, d *schema.ResourceData, config *transport_tpg.Config) (*composer.AirflowMetadataRetentionPolicyConfig, error) {
+	l := v.([]interface{})
+	if len(l) == 0 {
+		return nil, nil
+	}
+	raw := l[0]
+	original := raw.(map[string]interface{})
+	transformed := &composer.AirflowMetadataRetentionPolicyConfig{}
+
+	if v, ok := original["retention_mode"]; ok {
+		transformed.RetentionMode = v.(string)
+	}
+	if v, ok := original["retention_days"]; ok {
+		transformed.RetentionDays = int64(v.(int))
+	}
+
+	return transformed, nil
+}
+
 func expandComposerEnvironmentConfigWorkloadsConfig(v interface{}, d *schema.ResourceData, config *transport_tpg.Config) (*composer.WorkloadsConfig, error) {
 	l := v.([]interface{})
 	if len(l) == 0 {

google/services/composer/resource_composer_environment_test.go

@@ -1174,6 +1174,39 @@ func TestAccComposerEnvironmentComposer3_withNetworkSubnetworkAndAttachment_expe
 	})
 }
 
+func TestAccComposerEnvironmentComposer3_databaseRetention(t *testing.T) {
+	t.Parallel()
+
+	envName := fmt.Sprintf("%s-%d", testComposerEnvironmentPrefix, acctest.RandInt(t))
+	network := fmt.Sprintf("%s-%d", testComposerNetworkPrefix, acctest.RandInt(t))
+	subnetwork := network + "-1"
+
+	acctest.VcrTest(t, resource.TestCase{
+		PreCheck:                 func() { acctest.AccTestPreCheck(t) },
+		ProtoV5ProviderFactories: acctest.ProtoV5ProviderFactories(t),
+		CheckDestroy:             testAccComposerEnvironmentDestroyProducer(t),
+		Steps: []resource.TestStep{
+			{
+				Config: testAccComposerEnvironmentComposer3_databaseRetention(envName, network, subnetwork),
+			},
+			{
+				ResourceName:      "google_composer_environment.test",
+				ImportState:       true,
+				ImportStateVerify: true,
+			},
+			// This is a terrible clean-up step in order to get destroy to succeed,
+			// due to dangling firewall rules left by the Composer Environment blocking network deletion.
+			// TODO: Remove this check if firewall rules bug gets fixed by Composer.
+			{
+				PlanOnly:           true,
+				ExpectNonEmptyPlan: false,
+				Config:             testAccComposerEnvironmentComposer3_databaseRetention(envName, network, subnetwork),
+				Check:              testAccCheckClearComposerEnvironmentFirewalls(t, network),
+			},
+		},
+	})
+}
+
 func TestAccComposerEnvironmentComposer3_withNetworkAttachment(t *testing.T) {
 	t.Parallel()
 
@@ -3246,6 +3279,44 @@ resource "google_compute_subnetwork" "test" {
 `, name, networkAttachment, network, subnetwork)
 }
 
+func testAccComposerEnvironmentComposer3_databaseRetention(name, network, subnetwork string) string {
+	return fmt.Sprintf(`
+resource "google_composer_environment" "test" {
+  name   = "%s"
+  region = "us-central1"
+  config {
+    software_config {
+      image_version = "composer-3-airflow-2"
+    }
+    node_config {
+      network    = google_compute_network.test.id
+      subnetwork = google_compute_subnetwork.test.id
+    }
+	data_retention_config {
+      airflow_metadata_retention_config {
+        retention_mode = "RETENTION_MODE_ENABLED"
+        retention_days = 61
+      }
+    }
+  }
+}
+
+// use a separate network to avoid conflicts with other tests running in parallel
+// that use the default network/subnet
+resource "google_compute_network" "test" {
+  name                    = "%s"
+  auto_create_subnetworks = false
+}
+
+resource "google_compute_subnetwork" "test" {
+  name          = "%s"
+  ip_cidr_range = "10.2.0.0/16"
+  region        = "us-central1"
+  network       = google_compute_network.test.self_link
+}
+`, name, network, subnetwork)
+}
+
 // WARNING: This is not actually a check and is a terrible clean-up step because Composer Environments
 // have a bug that hasn't been fixed. Composer will add firewalls to non-default networks for environments
 // but will not remove them when the Environment is deleted.