Skip to content

Commit 7992863

Browse files
author
Jonathan P Moraes
authored
Lineage endpoint (#40)
* Adding direct lineage endpoint
1 parent 2f74015 commit 7992863

3 files changed

Lines changed: 347 additions & 159 deletions

File tree

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,4 +175,18 @@ static class Events {
175175

176176
int totalCount;
177177
}
178+
179+
@Timed
180+
@ResponseMetered
181+
@ExceptionMetered
182+
@GET
183+
@Consumes(APPLICATION_JSON)
184+
@Produces(APPLICATION_JSON)
185+
@Path("/lineage/direct")
186+
public Response getDirectLineage(
187+
@QueryParam("nodeId") @NotNull NodeId nodeId,
188+
@QueryParam("depth") @DefaultValue(DEFAULT_DEPTH) int depth) {
189+
throwIfNotExists(nodeId);
190+
return Response.ok(lineageService.directLineage(nodeId, depth)).build();
191+
}
178192
}

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

Lines changed: 150 additions & 104 deletions
Original file line numberDiff line numberDiff line change
@@ -35,78 +35,82 @@
3535
@RegisterRowMapper(UpstreamRunRowMapper.class)
3636
public interface LineageDao {
3737

38-
public record JobSummary(NamespaceName namespace, JobName name, UUID version) {}
38+
public record JobSummary(NamespaceName namespace, JobName name, UUID version) {
39+
}
3940

40-
public record RunSummary(RunId id, Instant start, Instant end, String status) {}
41+
public record RunSummary(RunId id, Instant start, Instant end, String status) {
42+
}
4143

4244
public record DatasetSummary(
43-
NamespaceName namespace, DatasetName name, UUID version, RunId producedByRunId) {}
45+
NamespaceName namespace, DatasetName name, UUID version, RunId producedByRunId) {
46+
}
4447

45-
public record UpstreamRunRow(JobSummary job, RunSummary run, DatasetSummary input) {}
48+
public record UpstreamRunRow(JobSummary job, RunSummary run, DatasetSummary input) {
49+
}
4650

