fix attempt output persistence for schedulerv2 (#9764)

* fix attempt output persistence for schedulerv2 * fix test * add test for failure case pushing output * fix comment
airbytehq · Jan 24, 2022 · f5b5fea · f5b5fea
1 parent ded87be
commit f5b5fea
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 6 deletions.
diff --git a/...s/src/main/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowImpl.java b/...s/src/main/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowImpl.java
@@ -195,7 +195,8 @@ private void reportSuccess(final ConnectionUpdaterInput connectionUpdaterInput)
   private void reportFailure(final ConnectionUpdaterInput connectionUpdaterInput) {
     jobCreationAndStatusUpdateActivity.attemptFailure(new AttemptFailureInput(
         connectionUpdaterInput.getJobId(),
-        connectionUpdaterInput.getAttemptId()));
+        connectionUpdaterInput.getAttemptId(),
+        standardSyncOutput.orElse(null)));
 
     final int maxAttempt = configFetchActivity.getMaxAttempt().getMaxAttempt();
     final int attemptNumber = connectionUpdaterInput.getAttemptNumber();

diff --git a/...io/airbyte/workers/temporal/scheduling/activities/JobCreationAndStatusUpdateActivity.java b/...io/airbyte/workers/temporal/scheduling/activities/JobCreationAndStatusUpdateActivity.java
@@ -111,6 +111,7 @@ class AttemptFailureInput {
 
     private long jobId;
     private int attemptId;
+    private StandardSyncOutput standardSyncOutput;
 
   }
 

diff --git a/...irbyte/workers/temporal/scheduling/activities/JobCreationAndStatusUpdateActivityImpl.java b/...irbyte/workers/temporal/scheduling/activities/JobCreationAndStatusUpdateActivityImpl.java
@@ -112,7 +112,7 @@ public void jobSuccess(final JobSuccessInput input) {
         final JobOutput jobOutput = new JobOutput().withSync(input.getStandardSyncOutput());
         jobPersistence.writeOutput(input.getJobId(), input.getAttemptId(), jobOutput);
       } else {
-        log.warn("The job {} doesn't have an input for the attempt {}", input.getJobId(), input.getAttemptId());
+        log.warn("The job {} doesn't have any output for the attempt {}", input.getJobId(), input.getAttemptId());
       }
       jobPersistence.succeedAttempt(input.getJobId(), input.getAttemptId());
       final Job job = jobPersistence.getJob(input.getJobId());
@@ -139,7 +139,13 @@ public void jobFailure(final JobFailureInput input) {
   public void attemptFailure(final AttemptFailureInput input) {
     try {
       jobPersistence.failAttempt(input.getJobId(), input.getAttemptId());
-      final Job job = jobPersistence.getJob(input.getJobId());
+
+      if (input.getStandardSyncOutput() != null) {
+        final JobOutput jobOutput = new JobOutput().withSync(input.getStandardSyncOutput());
+        jobPersistence.writeOutput(input.getJobId(), input.getAttemptId(), jobOutput);
+      } else {
+        log.warn("The job {} doesn't have any output for the attempt {}", input.getJobId(), input.getAttemptId());
+      }
     } catch (final IOException e) {
       throw new RetryableException(e);
     }

diff --git a/...irbyte/workers/temporal/scheduling/activities/JobCreationAndStatusUpdateActivityTest.java b/...irbyte/workers/temporal/scheduling/activities/JobCreationAndStatusUpdateActivityTest.java
@@ -77,6 +77,8 @@ public class JobCreationAndStatusUpdateActivityTest {
           new StandardSyncSummary()
               .withStatus(ReplicationStatus.COMPLETED));
 
+  private static final JobOutput jobOutput = new JobOutput().withSync(standardSyncOutput);
+
   @Nested
   class Creation {
 
@@ -147,7 +149,6 @@ class Update {
     @Test
     public void setJobSuccess() throws IOException {
       jobCreationAndStatusUpdateActivity.jobSuccess(new JobSuccessInput(JOB_ID, ATTEMPT_ID, standardSyncOutput));
-      final JobOutput jobOutput = new JobOutput().withSync(standardSyncOutput);
 
       Mockito.verify(mJobPersistence).writeOutput(JOB_ID, ATTEMPT_ID, jobOutput);
       Mockito.verify(mJobPersistence).succeedAttempt(JOB_ID, ATTEMPT_ID);
@@ -185,17 +186,18 @@ public void setJobFailureWrapException() throws IOException {
 
     @Test
     public void setAttemptFailure() throws IOException {
-      jobCreationAndStatusUpdateActivity.attemptFailure(new AttemptFailureInput(JOB_ID, ATTEMPT_ID));
+      jobCreationAndStatusUpdateActivity.attemptFailure(new AttemptFailureInput(JOB_ID, ATTEMPT_ID, standardSyncOutput));
 
       Mockito.verify(mJobPersistence).failAttempt(JOB_ID, ATTEMPT_ID);
+      Mockito.verify(mJobPersistence).writeOutput(JOB_ID, ATTEMPT_ID, jobOutput);
     }
 
     @Test
     public void setAttemptFailureWrapException() throws IOException {
       Mockito.doThrow(new IOException())
           .when(mJobPersistence).failAttempt(JOB_ID, ATTEMPT_ID);
 
-      Assertions.assertThatThrownBy(() -> jobCreationAndStatusUpdateActivity.attemptFailure(new AttemptFailureInput(JOB_ID, ATTEMPT_ID)))
+      Assertions.assertThatThrownBy(() -> jobCreationAndStatusUpdateActivity.attemptFailure(new AttemptFailureInput(JOB_ID, ATTEMPT_ID, null)))
           .isInstanceOf(RetryableException.class)
           .hasCauseInstanceOf(IOException.class);
     }