provide simple implementation of one-level lineage optimized for parent jobs

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

@@ -107,6 +107,22 @@ public Response getLineage(
     return Response.ok(lineageService.lineage(nodeId, depth, true)).build();
   }
 
+  @Timed
+  @ResponseMetered
+  @ExceptionMetered
+  @GET
+  @Consumes(APPLICATION_JSON)
+  @Produces(APPLICATION_JSON)
+  @Path("/lineage/direct")
+  public Response getSimpleLineage(
+      @QueryParam("parentJobNodeId") @NotNull NodeId parentJobNodeId) {
+    if (!parentJobNodeId.isJobType()) {
+      throw new IllegalArgumentException("Only job expected, got " + parentJobNodeId.getValue());
+    }
+    throwIfNotExists(parentJobNodeId);
+    return Response.ok(lineageService.parentDirectLineage(parentJobNodeId.asJobId())).build();
+  }
+
   @Timed
   @ResponseMetered
   @ExceptionMetered

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

@@ -183,7 +183,7 @@ public static String stringOrNull(final ResultSet results, final String column)
   public static String stringOrThrow(final ResultSet results, final String column)
       throws SQLException {
     if (results.getObject(column) == null) {
-      throw new IllegalArgumentException();
+      throw new IllegalArgumentException("no column found for " + column);
     }
     return results.getString(column);
   }

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

@@ -10,23 +10,42 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
+
+import org.jdbi.v3.sqlobject.config.RegisterRowMapper;
+import org.jdbi.v3.sqlobject.customizer.BindList;
+import org.jdbi.v3.sqlobject.statement.SqlQuery;
+
+import marquez.common.models.DatasetId;
+import marquez.common.models.JobId;
 import marquez.db.mappers.DatasetDataMapper;
 import marquez.db.mappers.JobDataMapper;
 import marquez.db.mappers.JobRowMapper;
 import marquez.db.mappers.RunMapper;
+import marquez.db.mappers.SimpleLineageEdgeMapper;
 import marquez.service.models.DatasetData;
 import marquez.service.models.JobData;
 import marquez.service.models.Run;
-import org.jdbi.v3.sqlobject.config.RegisterRowMapper;
-import org.jdbi.v3.sqlobject.customizer.BindList;
-import org.jdbi.v3.sqlobject.statement.SqlQuery;
 
 @RegisterRowMapper(DatasetDataMapper.class)
 @RegisterRowMapper(JobDataMapper.class)
 @RegisterRowMapper(RunMapper.class)
 @RegisterRowMapper(JobRowMapper.class)
