Flink fix terminal streaming events

Signed-off-by: Pawel Leszczynski <[email protected]>

Author: pawel-big-lebowski

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 ## [Unreleased](https://github.com/MarquezProject/marquez/compare/0.45.0...HEAD)
 
+### Fixed
+
+* Streaming API: fix behaviour for `COMPLETE`/`FAIL` events within streaming jobs [`#2768`](https://github.com/MarquezProject/marquez/pull/2768) [@pawel-big-lebowski]( https://github.com/pawel-big-lebowski) 
+  *New `job_version` is not created for a streaming job terminal event with no dataset information and existing version is kept.*
+
 ## [0.45.0](https://github.com/MarquezProject/marquez/compare/0.44.0...0.45.0) - 2024-03-07
 
 ### Added

api/src/main/java/marquez/db/OpenLineageDao.java

@@ -374,7 +374,7 @@ default UpdateLineageRow updateBaseMarquezModel(LineageEvent event, ObjectMapper
         insertDatasetFacets(daos, dataset, record, runUuid, event.getEventType(), now);
         insertInputDatasetFacets(daos, dataset, record, runUuid, event.getEventType(), now);
       }
-    } else {
+    } else if (!event.isTerminalEventForStreamingJobWithNoDatasets()) {
       // mark job_versions_io_mapping as obsolete
       daos.getJobVersionDao().markInputOrOutputDatasetAsPreviousFor(job.getUuid(), IoType.INPUT);
     }
@@ -390,7 +390,7 @@ default UpdateLineageRow updateBaseMarquezModel(LineageEvent event, ObjectMapper
         insertDatasetFacets(daos, dataset, record, runUuid, event.getEventType(), now);
         insertOutputDatasetFacets(daos, dataset, record, runUuid, event.getEventType(), now);
       }
-    } else {
+    } else if (!event.isTerminalEventForStreamingJobWithNoDatasets()) {
       // mark job_versions_io_mapping as obsolete
       daos.getJobVersionDao().markInputOrOutputDatasetAsPreviousFor(job.getUuid(), IoType.OUTPUT);
     }
@@ -791,6 +791,10 @@ default void updateMarquezOnStreamingJob(
         jobVersionDao.loadJobRowRunDetails(
             updateLineageRow.getJob(), updateLineageRow.getRun().getUuid());
 
+    if (event.isTerminalEventForStreamingJobWithNoDatasets()) {
+      return;
+    }
+
     if (!jobVersionDao.versionExists(jobRowRunDetails.jobVersion().getValue())) {
       // need to insert new job version
       BagOfJobVersionInfo bagOfJobVersionInfo =

api/src/main/java/marquez/service/models/LineageEvent.java

@@ -52,6 +52,20 @@ public class LineageEvent extends BaseEvent {
   @Valid @NotNull private String producer;
   @Valid private URI schemaURL;
 
+  @JsonIgnore
+  public boolean isTerminalEvent() {
+    return (eventType != null)
+        && (eventType.equalsIgnoreCase("COMPLETE") || eventType.equalsIgnoreCase("FAIL"));
+  }
+
+  @JsonIgnore
+  public boolean isTerminalEventForStreamingJobWithNoDatasets() {
+    return isTerminalEvent()
+        && (job != null && job.isStreamingJob())
+        && (outputs == null || outputs.isEmpty())
+        && (inputs == null || inputs.isEmpty());
+  }
+
   @AllArgsConstructor
   @NoArgsConstructor
   @Setter

api/src/test/java/marquez/service/LineageServiceTest.java

@@ -555,6 +555,41 @@ public void testGetLineageForRunningStreamingJob() {
     assertThat(lineageFromInput.getGraph()).isEqualTo(lineageFromOutput.getGraph());
   }
 
+  @Test
+  public void testGetLineageForCompleteStreamingJob() {
+    Dataset input = Dataset.builder().name("input-dataset").namespace(NAMESPACE).build();
+    Dataset output = Dataset.builder().name("output-dataset").namespace(NAMESPACE).build();
+
+    LineageTestUtils.createLineageRow(
+        openLineageDao,
+        "streamingjob",
+        "RUNNING",
+        JobFacet.builder()
+            .jobType(JobTypeJobFacet.builder().processingType("STREAMING").build())
+            .build(),
+        Arrays.asList(input),
+        Arrays.asList(output));
+
+    LineageTestUtils.createLineageRow(
+        openLineageDao,
+        "streamingjob",
+        "COMPLETE",
+        JobFacet.builder()
+            .jobType(JobTypeJobFacet.builder().processingType("STREAMING").build())
+            .build(),
+        Collections.emptyList(),
+        Collections.emptyList());
+
+    Lineage lineage =
+        lineageService.lineage(
+            NodeId.of(
+                new DatasetId(new NamespaceName(NAMESPACE), new DatasetName("output-dataset"))),
+            5,
+            true);
+
+    assertThat(lineage.getGraph()).hasSize(3); // 1 job + 2 datasets
+  }
+
   @Test
   public void testLineageForOrphanedDataset() {
     UpdateLineageRow writeJob =