Skip to content

Commit 0f48ac2

Browse files
Kess220Kess220
authored
Merge pull request #43 from nubank/lineage-api-downstream
[DGD-4893] Bring downstream data
2 parents 740cc2b + c1291d2 commit 0f48ac2

3 files changed

Lines changed: 92 additions & 60 deletions

File tree

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import marquez.service.models.JobData;
2626
import marquez.service.models.Run;
2727
import org.jdbi.v3.sqlobject.config.RegisterRowMapper;
28+
import org.jdbi.v3.sqlobject.customizer.Bind;
2829
import org.jdbi.v3.sqlobject.customizer.BindList;
2930
import org.jdbi.v3.sqlobject.statement.SqlQuery;
3031

@@ -282,4 +283,17 @@ WHERE j.uuid IN (<jobIds>) OR j.symlink_target_uuid IN (<jobIds>)
282283
default Set<JobData> getDirectLineage(@BindList Set<UUID> jobIds, int depth) {
283284
throw new UnsupportedOperationException("Use getDirectLineage and iterate in LineageService.");
284285
}
286+
287+
@SqlQuery("""
288+
SELECT DISTINCT j.uuid
289+
FROM jobs j
290+
INNER JOIN job_versions jv ON jv.job_uuid = j.uuid
291+
INNER JOIN job_versions_io_mapping io ON io.job_version_uuid = jv.uuid
292+
INNER JOIN datasets_view ds ON ds.uuid = io.dataset_uuid
293+
WHERE ds.name = :datasetName
294+
AND ds.namespace_name = :namespaceName
295+
AND io.io_type = 'INPUT'
296+
""")
297+
Set<UUID> getDownstreamJobs(@Bind("namespaceName") String namespaceName, @Bind("datasetName") String datasetName);
298+
285299
}