4751
/**
48-
* Fetch all of the jobs that consume or produce the datasets that are consumed or produced by the
49-
* input jobIds. This returns a single layer from the BFS using datasets as edges. Jobs that have
50-
* no input or output datasets will have no results. Jobs that have no upstream producers or
52+
* Fetch all of the jobs that consume or produce the datasets that are consumed
53+
* or produced by the
54+
* input jobIds. This returns a single layer from the BFS using datasets as
55+
* edges. Jobs that have
56+
* no input or output datasets will have no results. Jobs that have no upstream
57+
* producers or
5158
* downstream consumers will have the original jobIds returned.
5259
*
5360
* @param jobIds
5461
* @return
5562
*/
56-
@SqlQuery(
57-
"""
58-
WITH RECURSIVE
59-
job_io AS (
60-
SELECT
61-
io.job_uuid AS job_uuid,
62-
io.job_symlink_target_uuid AS job_symlink_target_uuid,
63-
ARRAY_AGG(DISTINCT io.dataset_uuid) FILTER (WHERE io.io_type='INPUT') AS inputs,
64-
ARRAY_AGG(DISTINCT io.dataset_uuid) FILTER (WHERE io.io_type='OUTPUT') AS outputs
65-
FROM job_versions_io_mapping io
66-
WHERE io.is_current_job_version = TRUE
67-
GROUP BY io.job_symlink_target_uuid, io.job_uuid
68-
),
69-
lineage(job_uuid, job_symlink_target_uuid, inputs, outputs) AS (
70-
SELECT job_uuid,
71-
job_symlink_target_uuid,
72-
COALESCE(inputs, Array[]::uuid[]) AS inputs,
73-
COALESCE(outputs, Array[]::uuid[]) AS outputs,
74-
0 AS depth
75-
FROM job_io
76-
WHERE job_uuid IN (<jobIds>) OR job_symlink_target_uuid IN (<jobIds>)
77-
UNION
78-
SELECT io.job_uuid, io.job_symlink_target_uuid, io.inputs, io.outputs, l.depth + 1
79-
FROM job_io io, lineage l
80-
WHERE (io.job_uuid != l.job_uuid) AND
81-
array_cat(io.inputs, io.outputs) && array_cat(l.inputs, l.outputs)
82-
AND depth < :depth),
83-
lineage_outside_job_io(job_uuid) AS (
84-
SELECT
85-
param_jobs.param_job_uuid as job_uuid,
86-
j.symlink_target_uuid,
87-
Array[]::uuid[] AS inputs,
88-
Array[]::uuid[] AS outputs,
89-
0 AS depth
90-
FROM (SELECT unnest(ARRAY[<jobIds>]::UUID[]) AS param_job_uuid) param_jobs
91-
LEFT JOIN lineage l on param_jobs.param_job_uuid = l.job_uuid
92-
INNER JOIN jobs j ON j.uuid = param_jobs.param_job_uuid
93-
WHERE l.job_uuid IS NULL
94-
)
95-
SELECT DISTINCT ON (j.uuid) j.*, inputs AS input_uuids, outputs AS output_uuids
96-
FROM (SELECT * FROM lineage UNION SELECT * FROM lineage_outside_job_io) l2
97-
INNER JOIN jobs_view j ON (j.uuid=l2.job_uuid OR j.uuid=l2.job_symlink_target_uuid)
98-
""")
63+
@SqlQuery("""
64+
WITH RECURSIVE
65+
job_io AS (
66+
SELECT
67+
io.job_uuid AS job_uuid,
68+
io.job_symlink_target_uuid AS job_symlink_target_uuid,
69+
ARRAY_AGG(DISTINCT io.dataset_uuid) FILTER (WHERE io.io_type='INPUT') AS inputs,
70+
ARRAY_AGG(DISTINCT io.dataset_uuid) FILTER (WHERE io.io_type='OUTPUT') AS outputs
71+
FROM job_versions_io_mapping io
72+
WHERE io.is_current_job_version = TRUE
73+
GROUP BY io.job_symlink_target_uuid, io.job_uuid
74+
),
75+
lineage(job_uuid, job_symlink_target_uuid, inputs, outputs) AS (
76+
SELECT job_uuid,
77+
job_symlink_target_uuid,
78+
COALESCE(inputs, Array[]::uuid[]) AS inputs,
79+
COALESCE(outputs, Array[]::uuid[]) AS outputs,
80+
0 AS depth
81+
FROM job_io
82+
WHERE job_uuid IN (<jobIds>) OR job_symlink_target_uuid IN (<jobIds>)
83+
UNION
84+
SELECT io.job_uuid, io.job_symlink_target_uuid, io.inputs, io.outputs, l.depth + 1
85+
FROM job_io io, lineage l
86+
WHERE (io.job_uuid != l.job_uuid) AND
87+
array_cat(io.inputs, io.outputs) && array_cat(l.inputs, l.outputs)
88+
AND depth < :depth),
89+
lineage_outside_job_io(job_uuid) AS (
90+
SELECT
91+
param_jobs.param_job_uuid as job_uuid,
92+
j.symlink_target_uuid,
93+
Array[]::uuid[] AS inputs,
94+
Array[]::uuid[] AS outputs,
95+
0 AS depth
96+
FROM (SELECT unnest(ARRAY[<jobIds>]::UUID[]) AS param_job_uuid) param_jobs
97+
LEFT JOIN lineage l on param_jobs.param_job_uuid = l.job_uuid
98+
INNER JOIN jobs j ON j.uuid = param_jobs.param_job_uuid
99+
WHERE l.job_uuid IS NULL
100+
)
101+
SELECT DISTINCT ON (j.uuid) j.*, inputs AS input_uuids, outputs AS output_uuids
102+
FROM (SELECT * FROM lineage UNION SELECT * FROM lineage_outside_job_io) l2
103+
INNER JOIN jobs_view j ON (j.uuid=l2.job_uuid OR j.uuid=l2.job_symlink_target_uuid)
104+
""")
99105
Set<JobData> getLineage(@BindList Set<UUID> jobIds, int depth);
100106

101-
@SqlQuery(
102-
"""
103-
SELECT j.*, NULL as input_uuids, NULL AS output_uuids FROM jobs_view j
104-
WHERE j.parent_job_uuid= :jobId
105-
LIMIT 1""")
107+
@SqlQuery("""
108+
SELECT j.*, NULL as input_uuids, NULL AS output_uuids FROM jobs_view j
109+
WHERE j.parent_job_uuid= :jobId
110+
LIMIT 1""")
106111
Optional<JobData> getParentJobData(UUID jobId);
107112

