Merge branch 'main' into bug/spec-missing-required-fields-in-dataset

Author: wslulciuc

CHANGELOG.md

@@ -6,6 +6,8 @@
 
 * Column-lineage endpoints supports point-in-time requests [`#2265`](https://github.com/MarquezProject/marquez/pull/2265) [@pawel-big-lebowski](https://github.com/pawel-big-lebowski)
     *Enable requesting `column-lineage` endpoint by a dataset version, job version or dataset field of a specific dataset version.*
+* Column lineage point in time java client [`#2269`](https://github.com/MarquezProject/marquez/pull/2269) [@pawel-big-lebowski](https://github.com/pawel-big-lebowski)
+    *Java client methods to retrieve point in time `column-lineage`. Please note that existing methods `getColumnLineageByDataset`, `getColumnLineageByDataset` and `getColumnLineageByDatasetField` were replaced by a single `getColumnLineage` taking `NodeId` as a parameter.*
 
 ### Fixed
 

api/src/main/java/marquez/api/BaseResource.java

@@ -18,9 +18,11 @@
 import marquez.api.exceptions.RunAlreadyExistsException;
 import marquez.api.exceptions.RunNotFoundException;
 import marquez.api.exceptions.SourceNotFoundException;
+import marquez.common.models.DatasetFieldId;
 import marquez.common.models.DatasetId;
 import marquez.common.models.DatasetName;
 import marquez.common.models.FieldName;
+import marquez.common.models.JobId;
 import marquez.common.models.JobName;
 import marquez.common.models.NamespaceName;
 import marquez.common.models.RunId;
@@ -37,6 +39,7 @@
 import marquez.service.ServiceFactory;
 import marquez.service.SourceService;
 import marquez.service.TagService;
+import marquez.service.models.NodeId;
 import marquez.service.models.Run;
 
 public class BaseResource {
@@ -74,6 +77,10 @@ void throwIfNotExists(@NonNull NamespaceName namespaceName) {
     }
   }
 
+  void throwIfNotExists(@NonNull DatasetId datasetId) {
+    throwIfNotExists(datasetId.getNamespace(), datasetId.getName());
+  }
+
   void throwIfNotExists(@NonNull NamespaceName namespaceName, @NonNull DatasetName datasetName) {
     if (!datasetService.exists(namespaceName.getValue(), datasetName.getValue())) {
       throw new DatasetNotFoundException(datasetName);
@@ -86,6 +93,13 @@ void throwIfSourceNotExists(SourceName sourceName) {
     }
   }
 
+  void throwIfNotExists(@NonNull DatasetFieldId datasetFieldId) {
+    throwIfNotExists(
+        datasetFieldId.getDatasetId().getNamespace(),
+        datasetFieldId.getDatasetId().getName(),
+        datasetFieldId.getFieldName());
+  }
+
   void throwIfNotExists(
       @NonNull NamespaceName namespaceName,
       @NonNull DatasetName datasetName,
@@ -96,6 +110,10 @@ void throwIfNotExists(
     }
   }
 
+  void throwIfNotExists(@NonNull JobId jobId) {
+    throwIfNotExists(jobId.getNamespace(), jobId.getName());
+  }
+
   void throwIfNotExists(@NonNull NamespaceName namespaceName, @NonNull JobName jobName) {
     if (!jobService.exists(namespaceName.getValue(), jobName.getValue())) {
       throw new JobNotFoundException(jobName);
@@ -137,6 +155,20 @@ void throwIfDatasetsNotExist(ImmutableSet<DatasetId> datasets) {
     }
   }
 
+  void throwIfNotExists(@NonNull NodeId nodeId) {
+    if (!nodeId.hasVersion()) {
+      if (nodeId.isDatasetType()) {
+        throwIfNotExists(nodeId.asDatasetId());
+      } else if (nodeId.isDatasetFieldType()) {
+        throwIfNotExists(nodeId.asDatasetFieldId());
+      } else if (nodeId.isJobType()) {
+        throwIfNotExists(nodeId.asJobId());
+      } else if (nodeId.isRunType()) {
+        throwIfNotExists(nodeId.asRunId());
+      }
+    }
+  }
+
   URI locationFor(@NonNull UriInfo uriInfo, @NonNull Run run) {
     return uriInfo
         .getBaseUriBuilder()

api/src/main/java/marquez/api/OpenLineageResource.java

@@ -96,6 +96,7 @@ private int determineStatusCode(Throwable e) {
   public Response getLineage(
       @QueryParam("nodeId") @NotNull NodeId nodeId,
       @QueryParam("depth") @DefaultValue(DEFAULT_DEPTH) int depth) {
+    throwIfNotExists(nodeId);
     return Response.ok(lineageService.lineage(nodeId, depth, true)).build();
   }
 

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

@@ -80,6 +80,14 @@ SELECT DISTINCT ON (j.uuid) j.*, inputs AS input_uuids, outputs AS output_uuids,
       WHERE ds.uuid IN (<dsUuids>)""")
   Set<DatasetData> getDatasetData(@BindList Set<UUID> dsUuids);
 
+  @SqlQuery(
+      """
+      SELECT ds.*, dv.fields, dv.lifecycle_state
+      FROM datasets_view ds
+      LEFT JOIN dataset_versions dv on dv.uuid = ds.current_version_uuid
+      WHERE ds.name = :datasetName AND ds.namespace_name = :namespaceName""")
+  DatasetData getDatasetData(String namespaceName, String datasetName);
+
   @SqlQuery(
       """
       SELECT j.uuid FROM jobs j

api/src/main/java/marquez/service/LineageService.java

@@ -7,6 +7,7 @@
 
 import com.google.common.base.Functions;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedSet;
 import com.google.common.collect.Maps;
 import java.util.Collections;
 import java.util.HashMap;
@@ -20,6 +21,7 @@
 import java.util.UUID;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import lombok.NonNull;
 import lombok.extern.slf4j.Slf4j;
 import marquez.common.models.DatasetId;
 import marquez.common.models.JobId;
@@ -48,14 +50,30 @@ public LineageService(LineageDao delegate, JobDao jobDao) {
 
   // TODO make input parameters easily extendable if adding more options like 'withJobFacets'
   public Lineage lineage(NodeId nodeId, int depth, boolean withRunFacets) {
+    log.debug("Attempting to get lineage for node '{}' with depth '{}'", nodeId.getValue(), depth);
     Optional<UUID> optionalUUID = getJobUuid(nodeId);
     if (optionalUUID.isEmpty()) {
-      throw new NodeIdNotFoundException("Could not find node");
+      log.warn(
+          "Failed to get job associated with node '{}', returning orphan graph...",
+          nodeId.getValue());
+      return toLineageWithOrphanDataset(nodeId.asDatasetId());
     }
     UUID job = optionalUUID.get();
-
+    log.debug("Attempting to get lineage for job '{}'", job);
     Set<JobData> jobData = getLineage(Collections.singleton(job), depth);
 
+    // Ensure job data is not empty, an empty set cannot be passed to LineageDao.getCurrentRuns() or
+    // LineageDao.getCurrentRunsWithFacets().
+    if (jobData.isEmpty()) {
+      // Log warning, then return an orphan lineage graph; a graph should contain at most one
+      // job->dataset relationship.
+      log.warn(
+          "Failed to get lineage for job '{}' associated with node '{}', returning orphan graph...",
+          job,
+          nodeId.getValue());
+      return toLineageWithOrphanDataset(nodeId.asDatasetId());
+    }
+
     List<Run> runs =
         withRunFacets
             ? getCurrentRunsWithFacets(
@@ -81,10 +99,26 @@ public Lineage lineage(NodeId nodeId, int depth, boolean withRunFacets) {
     if (!datasetIds.isEmpty()) {
       datasets.addAll(this.getDatasetData(datasetIds));
     }
+    if (nodeId.isDatasetType()
+        && datasets.stream().noneMatch(n -> n.getId().equals(nodeId.asDatasetId()))) {
+      log.warn(
+          "Found jobs {} which no longer share lineage with dataset '{}' - discarding",
+          jobData.stream().map(JobData::getId).toList(),
+          nodeId.getValue());
+      return toLineageWithOrphanDataset(nodeId.asDatasetId());
+    }
 
     return toLineage(jobData, datasets);
   }
 
+  private Lineage toLineageWithOrphanDataset(@NonNull DatasetId datasetId) {
+    final DatasetData datasetData =
+        getDatasetData(datasetId.getNamespace().getValue(), datasetId.getName().getValue());
+    return new Lineage(
+        ImmutableSortedSet.of(
+            Node.dataset().data(datasetData).id(NodeId.of(datasetData.getId())).build()));
+  }
+
   private Lineage toLineage(Set<JobData> jobData, Set<DatasetData> datasets) {
     Set<Node> nodes = new LinkedHashSet<>();
     // build mapping for later
@@ -214,7 +248,8 @@ public Optional<UUID> getJobUuid(NodeId nodeId) {
       return getJobFromInputOrOutput(
           datasetId.getName().getValue(), datasetId.getNamespace().getValue());
     } else {
-      throw new NodeIdNotFoundException("Node must be a dataset node or job node");
+      throw new NodeIdNotFoundException(
+          String.format("Node '%s' must be of type dataset or job!", nodeId.getValue()));
     }
   }
 }

api/src/main/resources/marquez/db/migration/V52__alter_dataset_symlinks.sql

@@ -1,2 +1,28 @@
 /* SPDX-License-Identifier: Apache-2.0 */
-ALTER TABLE dataset_symlinks ALTER COLUMN name TYPE VARCHAR;
+
+DO
+$$
+DECLARE
+    datasets_view_exists boolean;
+    datasets_view_definition text;
+BEGIN
+    SELECT EXISTS (
+        SELECT * FROM information_schema.views
+        WHERE table_schema='public' AND table_name='datasets_view'
+    ) INTO datasets_view_exists;
+
+    IF datasets_view_exists THEN
+        -- Altering is not allowed when the column is being used from views. So here,
+        -- we temporarily drop the view before altering and recreate it.
+        SELECT view_definition FROM information_schema.views
+        WHERE table_schema='public' AND table_name='datasets_view'
+        INTO datasets_view_definition;
+
+        DROP VIEW datasets_view;
+        ALTER TABLE dataset_symlinks ALTER COLUMN name TYPE VARCHAR;
+        EXECUTE format('CREATE VIEW datasets_view AS %s', datasets_view_definition);
+    ELSE
+        ALTER TABLE dataset_symlinks ALTER COLUMN name TYPE VARCHAR;
+    END IF;
+END
+$$ LANGUAGE plpgsql;

api/src/main/resources/marquez/db/migration/V54__lineage_events_created_at_indexed.sql

@@ -0,0 +1,7 @@
+-- new column will be created with 'null' filled for existing rows
+ALTER TABLE lineage_events ADD created_at TIMESTAMP WITH TIME ZONE;
+
+create index lineage_events_created_at_index on lineage_events (created_at desc NULLS LAST);
+
+-- The new default set to UTC now() will only apply in subsequent INSERT or UPDATE commands; it does not cause rows already in the table to change.
+ALTER TABLE lineage_events ALTER COLUMN created_at SET DEFAULT (now() AT TIME ZONE 'UTC')::timestamptz;

api/src/test/java/marquez/ColumnLineageIntegrationTest.java

@@ -14,7 +14,11 @@
 import java.util.Optional;
 import marquez.api.JdbiUtils;
 import marquez.client.MarquezClient;
+import marquez.client.models.DatasetFieldId;
+import marquez.client.models.DatasetId;
+import marquez.client.models.JobId;
 import marquez.client.models.Node;
+import marquez.client.models.NodeId;
 import marquez.db.LineageTestUtils;
 import marquez.db.OpenLineageDao;
 import marquez.jdbi.MarquezJdbiExternalPostgresExtension;
@@ -64,7 +68,8 @@ public void tearDown(Jdbi jdbi) {
 
   @Test
   public void testColumnLineageEndpointByDataset() {
-    MarquezClient.Lineage lineage = client.getColumnLineageByDataset("namespace", "dataset_b");
+    MarquezClient.Lineage lineage =
+        client.getColumnLineage(NodeId.of(new DatasetId("namespace", "dataset_b")));
 
     assertThat(lineage.getGraph()).hasSize(3);
     assertThat(getNodeByFieldName(lineage, "col_a")).isPresent();
@@ -75,7 +80,7 @@ public void testColumnLineageEndpointByDataset() {
   @Test
   public void testColumnLineageEndpointByDatasetField() {
     MarquezClient.Lineage lineage =
-        client.getColumnLineageByDataset("namespace", "dataset_b", "col_c");
+        client.getColumnLineage(NodeId.of(new DatasetFieldId("namespace", "dataset_b", "col_c")));
 
     assertThat(lineage.getGraph()).hasSize(3);
     assertThat(getNodeByFieldName(lineage, "col_a")).isPresent();
@@ -86,7 +91,8 @@ public void testColumnLineageEndpointByDatasetField() {
   @Test
   public void testColumnLineageEndpointWithDepthLimit() {
     MarquezClient.Lineage lineage =
-        client.getColumnLineageByDatasetField("namespace", "dataset_c", "col_d", 1, false);
+        client.getColumnLineage(
+            NodeId.of(new DatasetFieldId("namespace", "dataset_c", "col_d")), 1, false);
 
     assertThat(lineage.getGraph()).hasSize(2);
     assertThat(getNodeByFieldName(lineage, "col_c")).isPresent();
@@ -96,15 +102,16 @@ public void testColumnLineageEndpointWithDepthLimit() {
   @Test
   public void testColumnLineageEndpointWithDownstream() {
     MarquezClient.Lineage lineage =
-        client.getColumnLineageByDatasetField("namespace", "dataset_b", "col_c", 10, true);
+        client.getColumnLineage(NodeId.of(new JobId("namespace", "job1")), 10, true);
 
     assertThat(lineage.getGraph()).hasSize(4);
     assertThat(getNodeByFieldName(lineage, "col_d")).isPresent();
   }
 
   @Test
   public void testColumnLineageEndpointByJob() {
-    MarquezClient.Lineage lineage = client.getColumnLineageByJob("namespace", "job1");
+    MarquezClient.Lineage lineage =
+        client.getColumnLineage(NodeId.of(new JobId("namespace", "job1")), 1, false);
 
     assertThat(lineage.getGraph()).hasSize(3);
     assertThat(getNodeByFieldName(lineage, "col_a")).isPresent();

api/src/test/java/marquez/api/OpenLineageResourceTest.java

@@ -9,6 +9,7 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -20,6 +21,7 @@
 import javax.ws.rs.core.Response;
 import marquez.common.Utils;
 import marquez.db.OpenLineageDao;
+import marquez.service.JobService;
 import marquez.service.LineageService;
 import marquez.service.ServiceFactory;
 import marquez.service.models.Lineage;
@@ -36,6 +38,8 @@ class OpenLineageResourceTest {
   static {
     LineageService lineageService = mock(LineageService.class);
     OpenLineageDao openLineageDao = mock(OpenLineageDao.class);
+    JobService jobService = mock(JobService.class);
+    when(jobService.exists(anyString(), anyString())).thenReturn(true);
 
     Node testNode =
         Utils.fromJson(
@@ -45,7 +49,8 @@ class OpenLineageResourceTest {
     when(lineageService.lineage(any(NodeId.class), anyInt(), anyBoolean())).thenReturn(LINEAGE);
 
     ServiceFactory serviceFactory =
-        ApiTestUtils.mockServiceFactory(Map.of(LineageService.class, lineageService));
+        ApiTestUtils.mockServiceFactory(
+            Map.of(LineageService.class, lineageService, JobService.class, jobService));
 
     UNDER_TEST =
         ResourceExtension.builder()
@@ -58,7 +63,7 @@ public void testGetLineage() {
     final Lineage lineage =
         UNDER_TEST
             .target("/api/v1/lineage")
-            .queryParam("nodeId", "job:test")
+            .queryParam("nodeId", "job:test-namespace:test-job")
             .request()
             .get()
             .readEntity(Lineage.class);

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

@@ -17,6 +17,7 @@
 import marquez.api.JdbiUtils;
 import marquez.common.models.DatasetName;
 import marquez.common.models.DatasetVersionId;
+import marquez.common.models.JobId;
 import marquez.common.models.JobName;
 import marquez.common.models.NamespaceName;
 import marquez.db.DatasetDao;
@@ -393,6 +394,38 @@ public void testLineageWithWithCycle() {
         .matches(n -> n.isJobType() && n.asJobId().getName().getValue().equals("writeJob"));
   }
 
+  @Test
+  public void testLineageForOrphanedDataset() {
+    UpdateLineageRow writeJob =
+        LineageTestUtils.createLineageRow(
+            openLineageDao,
+            "writeJob",
+            "COMPLETE",
+            jobFacet,
+            Arrays.asList(),
+            Arrays.asList(dataset));
+
+    NodeId datasetNodeId =
+        NodeId.of(new NamespaceName(dataset.getNamespace()), new DatasetName(dataset.getName()));
+    Lineage lineage = lineageService.lineage(datasetNodeId, 2, false);
+    assertThat(lineage.getGraph())
+        .hasSize(2)
+        .extracting(Node::getId)
+        .containsExactlyInAnyOrder(
+            NodeId.of(new JobId(new NamespaceName(NAMESPACE), new JobName("writeJob"))),
+            datasetNodeId);
+
+    UpdateLineageRow updatedWriteJob =
+        LineageTestUtils.createLineageRow(
+            openLineageDao, "writeJob", "COMPLETE", jobFacet, Arrays.asList(), Arrays.asList());
+
+    lineage = lineageService.lineage(datasetNodeId, 2, false);
+    assertThat(lineage.getGraph())
+        .hasSize(1)
+        .extracting(Node::getId)
+        .containsExactlyInAnyOrder(datasetNodeId);
+  }
+
   private boolean jobNameEquals(Node node, String writeJob) {
     return node.getId().asJobId().getName().getValue().equals(writeJob);
   }