‎api/src/main/java/marquez/service/DelegatingDaos.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,8 @@ public static class DelegatingTagDao implements TagDao {
9191

9292
@RequiredArgsConstructor
9393
public static class DelegatingLineageDao implements LineageDao {
94-
@Delegate private final LineageDao delegate;
94+
@Delegate
95+
protected final LineageDao delegate;
9596
}
9697

9798
@RequiredArgsConstructor

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

Lines changed: 76 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,6 @@ public LineageService(LineageDao delegate, JobDao jobDao, RunDao runDao) {
6868
// TODO make input parameters easily extendable if adding more options like
6969
// 'withJobFacets'
7070
public Lineage lineage(NodeId nodeId, int depth) {
71-
log.debug("Attempting to get lineage for node '{}' with depth '{}'", nodeId.getValue(), depth);
7271
Optional<UUID> optionalUUID = getJobUuid(nodeId);
7372
if (optionalUUID.isEmpty()) {
7473
log.warn(
@@ -77,7 +76,6 @@ public Lineage lineage(NodeId nodeId, int depth) {
7776
return toLineageWithOrphanDataset(nodeId.asDatasetId());
7877
}
7978
UUID job = optionalUUID.get();
80-
log.debug("Attempting to get lineage for job '{}'", job);
8179
Set<JobData> jobData = getLineage(Collections.singleton(job), depth);
8280

8381
// Ensure job data is not empty, an empty set cannot be passed to
@@ -310,11 +308,10 @@ public UpstreamRunLineage upstream(@NotNull RunId runId, int depth) {
310308
*/
311309
public Lineage directLineage(NodeId nodeId, int depth) {
312310
depth += 1;
313-
log.debug("Attempting to get lineage for node '{}' with depth '{}'", nodeId.getValue(), depth);
314311
Optional<UUID> optionalUUID = getJobUuid(nodeId);
315312
if (optionalUUID.isEmpty()) {
316-
log.warn("Failed to get job for node '{}', returning orphan dataset...", nodeId.getValue());
317-
return toLineageWithOrphanDataset(nodeId.asDatasetId());
313+
log.warn("Failed to get job for node '{}', returning orphan dataset...", nodeId.getValue());
314+
return toLineageWithOrphanDataset(nodeId.asDatasetId());
318315
}
319316
UUID job = optionalUUID.get();
320317

@@ -324,61 +321,64 @@ public Lineage directLineage(NodeId nodeId, int depth) {
324321
Set<JobData> allJobData = new HashSet<>();
325322

326323
for (int level = 0; level < depth && !pending.isEmpty(); level++) {
327-
// Step 1: fetch direct lineage for the current pending job set
328-
Set<JobData> directLineage = getDirectLineage(pending);
329-
allJobData.addAll(directLineage);
330-
331-
// Collect all job UUIDs we discover this iteration
332-
Set<UUID> nextJobs = new HashSet<>();
333-
visited.addAll(pending);
334-
pending.clear();
335-
336-
// Step 2: for each job in current iteration:
337-
for (JobData jd : directLineage) {
338-
// Step 2a: gather any dataset inputs/outputs
339-
Set<UUID> dsUuids = new HashSet<>();
340-
dsUuids.addAll(jd.getInputUuids());
341-
dsUuids.addAll(jd.getOutputUuids());
342-
343-
// Fetch all dataset info once
344-
Set<DatasetData> dsList = getDatasetData(dsUuids);
345-
346-
// Step 2b: for each dataset, find the job(s) associated with it
347-
for (UUID dsUuid : dsUuids) {
348-
DatasetData ds = dsList.stream()
349-
.filter(d -> d.getUuid().equals(dsUuid))
350-
.findFirst()
351-
.orElse(null);
352-
if (ds == null) {
353-
continue;
354-
}
355-
// Only follow upstream/downstream direction based on whether the dataset
356-
// is an input or output
357-
if (jd.getInputUuids().contains(dsUuid)) {
358-
// IMPORTANT: Use the string values
359-
Optional<UUID> maybeJob = getJobFromInputOrOutput(
360-
ds.getName().getValue(),
361-
ds.getNamespace().getValue());
362-
maybeJob.ifPresent(nextJobs::add);
363-
}
324+
// Step 1: fetch direct lineage for the current pending job set
325+
Set<JobData> directLineage = getDirectLineage(pending);
326+
allJobData.addAll(directLineage);
327+
328+
// Collect all job UUIDs we discover this iteration
329+
Set<UUID> nextJobs = new HashSet<>();
330+
visited.addAll(pending);
331+
pending.clear();
332+
333+
// Step 2: for each job in current iteration:
334+
for (JobData jd : directLineage) {
335+
// Step 2a: gather any dataset inputs/outputs
336+
Set<UUID> dsUuids = new HashSet<>();
337+
dsUuids.addAll(jd.getInputUuids());
338+
dsUuids.addAll(jd.getOutputUuids());
339+
340+
// Fetch all dataset info once
341+
Set<DatasetData> dsList = getDatasetData(dsUuids);
342+
343+
// Step 2b: for each dataset, find the job(s) associated with it
344+
for (UUID dsUuid : dsUuids) {
345+
DatasetData ds = dsList.stream()
346+
.filter(d -> d.getUuid().equals(dsUuid))
347+
.findFirst()
348+
.orElse(null);
349+
if (ds == null) {
350+
continue;
351+
}
352+
353+
// Follow upstream (jobs producing this dataset)
354+
if (jd.getInputUuids().contains(dsUuid)) {
355+
Set<UUID> upstreamJobs = getUpstreamJobs(ds);
356+
nextJobs.addAll(upstreamJobs);
357+
}
358+
359+
// Follow downstream (jobs consuming this dataset)
360+
if (jd.getOutputUuids().contains(dsUuid)) {
361+
Set<UUID> downstreamJobs = getDownstreamJobs(ds);
362+
nextJobs.addAll(downstreamJobs);
363+
}
364+
}
364365
}
365-
}
366366

367-
// Step 3: remove visited jobs
368-
nextJobs.removeAll(visited);
367+
// Step 3: remove visited jobs
368+
nextJobs.removeAll(visited);
369369

370-
// Step 4: pending becomes nextJobs
371-
pending.addAll(nextJobs);
370+
// Step 4: pending becomes nextJobs
371+
pending.addAll(nextJobs);
372372
}
373373

374374
if (allJobData.isEmpty()) {
375-
log.warn("No lineage for job '{}' node '{}', returning orphan dataset...", job, nodeId.getValue());
376-
return toLineageWithOrphanDataset(nodeId.asDatasetId());
375+
log.warn("No lineage for job '{}' node '{}', returning orphan dataset...", job, nodeId.getValue());
376+
return toLineageWithOrphanDataset(nodeId.asDatasetId());
377377
}
378378

379379
// Enrich with run data
380380
for (JobData j : allJobData) {
381-
runDao.findRunByUuid(j.getCurrentRunUuid()).ifPresent(j::setLatestRun);
381+
runDao.findRunByUuid(j.getCurrentRunUuid()).ifPresent(j::setLatestRun);
382382
}
383383

384384
// Retrieve all datasets from the discovered lineage
@@ -389,21 +389,38 @@ public Lineage directLineage(NodeId nodeId, int depth) {
389389
? new HashSet<>()
390390
: getDatasetData(datasetIds);
391391

392-
// If this started at a dataset node, ensure it’s included in the final lineage
392+
// If this started at a dataset node, ensure it's included in the final lineage
393393
if (nodeId.isDatasetType()) {
394-
DatasetId datasetId = nodeId.asDatasetId();
395-
DatasetData datasetData = getDatasetData(
396-
datasetId.getNamespace().getValue(),
397-
datasetId.getName().getValue());
398-
if (!datasetIds.contains(datasetData.getUuid())) {
399-
log.warn("Found jobs with no lineage for dataset '{}', discarding", nodeId.getValue());
400-
return toLineageWithOrphanDataset(datasetId);
401-
}
394+
DatasetId datasetId = nodeId.asDatasetId();
395+
DatasetData datasetData = getDatasetData(
396+
datasetId.getNamespace().getValue(),
397+
datasetId.getName().getValue());
398+
if (!datasetIds.contains(datasetData.getUuid())) {
399+
log.warn("Found jobs with no lineage for dataset '{}', discarding", nodeId.getValue());
400+
return toLineageWithOrphanDataset(datasetId);
401+
}
402402
}
403403

404+
405+
404406
return toLineage(allJobData, datasets);
405407
}
406408

409+
private Set<UUID> getUpstreamJobs(DatasetData dataset) {
410+
return delegate.getJobFromInputOrOutput(
411+
dataset.getName().getValue(),
412+
dataset.getNamespace().getValue()
413+
).map(Collections::singleton)
414+
.orElse(Collections.emptySet());
415+
}
416+
417+
private Set<UUID> getDownstreamJobs(DatasetData dataset) {
418+
return delegate.getDownstreamJobs(
419+
dataset.getNamespace().getValue(),
420+
dataset.getName().getValue()
421+
);
422+
}
423+
407424
/**
408425
* Fetch dataset information for the given set of dataset UUIDs.
409426
*/

0 commit comments

Comments
 (0)