Allow Dataflow flex template jobs to be updated (#3069)

* Allow Dataflow flex template jobs to be updated

* Update Flex Template tests to run a real job

* Add test for updating flex template streaming jobs

* Use resourceSchema parameter in resourceDataflowJobIsVirtualUpdate

Co-authored-by: Andre Marianiello <[email protected]>

Author: andremarianiello

google-beta/resource_dataflow_flex_template_job.go

@@ -30,7 +30,6 @@ func resourceDataflowFlexTemplateJob() *schema.Resource {
 			"container_spec_gcs_path": {
 				Type:     schema.TypeString,
 				Required: true,
-				ForceNew: true,
 			},
 
 			"name": {
@@ -58,15 +57,13 @@ func resourceDataflowFlexTemplateJob() *schema.Resource {
 				Type:             schema.TypeMap,
 				Optional:         true,
 				DiffSuppressFunc: resourceDataflowJobLabelDiffSuppress,
-				ForceNew:         true,
 				// TODO add support for labels when the API supports it
 				Deprecated: "Deprecated until the API supports this field",
 			},
 
 			"parameters": {
 				Type:     schema.TypeMap,
 				Optional: true,
-				ForceNew: true,
 			},
 
 			"project": {
@@ -177,10 +174,96 @@ func resourceDataflowFlexTemplateJobRead(d *schema.ResourceData, meta interface{
 	return nil
 }
 
-// resourceDataflowFlexTemplateJobUpdate is a blank method to enable updating
-// the on_delete virtual field
+func waitForDataflowJobState(d *schema.ResourceData, config *Config, jobID, userAgent string, timeout time.Duration, targetState string) error {
+	return resource.Retry(timeout, func() *resource.RetryError {
+		project, err := getProject(d, config)
+		if err != nil {
+			return resource.NonRetryableError(err)
+		}
+
+		region, err := getRegion(d, config)
+		if err != nil {
+			return resource.NonRetryableError(err)
+		}
+
+		job, err := resourceDataflowJobGetJob(config, project, region, userAgent, jobID)
+		if err != nil {
+			if isRetryableError(err) {
+				return resource.RetryableError(err)
+			}
+			return resource.NonRetryableError(err)
+		}
+
+		state := job.CurrentState
+		if state == targetState {
+			log.Printf("[DEBUG] the job with ID %q has state %q.", jobID, state)
+			return nil
+		}
+		if _, terminated := dataflowTerminalStatesMap[state]; terminated {
+			return resource.NonRetryableError(fmt.Errorf("the job with ID %q has terminated with state %q instead of expected state %q", jobID, state, targetState))
+		} else {
+			log.Printf("[DEBUG] the job with ID %q has state %q.", jobID, state)
+			return resource.RetryableError(fmt.Errorf("the job with ID %q has state %q, waiting for %q", jobID, state, targetState))
+		}
+	})
+}
+
+// resourceDataflowFlexTemplateJobUpdate updates a Flex Template Job resource.
 func resourceDataflowFlexTemplateJobUpdate(d *schema.ResourceData, meta interface{}) error {
-	return nil
+	// Don't send an update request if only virtual fields have changes
+	if resourceDataflowJobIsVirtualUpdate(d, resourceDataflowFlexTemplateJob().Schema) {
+		return nil
+	}
+
+	config := meta.(*Config)
+	userAgent, err := generateUserAgentString(d, config.userAgent)
+	if err != nil {
+		return err
+	}
+
+	project, err := getProject(d, config)
+	if err != nil {
+		return err
+	}
+
+	region, err := getRegion(d, config)
+	if err != nil {
+		return err
+	}
+
+	// wait until current job is running or terminated
+	err = waitForDataflowJobState(d, config, d.Id(), userAgent, d.Timeout(schema.TimeoutUpdate), "JOB_STATE_RUNNING")
+	if err != nil {
+		return fmt.Errorf("Error waiting for job with job ID %q to be running: %v", d.Id(), err)
+	}
+
+	request := dataflow.LaunchFlexTemplateRequest{
+		LaunchParameter: &dataflow.LaunchFlexTemplateParameter{
+			ContainerSpecGcsPath: d.Get("container_spec_gcs_path").(string),
+			JobName:              d.Get("name").(string),
+			Parameters:           expandStringMap(d, "parameters"),
+			Update:               true,
+		},
+	}
+
+	response, err := config.NewDataflowClient(userAgent).Projects.Locations.FlexTemplates.Launch(project, region, &request).Do()
+	if err != nil {
+		return err
+	}
+
+	// don't set id until previous job is successfully updated
+	err = waitForDataflowJobState(d, config, d.Id(), userAgent, d.Timeout(schema.TimeoutUpdate), "JOB_STATE_UPDATED")
+	if err != nil {
+		return fmt.Errorf("Error waiting for Job with job ID %q to be updated: %v", d.Id(), err)
+	}
+
+	job := response.Job
+	d.SetId(job.Id)
+	if err := d.Set("job_id", job.Id); err != nil {
+		return fmt.Errorf("Error setting job_id: %s", err)
+	}
+
+	return resourceDataflowFlexTemplateJobRead(d, meta)
 }
 
 func resourceDataflowFlexTemplateJobDelete(d *schema.ResourceData, meta interface{}) error {

google-beta/resource_dataflow_flex_template_job_test.go

@@ -18,7 +18,6 @@ func TestAccDataflowFlexTemplateJob_basic(t *testing.T) {
 	t.Parallel()
 
 	randStr := randString(t, 10)
-	bucket := "tf-test-dataflow-gcs-" + randStr
 	job := "tf-test-dataflow-job-" + randStr
 
 	vcrTest(t, resource.TestCase{
@@ -27,9 +26,39 @@ func TestAccDataflowFlexTemplateJob_basic(t *testing.T) {
 		CheckDestroy: testAccCheckDataflowJobDestroyProducer(t),
 		Steps: []resource.TestStep{
 			{
-				Config: testAccDataflowFlexTemplateJob_basic(bucket, job),
+				Config: testAccDataflowFlexTemplateJob_basic(job, "mytopic"),
 				Check: resource.ComposeTestCheckFunc(
-					testAccDataflowJobExists(t, "google_dataflow_flex_template_job.big_data"),
+					testAccDataflowJobExists(t, "google_dataflow_flex_template_job.job"),
+				),
+			},
+		},
+	})
+}
+
+func TestAccDataflowFlexTemplateJob_streamUpdate(t *testing.T) {
+	// This resource uses custom retry logic that cannot be sped up without
+	// modifying the actual resource
+	skipIfVcr(t)
+	t.Parallel()
+
+	randStr := randString(t, 10)
+	job := "tf-test-dataflow-job-" + randStr
+
+	vcrTest(t, resource.TestCase{
+		PreCheck:     func() { testAccPreCheck(t) },
+		Providers:    testAccProviders,
+		CheckDestroy: testAccCheckDataflowJobDestroyProducer(t),
+		Steps: []resource.TestStep{
+			{
+				Config: testAccDataflowFlexTemplateJob_basic(job, "mytopic"),
+				Check: resource.ComposeTestCheckFunc(
+					testAccDataflowJobExists(t, "google_dataflow_flex_template_job.job"),
+				),
+			},
+			{
+				Config: testAccDataflowFlexTemplateJob_basic(job, "mytopic2"),
+				Check: resource.ComposeTestCheckFunc(
+					testAccDataflowJobHasOption(t, "google_dataflow_flex_template_job.job", "topic", "projects/myproject/topics/mytopic2"),
 				),
 			},
 		},
@@ -43,7 +72,6 @@ func TestAccDataflowFlexTemplateJob_withServiceAccount(t *testing.T) {
 	t.Parallel()
 
 	randStr := randString(t, 10)
-	bucket := "tf-test-dataflow-gcs-" + randStr
 	job := "tf-test-dataflow-job-" + randStr
 	accountId := "tf-test-dataflow-sa" + randStr
 	zone := "us-central1-b"
@@ -54,10 +82,10 @@ func TestAccDataflowFlexTemplateJob_withServiceAccount(t *testing.T) {
 		CheckDestroy: testAccCheckDataflowJobDestroyProducer(t),
 		Steps: []resource.TestStep{
 			{
-				Config: testAccDataflowFlexTemplateJob_serviceAccount(bucket, job, accountId, zone),
+				Config: testAccDataflowFlexTemplateJob_serviceAccount(job, accountId, zone),
 				Check: resource.ComposeTestCheckFunc(
-					testAccDataflowJobExists(t, "google_dataflow_flex_template_job.big_data"),
-					testAccDataflowFlexTemplateJobHasServiceAccount(t, "google_dataflow_flex_template_job.big_data", accountId, zone),
+					testAccDataflowJobExists(t, "google_dataflow_flex_template_job.job"),
+					testAccDataflowFlexTemplateJobHasServiceAccount(t, "google_dataflow_flex_template_job.job", accountId, zone),
 				),
 			},
 		},
@@ -123,136 +151,87 @@ func testAccDataflowFlexTemplateJobGetGeneratedInstance(t *testing.T, s *terrafo
 	return instance, nil
 }
 
-// note: this config creates a job that doesn't actually do anything
-func testAccDataflowFlexTemplateJob_basic(bucket, job string) string {
+// note: this config creates a job that doesn't actually do anything, but still runs
+func testAccDataflowFlexTemplateJob_basic(job, topicName string) string {
 	return fmt.Sprintf(`
-resource "google_storage_bucket" "temp" {
-  name = "%s"
-  force_destroy = true
+data "google_storage_bucket_object" "flex_template" {
+  name   = "latest/flex/Streaming_Data_Generator"
+  bucket = "dataflow-templates"
 }
 
-resource "google_storage_bucket_object" "flex_template" {
-  name   = "flex_template.json"
-  bucket = google_storage_bucket.temp.name
-  content = <<EOF
-{
-    "image": "my-image",
-    "metadata": {
-        "description": "An Apache Beam streaming pipeline that reads JSON encoded messages from Pub/Sub, uses Beam SQL to transform the message data, and writes the results to a BigQuery",
-        "name": "Streaming Beam SQL",
-        "parameters": [
-            {
-                "helpText": "Pub/Sub subscription to read from.",
-                "label": "Pub/Sub input subscription.",
-                "name": "inputSubscription",
-                "regexes": [
-                    "[-_.a-zA-Z0-9]+"
-                ]
-            },
-            {
-                "helpText": "BigQuery table spec to write to, in the form 'project:dataset.table'.",
-                "is_optional": true,
-                "label": "BigQuery output table",
-                "name": "outputTable",
-                "regexes": [
-                    "[^:]+:[^.]+[.].+"
-                ]
-            }
-        ]
-    },
-    "sdkInfo": {
-        "language": "JAVA"
-    }
-}
-EOF
-}
-
-resource "google_dataflow_flex_template_job" "big_data" {
+resource "google_dataflow_flex_template_job" "job" {
   name = "%s"
-  container_spec_gcs_path = "${google_storage_bucket.temp.url}/${google_storage_bucket_object.flex_template.name}"
+  container_spec_gcs_path = "gs://${data.google_storage_bucket_object.flex_template.bucket}/${data.google_storage_bucket_object.flex_template.name}"
   on_delete = "cancel"
   parameters = {
-    inputSubscription = "my-subscription"
-    outputTable  = "my-project:my-dataset.my-table"
+    schemaLocation = "gs://mybucket/schema.json"
+    qps = "1"
+    topic = "projects/myproject/topics/%s"
   }
 }
-`, bucket, job)
+`, job, topicName)
 }
 
-// note: this config creates a job that doesn't actually do anything
-func testAccDataflowFlexTemplateJob_serviceAccount(bucket, job, accountId, zone string) string {
+// note: this config creates a job that doesn't actually do anything, but still runs
+func testAccDataflowFlexTemplateJob_serviceAccount(job, accountId, zone string) string {
 	return fmt.Sprintf(`
-resource "google_storage_bucket" "temp" {
-  name = "%s"
-  force_destroy = true
-}
-
 resource "google_service_account" "dataflow-sa" {
   account_id   = "%s"
   display_name = "DataFlow Service Account"
 }
 
-resource "google_storage_bucket_iam_member" "dataflow-gcs" {
-  bucket = google_storage_bucket.temp.name
-  role   = "roles/storage.objectAdmin"
-  member = "serviceAccount:${google_service_account.dataflow-sa.email}"
-}
-
 resource "google_project_iam_member" "dataflow-worker" {
   role   = "roles/dataflow.worker"
   member = "serviceAccount:${google_service_account.dataflow-sa.email}"
 }
 
-resource "google_storage_bucket_object" "flex_template" {
-  name   = "flex_template.json"
-  bucket = google_storage_bucket.temp.name
-  content = <<EOF
-{
-    "image": "my-image",
-    "metadata": {
-        "description": "An Apache Beam streaming pipeline that reads JSON encoded messages from Pub/Sub, uses Beam SQL to transform the message data, and writes the results to a BigQuery",
-        "name": "Streaming Beam SQL",
-        "parameters": [
-            {
-                "helpText": "Pub/Sub subscription to read from.",
-                "label": "Pub/Sub input subscription.",
-                "name": "inputSubscription",
-                "regexes": [
-                    "[-_.a-zA-Z0-9]+"
-                ]
-            },
-            {
-                "helpText": "BigQuery table spec to write to, in the form 'project:dataset.table'.",
-                "is_optional": true,
-                "label": "BigQuery output table",
-                "name": "outputTable",
-                "regexes": [
-                    "[^:]+:[^.]+[.].+"
-                ]
-            }
-        ]
-    },
-    "sdkInfo": {
-        "language": "JAVA"
-    }
-}
-EOF
+data "google_storage_bucket_object" "flex_template" {
+  name   = "latest/flex/Streaming_Data_Generator"
+  bucket = "dataflow-templates"
 }
 
-resource "google_dataflow_flex_template_job" "big_data" {
+resource "google_dataflow_flex_template_job" "job" {
   name = "%s"
-  container_spec_gcs_path = "${google_storage_bucket.temp.url}/${google_storage_bucket_object.flex_template.name}"
+  container_spec_gcs_path = "gs://${data.google_storage_bucket_object.flex_template.bucket}/${data.google_storage_bucket_object.flex_template.name}"
   on_delete = "cancel"
   parameters = {
-    inputSubscription = "my-subscription"
-    outputTable  = "my-project:my-dataset.my-table"
+    schemaLocation = "gs://mybucket/schema.json"
+    qps = "1"
+    topic = "projects/myproject/topics/mytopic"
     serviceAccount = google_service_account.dataflow-sa.email
     zone = "%s"
   }
-  depends_on = [
-    google_storage_bucket_iam_member.dataflow-gcs,
-    google_project_iam_member.dataflow-worker
-  ]
 }
-`, bucket, accountId, job, zone)
+`, accountId, job, zone)
+}
+
+func testAccDataflowJobHasOption(t *testing.T, res, option, expectedValue string) resource.TestCheckFunc {
+	return func(s *terraform.State) error {
+		rs, ok := s.RootModule().Resources[res]
+		if !ok {
+			return fmt.Errorf("resource %q not found in state", res)
+		}
+
+		if rs.Primary.ID == "" {
+			return fmt.Errorf("No ID is set")
+		}
+		config := googleProviderConfig(t)
+
+		job, err := config.NewDataflowClient(config.userAgent).Projects.Jobs.Get(config.Project, rs.Primary.ID).View("JOB_VIEW_ALL").Do()
+		if err != nil {
+			return fmt.Errorf("dataflow job does not exist")
+		}
+
+		sdkPipelineOptions, err := ConvertToMap(job.Environment.SdkPipelineOptions)
+		if err != nil {
+			return fmt.Errorf("error from ConvertToMap: %s", err)
+		}
+		optionsMap := sdkPipelineOptions["options"].(map[string]interface{})
+
+		if optionsMap[option] != expectedValue {
+			return fmt.Errorf("Option %s do not match. Got %s while expecting %s", option, optionsMap[option], expectedValue)
+		}
+
+		return nil
+	}
 }