Fix draining job lookup for Dataflow updates

Author: lalitium

CHANGES.md

@@ -75,6 +75,7 @@
 ## Breaking Changes
 
 * (Python) Removed `google-perftools` from the SDK container images. Users who wish to use `--profiler_agent=tcmalloc` should install google-perftools APT package in their custom container images separately ([#39323](https://github.com/apache/beam/issues/39323)).
+* (Java) Added `DRAINING` and `DRAINED` states to `PipelineResult`, including runner state mappings and Dataflow update handling ([#39020](https://github.com/apache/beam/pull/39020)).
 
 ## Deprecations
 

runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java

@@ -2535,8 +2535,9 @@ private String getJobIdFromName(String jobName) {
         listResult = dataflowClient.listJobs(token);
         token = listResult.getNextPageToken();
         for (Job job : listResult.getJobs()) {
+          State state = MonitoringUtil.toState(job.getCurrentState());
           if (job.getName().equals(jobName)
-              && MonitoringUtil.toState(job.getCurrentState()).equals(State.RUNNING)) {
+              && (state.equals(State.RUNNING) || state.equals(State.DRAINING))) {
             return job.getId();
           }
         }

runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java

@@ -285,6 +285,11 @@ private static Pipeline buildDataflowPipelineWithLargeGraph(DataflowPipelineOpti
   }
 
   static Dataflow buildMockDataflow(Dataflow.Projects.Locations.Jobs mockJobs) throws IOException {
+    return buildMockDataflow(mockJobs, "JOB_STATE_RUNNING");
+  }
+
+  static Dataflow buildMockDataflow(Dataflow.Projects.Locations.Jobs mockJobs, String currentState)
+      throws IOException {
     Dataflow mockDataflowClient = mock(Dataflow.class);
     Dataflow.Projects mockProjects = mock(Dataflow.Projects.class);
     Dataflow.Projects.Locations mockLocations = mock(Dataflow.Projects.Locations.class);
@@ -308,7 +313,7 @@ static Dataflow buildMockDataflow(Dataflow.Projects.Locations.Jobs mockJobs) thr
                         new Job()
                             .setName("oldjobname")
                             .setId("oldJobId")
-                            .setCurrentState("JOB_STATE_RUNNING"))));
+                            .setCurrentState(currentState))));
 
     Job resultJob = new Job();
     resultJob.setId("newid");
@@ -375,14 +380,18 @@ static GcsUtil buildMockGcsUtil() throws IOException {
   }
 
   private DataflowPipelineOptions buildPipelineOptions() throws IOException {
+    return buildPipelineOptions("JOB_STATE_RUNNING");
+  }
+
+  private DataflowPipelineOptions buildPipelineOptions(String currentState) throws IOException {
     DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class);
     options.setRunner(DataflowRunner.class);
     options.setProject(PROJECT_ID);
     options.setTempLocation(VALID_TEMP_BUCKET);
     options.setRegion(REGION_ID);
     // Set FILES_PROPERTY to empty to prevent a default value calculated from classpath.
     options.setFilesToStage(new ArrayList<>());
-    options.setDataflowClient(buildMockDataflow(mockJobs));
+    options.setDataflowClient(buildMockDataflow(mockJobs, currentState));
     options.setGcsUtil(mockGcsUtil);
     options.setGcpCredential(new TestCredential());
 
@@ -793,6 +802,19 @@ public void testUpdate() throws IOException {
     assertValidJob(jobCaptor.getValue());
   }
 
+  @Test
+  public void testUpdateDrainingJob() throws IOException {
+    DataflowPipelineOptions options = buildPipelineOptions("JOB_STATE_DRAINING");
+    options.setUpdate(true);
+    options.setJobName("oldJobName");
+    Pipeline p = buildDataflowPipeline(options);
+    p.run();
+
+    ArgumentCaptor<Job> jobCaptor = ArgumentCaptor.forClass(Job.class);
+    Mockito.verify(mockJobs).create(eq(PROJECT_ID), eq(REGION_ID), jobCaptor.capture());
+    assertEquals("oldJobId", jobCaptor.getValue().getReplaceJobId());
+  }
+
   @Test
   public void testUploadGraph() throws IOException {
     DataflowPipelineOptions options = buildPipelineOptions();