+@RegisterRowMapper(SimpleLineageEdgeMapper.class)
 public interface LineageDao {
 
+  public record SimpleLineage(Collection<SimpleLineageEdge> edges) {
+  }
+
+  public record SimpleLineageEdge(
+      JobId job1,
+      String direction,
+      DatasetId dataset,
+      String direction2,
+      JobId job2,
+      JobId job2parent
+  ) {
+
+  }
   /**
    * Fetch all of the jobs that consume or produce the datasets that are consumed or produced by the
    * input jobIds. This returns a single layer from the BFS using datasets as edges. Jobs that have
@@ -79,6 +98,38 @@ SELECT DISTINCT ON (j.uuid) j.*, inputs AS input_uuids, outputs AS output_uuids
   """)
   Set<JobData> getLineage(@BindList Set<UUID> jobIds, int depth);
 
+  /**
+   *  1 level of lineage for all the children jobs of the given parent
+   *
+   * @param parentJobNamespace the namespace of the parent
+   * @param parentJobName the name of the parent
+   * @return edges form job to dataset to job
+   */
+  @SqlQuery(
+      """
+      SELECT
+          jobs.namespace_name AS job_namespace, jobs.simple_name AS job_name,
+          jvim.io_type AS io1,
+          d.namespace_name AS ds_namespace, d."name" AS ds_name,
+          jvim2.io_type AS io2,
+          jv2.namespace_name AS job2_namespace, jv2.job_name  AS job2_name,
+          pj.namespace_name AS job2_parent_namespace, pj.simple_name AS job2_parent_name
+      FROM jobs
+      INNER JOIN job_versions jv ON jv.uuid = jobs.current_version_uuid
+      LEFT JOIN job_versions_io_mapping jvim ON jvim.job_version_uuid = jobs.current_version_uuid
+      LEFT JOIN datasets d ON d.uuid = jvim.dataset_uuid
+      LEFT JOIN job_versions_io_mapping jvim2 ON jvim2.dataset_uuid = d.uuid AND jvim2.job_version_uuid <> jvim.job_version_uuid and jvim2.io_type <> jvim.io_type
+      LEFT JOIN job_versions jv2 ON jv2.uuid = jvim2.job_version_uuid
+      LEFT JOIN jobs j2 ON jv2.job_uuid = j2.uuid
+      LEFT JOIN jobs pj ON j2.parent_job_uuid = pj.uuid
+      WHERE jobs.parent_job_uuid IN (
+            SELECT uuid AS parent_job_uuid
+            FROM jobs
+            WHERE namespace_name=:parentJobNamespace and simple_name=:parentJobName
+      );
+  """)
+  Collection<SimpleLineageEdge> getDirectLineageFromParent(String parentJobNamespace, String parentJobName);
+
   @SqlQuery(
       """
       SELECT ds.*, dv.fields, dv.lifecycle_state

api/src/main/java/marquez/db/mappers/SimpleLineageEdgeMapper.java

@@ -0,0 +1,36 @@
+package marquez.db.mappers;
+
+import static marquez.db.Columns.stringOrNull;
+import static marquez.db.Columns.stringOrThrow;
+
+import java.sql.ResultSet;
+import java.sql.SQLException;
+
+import org.jdbi.v3.core.mapper.RowMapper;
+import org.jdbi.v3.core.statement.StatementContext;
+
+import lombok.NonNull;
+import marquez.common.models.DatasetId;
+import marquez.common.models.DatasetName;
+import marquez.common.models.JobId;
+import marquez.common.models.JobName;
+import marquez.common.models.NamespaceName;
+import marquez.db.LineageDao.SimpleLineageEdge;
+
+public final class SimpleLineageEdgeMapper implements RowMapper<SimpleLineageEdge> {
+  @Override
+  public SimpleLineageEdge map(@NonNull ResultSet results, @NonNull StatementContext context)
+      throws SQLException {
+    JobId job1 = JobId.of(NamespaceName.of(stringOrThrow(results, "job_namespace")), JobName.of(stringOrThrow(results, "job_name")));
+    String io1 = stringOrNull(results, "io1");
+    String ds_namespace = stringOrNull(results, "ds_namespace");
+    DatasetId ds = ds_namespace == null ? null : new DatasetId(NamespaceName.of(ds_namespace), DatasetName.of(stringOrNull(results, "ds_name")));
+    String io2 = stringOrNull(results, "io2");
+    String job2_namespace = stringOrNull(results, "job2_namespace");
+    JobId job2 = job2_namespace == null ? null : JobId.of(NamespaceName.of(job2_namespace), JobName.of(stringOrThrow(results, "job2_name")));
+    String job2parent_namespace = stringOrNull(results, "job2_parent_namespace");
+    JobId job2parent = job2parent_namespace == null ? null : JobId.of(NamespaceName.of(job2parent_namespace), JobName.of(stringOrThrow(results, "job2_parent_name")));
+    return new  SimpleLineageEdge(job1, io1, ds, io2, job2, job2parent);
+  }
+}
+

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

@@ -5,10 +5,12 @@
 
 package marquez.service;
 
-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 static java.util.stream.Collectors.filtering;
+import static java.util.stream.Collectors.groupingBy;
+import static java.util.stream.Collectors.mapping;
+import static java.util.stream.Collectors.toList;
+
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -21,6 +23,12 @@
 import java.util.UUID;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+
+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 lombok.NonNull;
 import lombok.extern.slf4j.Slf4j;
 import marquez.common.models.DatasetId;
@@ -41,13 +49,73 @@
 
 @Slf4j
 public class LineageService extends DelegatingLineageDao {
+
+  public record JobWithParent(JobId job, JobId parent) {
+  }
+
+  public record DatasetLineage(DatasetId dataset, Collection<JobWithParent> consumers, Collection<JobWithParent> producers) {
+  }
+
+  public record ChildLineage(JobId job, Collection<DatasetLineage> inputs, Collection<DatasetLineage> outputs) {
+  }
+
+  public record ParentLineage(JobId parent, Collection<ChildLineage> children) {
+  }
+
   private final JobDao jobDao;
 
   public LineageService(LineageDao delegate, JobDao jobDao) {
     super(delegate);
     this.jobDao = jobDao;
   }
 
+  /**
+   * This method is specialized for returning one level of lineage from a parent job.
+   * It finds all the children of the provided parent node
+   * It then finds the input and output datasets those children write to.
+   * It finally returns the other jobs consuming or producing those datasets (and their parent).
+   * @param parentJobId the parent job
+   * @return 1 level of lineage for all the children jobs of the given parent
+   */
+  public ParentLineage parentDirectLineage(JobId parentJobId) {
+    log.debug("Attempting to get lineage for parent job '{}'", parentJobId);
+
+    Collection<SimpleLineageEdge> directLineageFromParent =
+        getDirectLineageFromParent(parentJobId.getNamespace().getValue(), parentJobId.getName().getValue());
+
+
+    Map<JobId, Map<String, Map<DatasetId, Map<String, List<JobWithParent>>>>> grouped =
+        directLineageFromParent.stream().collect(
+        groupingBy(SimpleLineageEdge::job1,
+            filtering(e -> e.direction() != null,
+            groupingBy(SimpleLineageEdge::direction,
+                filtering(e -> e.dataset() != null,
+                groupingBy(SimpleLineageEdge::dataset,
+                    filtering(e -> e.direction2() != null,
+                    groupingBy(SimpleLineageEdge::direction2,
+                        mapping(e -> new JobWithParent(e.job2(), e.job2parent()),
+                        toList())))))))));
+
+    List<ChildLineage> children = grouped.entrySet().stream().map(
+        e -> new ChildLineage(
+              e.getKey(),
+              toDatasetLineages(e.getValue().get("INPUT")),
+              toDatasetLineages(e.getValue().get("OUTPUT"))
+            )
+      ).collect(toList());
+    return new ParentLineage(parentJobId, children);
+  }
+
+  private Collection<DatasetLineage> toDatasetLineages(Map<DatasetId, Map<String, List<JobWithParent>>> datasets) {
+    return datasets == null ? null : datasets.entrySet().stream().map(
+        e -> new DatasetLineage(
+              e.getKey(),
+              e.getValue().get("INPUT"),
+              e.getValue().get("OUTPUT")
+            )
+      ).collect(toList());
+  }
+
   // 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);

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

@@ -79,6 +79,7 @@ public JavaType typeFromId(DatabindContext context, String id) throws IOExceptio
             .filter(s -> s.getName().equals(type))
             .findAny()
             .map(EventSchemaURL::getSubType)
+            .map(p -> (Class)p)
             .orElse(LINEAGE_EVENT.subType);
 
     return context.constructSpecializedType(superType, subType);