Support for Composer scheduler count flag (beta) (hashicorp#10026) (h…

…ashicorp#5180)

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

.changelog/5180.txt

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

google/resource_composer_environment.go

@@ -42,6 +42,7 @@ var (
 		"config.0.software_config.0.env_variables",
 		"config.0.software_config.0.image_version",
 		"config.0.software_config.0.python_version",
+		"config.0.software_config.0.scheduler_count",
 	}
 
 	composerConfigKeys = []string{
@@ -284,6 +285,13 @@ func resourceComposerEnvironment() *schema.Resource {
 										ForceNew:     true,
 										Description:  `The major version of Python used to run the Apache Airflow scheduler, worker, and webserver processes. Can be set to '2' or '3'. If not specified, the default is '2'. Cannot be updated.`,
 									},
+									"scheduler_count": {
+										Type:         schema.TypeInt,
+										Optional:     true,
+										AtLeastOneOf: composerSoftwareConfigKeys,
+										Computed:     true,
+										Description:  `The number of schedulers for Airflow. This field is supported for Cloud Composer environments in versions composer-1.*.*-airflow-2.*.*.`,
+									},
 								},
 							},
 						},
@@ -788,6 +796,7 @@ func flattenComposerEnvironmentConfigSoftwareConfig(softwareCfg *composer.Softwa
 	transformed["airflow_config_overrides"] = softwareCfg.AirflowConfigOverrides
 	transformed["pypi_packages"] = softwareCfg.PypiPackages
 	transformed["env_variables"] = softwareCfg.EnvVariables
+	transformed["scheduler_count"] = softwareCfg.SchedulerCount
 	return []interface{}{transformed}
 }
 
@@ -1070,6 +1079,7 @@ func expandComposerEnvironmentConfigSoftwareConfig(v interface{}, d *schema.Reso
 	transformed.AirflowConfigOverrides = expandComposerEnvironmentConfigSoftwareConfigStringMap(original, "airflow_config_overrides")
 	transformed.PypiPackages = expandComposerEnvironmentConfigSoftwareConfigStringMap(original, "pypi_packages")
 	transformed.EnvVariables = expandComposerEnvironmentConfigSoftwareConfigStringMap(original, "env_variables")
+	transformed.SchedulerCount = int64(original["scheduler_count"].(int))
 	return transformed, nil
 }
 

google/resource_composer_environment_test.go

@@ -239,6 +239,38 @@ func TestAccComposerEnvironment_withSoftwareConfig(t *testing.T) {
 	})
 }
 