108-
@SqlQuery(
109-
"""
113+
@SqlQuery("""
110114
SELECT ds.*, dv.fields, dv.lifecycle_state
111115
FROM datasets_view ds
112116
LEFT JOIN dataset_versions dv ON dv.uuid = ds.current_version_uuid
@@ -115,8 +119,7 @@ INNER JOIN jobs_view j ON (j.uuid=l2.job_uuid OR j.uuid=l2.job_symlink_target_uu
115119
AND ds.uuid IN (<dsUuids>)""")
116120
Set<DatasetData> getDatasetData(@BindList Set<UUID> dsUuids);
117121

118-
@SqlQuery(
119-
"""
122+
@SqlQuery("""
120123
SELECT ds.*, dv.fields, dv.lifecycle_state
121124
FROM datasets_view ds
122125
LEFT JOIN dataset_versions dv on dv.uuid = ds.current_version_uuid
@@ -126,8 +129,7 @@ AND ds.uuid IN (<dsUuids>)""")
126129
AND CAST((:namespaceName, :datasetName) AS DATASET_NAME) = ANY(d.dataset_symlinks)""")
127130
DatasetData getDatasetData(String namespaceName, String datasetName);
128131

129-
@SqlQuery(
130-
"""
132+
@SqlQuery("""
131133
SELECT j.uuid FROM jobs j
132134
INNER JOIN job_versions jv ON jv.job_uuid = j.uuid
133135
INNER JOIN job_versions_io_mapping io ON io.job_version_uuid = jv.uuid
@@ -137,49 +139,47 @@ AND CAST((:namespaceName, :datasetName) AS DATASET_NAME) = ANY(d.dataset_symlink
137139
LIMIT 1""")
138140
Optional<UUID> getJobFromInputOrOutput(String datasetName, String namespaceName);
139141

140-
@SqlQuery(
141-
"""
142-
WITH latest_runs AS (
143-
SELECT DISTINCT on(r.job_name, r.namespace_name) r.*, jv.version
144-
FROM runs_view r
145-
INNER JOIN job_versions jv ON jv.uuid=r.job_version_uuid
146-
INNER JOIN jobs_view j ON j.uuid=jv.job_uuid
147-
WHERE j.uuid in (<jobUuid>) OR j.symlink_target_uuid IN (<jobUuid>)
148-
ORDER BY r.job_name, r.namespace_name, created_at DESC
149-
)
150-
SELECT r.*, ra.args, f.facets,
151-
r.version AS job_version, ri.input_versions, ro.output_versions
152-
from latest_runs AS r
153-
LEFT JOIN run_args AS ra ON ra.uuid = r.run_args_uuid
154-
LEFT JOIN LATERAL (
155-
SELECT im.run_uuid,
156-
JSON_AGG(json_build_object('namespace', dv.namespace_name,
157-
'name', dv.dataset_name,
158-
'version', dv.version)) AS input_versions
159-
FROM runs_input_mapping im
160-
INNER JOIN dataset_versions dv on im.dataset_version_uuid = dv.uuid
161-
WHERE im.run_uuid=r.uuid
162-
GROUP BY im.run_uuid
163-
) ri ON ri.run_uuid=r.uuid
164-
LEFT JOIN LATERAL (
165-
SELECT rf.run_uuid, JSON_AGG(rf.facet ORDER BY rf.lineage_event_time ASC) AS facets
166-
FROM run_facets_view AS rf
167-
WHERE rf.run_uuid=r.uuid
168-
GROUP BY rf.run_uuid
169-
) AS f ON r.uuid=f.run_uuid
170-
LEFT JOIN LATERAL (
171-
SELECT run_uuid, JSON_AGG(json_build_object('namespace', namespace_name,
172-
'name', dataset_name,
173-
'version', version)) AS output_versions
174-
FROM dataset_versions
175-
WHERE run_uuid=r.uuid
176-
GROUP BY run_uuid
177-
) ro ON ro.run_uuid=r.uuid
178-
""")
142+
@SqlQuery("""
143+
WITH latest_runs AS (
144+
SELECT DISTINCT on(r.job_name, r.namespace_name) r.*, jv.version
145+
FROM runs_view r
146+
INNER JOIN job_versions jv ON jv.uuid=r.job_version_uuid
147+
INNER JOIN jobs_view j ON j.uuid=jv.job_uuid
148+
WHERE j.uuid in (<jobUuid>) OR j.symlink_target_uuid IN (<jobUuid>)
149+
ORDER BY r.job_name, r.namespace_name, created_at DESC
150+
)
151+
SELECT r.*, ra.args, f.facets,
152+
r.version AS job_version, ri.input_versions, ro.output_versions
153+
from latest_runs AS r
154+
LEFT JOIN run_args AS ra ON ra.uuid = r.run_args_uuid
155+
LEFT JOIN LATERAL (
156+
SELECT im.run_uuid,
157+
JSON_AGG(json_build_object('namespace', dv.namespace_name,
158+
'name', dv.dataset_name,
159+
'version', dv.version)) AS input_versions
160+
FROM runs_input_mapping im
161+
INNER JOIN dataset_versions dv on im.dataset_version_uuid = dv.uuid
162+
WHERE im.run_uuid=r.uuid
163+
GROUP BY im.run_uuid
164+
) ri ON ri.run_uuid=r.uuid
165+
LEFT JOIN LATERAL (
166+
SELECT rf.run_uuid, JSON_AGG(rf.facet ORDER BY rf.lineage_event_time ASC) AS facets
167+
FROM run_facets_view AS rf
168+
WHERE rf.run_uuid=r.uuid
169+
GROUP BY rf.run_uuid
170+
) AS f ON r.uuid=f.run_uuid
171+
LEFT JOIN LATERAL (
172+
SELECT run_uuid, JSON_AGG(json_build_object('namespace', namespace_name,
173+
'name', dataset_name,
174+
'version', version)) AS output_versions
175+
FROM dataset_versions
176+
WHERE run_uuid=r.uuid
177+
GROUP BY run_uuid
178+
) ro ON ro.run_uuid=r.uuid
179+
""")
179180
List<Run> getCurrentRunsWithFacets(@BindList Collection<UUID> jobUuid);
180181

181-
@SqlQuery(
182-
"""
182+
@SqlQuery("""
183183
WITH latest_runs AS (SELECT current_run_uuid, current_version_uuid AS job_version
184184
FROM jobs j
185185
WHERE j.uuid in (<jobUuid>) OR j.symlink_target_uuid IN (<jobUuid>))
@@ -190,8 +190,7 @@ WHERE j.uuid in (<jobUuid>) OR j.symlink_target_uuid IN (<jobUuid>))
190190
""")
191191
List<Run> getCurrentRuns(@BindList Collection<UUID> jobUuid);
192192

193-
@SqlQuery(
194-
"""
193+
@SqlQuery("""
195194
WITH RECURSIVE
196195
upstream_runs(
197196
r_uuid, -- run uuid
@@ -236,4 +235,51 @@ SELECT DISTINCT ON (upstream_runs.r_uuid, upstream_runs.dataset_version_uuid, up
236235
ORDER BY depth ASC, job_name ASC;
237236
""")
238237
List<UpstreamRunRow> getUpstreamRuns(@NotNull UUID runId, int depth);
238+
239+
/**
240+
* Fetch all of the jobs that consume or produce the datasets that are consumed
241+
* or produced by the
242+
* input jobIds. This returns a single layer from the BFS using datasets as
243+
* edges. Jobs that have
244+
* no input or output datasets will have no results. Jobs that have no upstream
245+
* producers or
246+
* downstream consumers will have the original jobIds returned.
247+
* @param jobIds
248+
* @return
249+
*/
250+
@SqlQuery("""
251+
SELECT DISTINCT j.*,
252+
job_io.inputs AS input_uuids,
253+
job_io.outputs AS output_uuids
254+
FROM (
255+
SELECT
256+
io.job_uuid AS job_uuid,
257+
io.job_symlink_target_uuid AS job_symlink_target_uuid,
258+
ARRAY_AGG(DISTINCT io.dataset_uuid) FILTER (WHERE io_type='INPUT') AS inputs,
259+
ARRAY_AGG(DISTINCT io.dataset_uuid) FILTER (WHERE io_type='OUTPUT') AS outputs
260+
FROM job_versions_io_mapping io
261+
WHERE io.is_current_job_version = TRUE
262+
GROUP BY io.job_symlink_target_uuid, io.job_uuid
263+
) job_io
264+
INNER JOIN jobs_view j
265+
ON (j.uuid=job_io.job_uuid OR j.uuid=job_io.job_symlink_target_uuid)
266+
WHERE j.uuid IN (<jobIds>) OR j.symlink_target_uuid IN (<jobIds>)
267+
""")
268+
Set<JobData> getDirectLineage(@BindList Set<UUID> jobIds);
269+
270+
/**
271+
* Fetch all of the jobs that consume or produce the datasets that are consumed
272+
* or produced by the
273+
* input jobIds. This returns a single layer from the BFS using datasets as
274+
* edges. Jobs that have
275+
* no input or output datasets will have no results. Jobs that have no upstream
276+
* producers or
277+
* downstream consumers will have the original jobIds returned.
278+
*
279+
* @param jobIds
280+
* @return
281+
*/
282+
default Set<JobData> getDirectLineage(@BindList Set<UUID> jobIds, int depth) {
283+
throw new UnsupportedOperationException("Use getDirectLineage and iterate in LineageService.");
284+
}
239285
}

0 commit comments

Comments
 (0)