From cb3e6d0b8a2cc195f539dcd75f00afd529bd27bb Mon Sep 17 00:00:00 2001
From: Joshua Carp <jm.carp@gmail.com>
Date: Mon, 13 May 2019 15:24:28 -0400
Subject: [PATCH] Make cloud composer environment image version updateable.

---
 .../resource_composer_environment.go.erb      | 20 +++++++++++++++++--
 .../resource_composer_environment_test.go.erb |  8 ++++++--
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/third_party/terraform/resources/resource_composer_environment.go.erb b/third_party/terraform/resources/resource_composer_environment.go.erb
index 70a2dc6b707f..d2563f36e47a 100644
--- a/third_party/terraform/resources/resource_composer_environment.go.erb
+++ b/third_party/terraform/resources/resource_composer_environment.go.erb
@@ -185,7 +185,6 @@ func resourceComposerEnvironment() *schema.Resource {
 										Computed: true,
 <% unless version == 'ga' -%>
 										Optional: true,
-										ForceNew: true,
 										ValidateFunc: validateRegexp(composerEnvironmentVersionRegexp),
 										DiffSuppressFunc: composerImageVersionDiffSuppress,
 									},
@@ -333,8 +332,25 @@ func resourceComposerEnvironmentUpdate(d *schema.ResourceData, meta interface{})
 			return err
 		}
 
-		if d.HasChange("config.0.software_config.0.airflow_config_overrides") {
+<% unless version == 'ga' -%>
+		if d.HasChange("config.0.software_config.0.image_version") {
+			patchObj := &composer.Environment{
+				Config: &composer.EnvironmentConfig{
+					SoftwareConfig: &composer.SoftwareConfig{},
+				},
+			}
+			if config != nil && config.SoftwareConfig != nil {
+				patchObj.Config.SoftwareConfig.ImageVersion = config.SoftwareConfig.ImageVersion
+			}
+			err = resourceComposerEnvironmentPatchField("config.softwareConfig.imageVersion", patchObj, d, tfConfig)
+			if err != nil {
+				return err
+			}
+			d.SetPartial("config")
+		}
+<% end -%>
 
+		if d.HasChange("config.0.software_config.0.airflow_config_overrides") {
 			patchObj := &composer.Environment{
 				Config: &composer.EnvironmentConfig{
 					SoftwareConfig: &composer.SoftwareConfig{
diff --git a/third_party/terraform/tests/resource_composer_environment_test.go.erb b/third_party/terraform/tests/resource_composer_environment_test.go.erb
index d7ab125eecf9..fa206d3e34d3 100644
--- a/third_party/terraform/tests/resource_composer_environment_test.go.erb
+++ b/third_party/terraform/tests/resource_composer_environment_test.go.erb
@@ -279,8 +279,8 @@ func testAccComposerEnvironmentDestroy(s *terraform.State) error {
 func testAccComposerEnvironment_basic(name string) string {
 	return fmt.Sprintf(`
 resource "google_composer_environment" "test" {
-  name           = "%s"
-  region         = "us-central1"
+	name           = "%s"
+	region         = "us-central1"
 }
 `, name)
 }
@@ -295,6 +295,10 @@ resource "google_composer_environment" "test" {
 		node_count = 4
 
 		software_config {
+<% unless version == 'ga' -%>
+			image_version = "composer-1.7.0-airflow-1.10.2"
+<% end -%>
+
 			airflow_config_overrides = {
 			  core-load_example = "True"
 			}