+func TestAccComposerEnvironmentAirflow2_withSoftwareConfig(t *testing.T) {
+	t.Parallel()
+	envName := fmt.Sprintf("%s-%d", testComposerEnvironmentPrefix, randInt(t))
+	network := fmt.Sprintf("%s-%d", testComposerNetworkPrefix, randInt(t))
+	subnetwork := network + "-1"
+
+	vcrTest(t, resource.TestCase{
+		PreCheck:     func() { testAccPreCheck(t) },
+		Providers:    testAccProviders,
+		CheckDestroy: testAccComposerEnvironmentDestroyProducer(t),
+		Steps: []resource.TestStep{
+			{
+				Config: testAccComposerEnvironment_airflow2SoftwareCfg(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(emilyye): Remove this check if firewall rules bug gets fixed by Composer.
+			{
+				PlanOnly:           true,
+				ExpectNonEmptyPlan: false,
+				Config:             testAccComposerEnvironment_airflow2SoftwareCfg(envName, network, subnetwork),
+				Check:              testAccCheckClearComposerEnvironmentFirewalls(t, network),
+			},
+		},
+	})
+}
+
 // Checks behavior of config for creation for attributes that must
 // be updated during create.
 func TestAccComposerEnvironment_withUpdateOnCreate(t *testing.T) {
@@ -374,6 +406,16 @@ resource "google_compute_subnetwork" "test" {
 
 func testAccComposerEnvironment_update(name, network, subnetwork string) string {
 	return fmt.Sprintf(`
+data "google_composer_image_versions" "all" {
+}
+
+locals {
+	composer_version = "1"  # both composer_version and airflow_version are parts of regex, so if either 1 or 2 version is ok "[12]" should be used,  
+	airflow_version = "1"   # if sub-version is needed remember to escape "." with "\\." for example 1.2 should be written as "1\\.2"
+	reg_ex = join("", ["composer-", local.composer_version, "\\.[\\d+\\.]*\\d+.*-airflow-", local.airflow_version, "\\.[\\d+\\.]*\\d+"])
+	matching_images = [for v in data.google_composer_image_versions.all.image_versions[*].image_version_id: v if length(regexall(local.reg_ex, v)) > 0]
+}
+
 resource "google_composer_environment" "test" {
 	name   = "%s"
 	region = "us-central1"
@@ -387,7 +429,7 @@ resource "google_composer_environment" "test" {
 		}
 
 		software_config {
-			image_version = "composer-1.16.14-airflow-1.10.15"
+			image_version = local.matching_images[0]
 
 			airflow_config_overrides = {
 				core-load_example = "True"
@@ -472,6 +514,16 @@ resource "google_project_iam_member" "composer-worker" {
 
 func testAccComposerEnvironment_softwareCfg(name, network, subnetwork string) string {
 	return fmt.Sprintf(`
+data "google_composer_image_versions" "all" {
+}
+
+locals {
+	composer_version = "1"  # both composer_version and airflow_version are parts of regex, so if either 1 or 2 version is ok "[12]" should be used,  
+	airflow_version = "1"   # if sub-version is needed remember to escape "." with "\\." for example 1.2 should be written as "1\\.2"
+	reg_ex = join("", ["composer-", local.composer_version, "\\.[\\d+\\.]*\\d+.*-airflow-", local.airflow_version, "\\.[\\d+\\.]*\\d+"])
+	matching_images = [for v in data.google_composer_image_versions.all.image_versions[*].image_version_id: v if length(regexall(local.reg_ex, v)) > 0]
+}
+
 resource "google_composer_environment" "test" {
 	name   = "%s"
 	region = "us-central1"
@@ -482,7 +534,7 @@ resource "google_composer_environment" "test" {
 			zone       = "us-central1-a"
 		}
 		software_config {
-			image_version  = "composer-1.16.14-airflow-1.10.15"
+			image_version  = local.matching_images[0]
 			python_version = "3"
 		}
 	}
@@ -539,6 +591,51 @@ resource "google_compute_subnetwork" "test" {
 `, name, network, subnetwork)
 }
 
+func testAccComposerEnvironment_airflow2SoftwareCfg(name, network, subnetwork string) string {
+	return fmt.Sprintf(`
+data "google_composer_image_versions" "all" {
+}
+
+locals {
+	composer_version = "1"  # both composer_version and airflow_version are parts of regex, so if either 1 or 2 version is ok "[12]" should be used,  
+	airflow_version = "2"   # if sub-version is needed remember to escape "." with "\\." for example 1.2 should be written as "1\\.2"
+	reg_ex = join("", ["composer-", local.composer_version, "\\.[\\d+\\.]*\\d+.*-airflow-", local.airflow_version, "\\.[\\d+\\.]*\\d+"])
+	matching_images = [for v in data.google_composer_image_versions.all.image_versions[*].image_version_id: v if length(regexall(local.reg_ex, v)) > 0]
+}
+
+
+resource "google_composer_environment" "test" {
+	name   = "%s"
+	region = "us-central1"
+	config {
+		node_config {
+			network    = google_compute_network.test.self_link
+			subnetwork = google_compute_subnetwork.test.self_link
+			zone       = "us-central1-a"
+		}
+		software_config {
+			image_version  = local.matching_images[0]
+			scheduler_count = 2
+		}
+	}
+}
+
+// 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)
+}
+
 /**
  * CLEAN UP HELPER FUNCTIONS
  * Because the environments are flaky and bucket deletion rates can be

website/docs/guides/version_3_upgrade.html.markdown

@@ -460,10 +460,10 @@ In an attempt to avoid allowing empty blocks in config files, at least one of `s
 
 ## Resource: `google_composer_environment`
 
-### At least one of `airflow_config_overrides`, `pypi_packages`, `env_variables`, `image_version`, or `python_version` is now required on `google_composer_environment.config.software_config`
+### At least one of `airflow_config_overrides`, `pypi_packages`, `env_variables`, `image_version`, `python_version` or `scheduler_count` is now required on `google_composer_environment.config.software_config`
 
 In an attempt to avoid allowing empty blocks in config files, at least one of `airflow_config_overrides`,
-`pypi_packages`, `env_variables`, `image_version`, or `python_version` is now required on the
+`pypi_packages`, `env_variables`, `image_version`, `python_version` or `scheduler_count` is now required on the
 `config.software_config` block.
 
 ### `use_ip_aliases` is now required on block `google_composer_environment.ip_allocation_policy`

website/docs/r/composer_environment.html.markdown

@@ -314,6 +314,9 @@ The `software_config` block supports:
   The major version of Python used to run the Apache Airflow scheduler, worker, and webserver processes.
   Can be set to '2' or '3'. If not specified, the default is '2'. Cannot be updated.
 
+* `scheduler_count` (Optional) -
+  The number of schedulers for Airflow. This field is supported for Cloud Composer environments in versions composer-1.*.*-airflow-2.*.*.`
+
 See [documentation](https://cloud.google.com/composer/docs/how-to/managing/configuring-private-ip) for setting up private environments. The `private_environment_config` block supports:
 
 * `enable_private_endpoint` -