Skip to content

Commit 5d49c51

Browse files
pawel-big-lebowskipawel-big-lebowski
committed
include column lineage in dataset resource
Signed-off-by: Pawel Leszczynski <leszczynski.pawel@gmail.com>
1 parent 791e1bf commit 5d49c51

19 files changed

Lines changed: 330 additions & 13 deletions

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
* Implemented dataset symlink feature which allows providing multiple names for a dataset and adds edges to lineage graph based on symlinks [`#2066`](https://github.com/MarquezProject/marquez/pull/2066) [@pawel-big-lebowski](https://github.com/pawel-big-lebowski)
77
* Store column lineage facets in separate table [`#2096`](https://github.com/MarquezProject/marquez/pull/2096) [@mzareba382](https://github.com/mzareba382) [@pawel-big-lebowski](https://github.com/pawel-big-lebowski)
88
* Lineage graph endpoint for column lineage [`#2124`](https://github.com/MarquezProject/marquez/pull/2124) [@pawel-big-lebowski](https://github.com/pawel-big-lebowski)
9+
* Enrich returned dataset resource with column lineage information [`#2113`](https://github.com/MarquezProject/marquez/pull/2113) [@pawel-big-lebowski](https://github.com/pawel-big-lebowski)
910

1011
### Fixed
1112
* Add support for `parentRun` facet as reported by older Airflow OpenLineage versions [@collado-mike](https://github.com/collado-mike)

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import com.codahale.metrics.annotation.ResponseMetered;
1313
import com.codahale.metrics.annotation.Timed;
1414
import com.fasterxml.jackson.annotation.JsonProperty;
15+
import java.util.Arrays;
1516
import java.util.List;
1617
import java.util.Locale;
1718
import javax.validation.Valid;
@@ -85,10 +86,11 @@ public Response getDataset(
8586
@PathParam("dataset") DatasetName datasetName) {
8687
throwIfNotExists(namespaceName);
8788

88-
final Dataset dataset =
89+
Dataset dataset =
8990
datasetService
9091
.findWithTags(namespaceName.getValue(), datasetName.getValue())
9192
.orElseThrow(() -> new DatasetNotFoundException(datasetName));
93+
columnLineageService.enrichWithColumnLineage(Arrays.asList(dataset));
9294
return Response.ok(dataset).build();
9395
}
9496

@@ -147,6 +149,7 @@ public Response list(
147149

148150
final List<Dataset> datasets =
149151
datasetService.findAllWithTags(namespaceName.getValue(), limit, offset);
152+
columnLineageService.enrichWithColumnLineage(datasets);
150153
final int totalCount = datasetService.countFor(namespaceName.getValue());
151154
return Response.ok(new ResultsPage<>("datasets", datasets, totalCount)).build();
152155
}

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

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,4 +150,47 @@ Set<ColumnLineageNodeData> getLineage(
150150
int depth,
151151
@BindList(onEmpty = NULL_STRING) List<UUID> datasetFieldUuids,
152152
Instant createdAtUntil);
153+
154+
@SqlQuery(
155+
"""
156+
WITH selected_column_lineage AS (
157+
SELECT cl.*
158+
FROM column_lineage cl
159+
JOIN dataset_fields df ON df.uuid = cl.output_dataset_field_uuid
160+
JOIN datasets_view dv ON dv.uuid = df.dataset_uuid
161+
WHERE ARRAY[<values>]::DATASET_NAME[] && dv.dataset_symlinks
162+
),
163+
dataset_fields_view AS (
164+
SELECT d.namespace_name as namespace_name, d.name as dataset_name, df.name as field_name, df.type, df.uuid
165+
FROM dataset_fields df
166+
INNER JOIN datasets_view d ON d.uuid = df.dataset_uuid
167+
)
168+
SELECT
169+
output_fields.namespace_name,
170+
output_fields.dataset_name,
171+
output_fields.field_name,
172+
output_fields.type,
173+
ARRAY_AGG(ARRAY[input_fields.namespace_name, input_fields.dataset_name, input_fields.field_name]) AS inputFields,
174+
c.transformation_description,
175+
c.transformation_type,
176+
c.created_at,
177+
c.updated_at
178+
FROM selected_column_lineage c
179+
INNER JOIN dataset_fields_view output_fields ON c.output_dataset_field_uuid = output_fields.uuid
180+
LEFT JOIN dataset_fields_view input_fields ON c.input_dataset_field_uuid = input_fields.uuid
181+
GROUP BY
182+
output_fields.namespace_name,
183+
output_fields.dataset_name,
184+
output_fields.field_name,
185+
output_fields.type,
186+
c.transformation_description,
187+
c.transformation_type,
188+
c.created_at,
189+
c.updated_at
190+
""")
191+
Set<ColumnLineageNodeData> getLineageRowsForDatasets(
192+
@BindBeanList(
193+
propertyNames = {"left", "right"},
194+
value = "values")
195+
List<Pair<String, String>> datasets);
153196
}

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

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -34,11 +34,14 @@
3434
@RegisterRowMapper(FieldDataMapper.class)
3535
public interface DatasetFieldDao extends BaseDao {
3636
@SqlQuery(
37-
"SELECT EXISTS ("
38-
+ "SELECT 1 FROM dataset_fields AS df "
39-
+ "INNER JOIN datasets_view AS d "
40-
+ " ON d.uuid = df.dataset_uuid AND d.name = :datasetName AND d.namespace_name = :namespaceName "
41-
+ "WHERE df.name = :name)")
37+
"""
38+
SELECT EXISTS (
39+
SELECT 1 FROM dataset_fields AS df
40+
INNER JOIN datasets_view AS d ON d.uuid = df.dataset_uuid
41+
WHERE CAST((:namespaceName, :datasetName) AS DATASET_NAME) = ANY(d.dataset_symlinks)
42+
AND df.name = :name
43+
)
44+
""")
4245
boolean exists(String namespaceName, String datasetName, String name);
4346

4447
default Dataset updateTags(
@@ -97,20 +100,20 @@ default Dataset updateTags(
97100
"""
98101
SELECT df.uuid
99102
FROM dataset_fields df
100-
INNER JOIN datasets_view AS d
101-
ON d.uuid = df.dataset_uuid AND d.name = :datasetName AND d.namespace_name = :namespace
103+
JOIN datasets_view AS d ON d.uuid = df.dataset_uuid
104+
WHERE CAST((:namespaceName, :datasetName) AS DATASET_NAME) = ANY(d.dataset_symlinks)
102105
""")
103-
List<UUID> findDatasetFieldsUuids(String namespace, String datasetName);
106+
List<UUID> findDatasetFieldsUuids(String namespaceName, String datasetName);
104107

105108
@SqlQuery(
106109
"""
107110
SELECT df.uuid
108111
FROM dataset_fields df
109-
INNER JOIN datasets_view AS d
110-
ON d.uuid = df.dataset_uuid AND d.name = :datasetName AND d.namespace_name = :namespace
111-
WHERE df.name = :name
112+
JOIN datasets_view AS d ON d.uuid = df.dataset_uuid
113+
WHERE CAST((:namespaceName, :datasetName) AS DATASET_NAME) = ANY(d.dataset_symlinks)
114+
AND df.name = :name
112115
""")
113-
Optional<UUID> findUuid(String namespace, String datasetName, String name);
116+
Optional<UUID> findUuid(String namespaceName, String datasetName, String name);
114117

115118
@SqlQuery(
116119
"SELECT f.*, "

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

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,14 @@
2222
import marquez.db.ColumnLineageDao;
2323
import marquez.db.DatasetFieldDao;
2424
import marquez.db.models.ColumnLineageNodeData;
25+
import marquez.service.models.ColumnLineage;
26+
import marquez.service.models.ColumnLineageInputField;
27+
import marquez.service.models.Dataset;
2528
import marquez.service.models.Edge;
2629
import marquez.service.models.Lineage;
2730
import marquez.service.models.Node;
2831
import marquez.service.models.NodeId;
32+
import org.apache.commons.lang3.tuple.Pair;
2933

3034
@Slf4j
3135
public class ColumnLineageService extends DelegatingDaos.DelegatingColumnLineageDao {
@@ -125,4 +129,51 @@ List<UUID> getColumnNodeUuids(NodeId nodeId) {
125129
}
126130
return columnNodeUuids;
127131
}
132+
133+
public void enrichWithColumnLineage(List<Dataset> datasets) {
134+
if (datasets.isEmpty()) {
135+
return;
136+
}
137+
138+
Set<ColumnLineageNodeData> lineageRowsForDatasets =
139+
getLineageRowsForDatasets(
140+
datasets.stream()
141+
.map(d -> Pair.of(d.getNamespace().getValue(), d.getName().getValue()))
142+
.collect(Collectors.toList()));
143+
144+
Map<Dataset, List<ColumnLineage>> datasetLineage = new HashMap<>();
145+
lineageRowsForDatasets.stream()
146+
.forEach(
147+
nodeData -> {
148+
Dataset dataset =
149+
datasets.stream()
150+
.filter(d -> d.getNamespace().getValue().equals(nodeData.getNamespace()))
151+
.filter(d -> d.getName().getValue().equals(nodeData.getName()))
152+
.findAny()
153+
.get();
154+
155+
if (!datasetLineage.containsKey(dataset)) {
156+
datasetLineage.put(dataset, new LinkedList<>());
157+
}
158+
datasetLineage
159+
.get(dataset)
160+
.add(
161+
ColumnLineage.builder()
162+
.name(nodeData.getField())
163+
.transformationDescription(nodeData.getTransformationDescription())
164+
.transformationType(nodeData.getTransformationType())
165+
.inputFields(
166+
nodeData.getInputFields().stream()
167+
.map(
168+
f ->
169+
new ColumnLineageInputField(
170+
f.getNamespace(), f.getName(), f.getField()))
171+
.collect(Collectors.toList()))
172+
.build());
173+
});
174+
175+
datasets.stream()
176+
.filter(dataset -> datasetLineage.containsKey(dataset))
177+
.forEach(dataset -> dataset.setColumnLineage(datasetLineage.get(dataset)));
178+
}
128179
}

‎api/src/main/java/marquez/service/models/ColumnLineage.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
/*
2+
* Copyright 2018-2022 contributors to the Marquez project
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package marquez.service.models;
7+
8+
import java.util.List;
9+
import javax.validation.constraints.NotNull;
10+
import lombok.Builder;
11+
import lombok.EqualsAndHashCode;
12+
import lombok.Getter;
13+
import lombok.ToString;
14+
15+
@EqualsAndHashCode
16+
@ToString
17+
@Builder
18+
@Getter
19+
public class ColumnLineage {
20+
@NotNull private String name;
21+
@NotNull private List<ColumnLineageInputField> inputFields;
22+
@NotNull private String transformationDescription;
23+
@NotNull private String transformationType;
24+
}

‎api/src/main/java/marquez/service/models/ColumnLineageInputField.java‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
/*
2+
* Copyright 2018-2022 contributors to the Marquez project
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package marquez.service.models;
7+
8+
import javax.validation.constraints.NotNull;
9+
import lombok.AllArgsConstructor;
10+
import lombok.EqualsAndHashCode;
11+
import lombok.Getter;
12+
import lombok.ToString;
13+
14+
@EqualsAndHashCode
15+
@ToString
16+
@Getter
17+
@AllArgsConstructor
18+
public class ColumnLineageInputField {
19+
@NotNull private String datasetNamespace;
20+
@NotNull private String datasetName;
21+
@NotNull private String fieldName;
22+
}

‎api/src/main/java/marquez/service/models/Dataset.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ public abstract class Dataset {
5353
@Nullable private final String lastLifecycleState;
5454
@Nullable private final String description;
5555
@Nullable private final UUID currentVersion;
56+
@Getter @Setter @Nullable private List<ColumnLineage> columnLineage;
5657
@Getter ImmutableMap<String, Object> facets;
5758
@Getter private final boolean isDeleted;
5859

@@ -70,6 +71,7 @@ public Dataset(
7071
@Nullable final String lastLifecycleState,
7172
@Nullable final String description,
7273
@Nullable final UUID currentVersion,
74+
@Nullable final ImmutableList<ColumnLineage> columnLineage,
7375
@Nullable final ImmutableMap<String, Object> facets,
7476
boolean isDeleted) {
7577
this.id = id;
@@ -86,6 +88,7 @@ public Dataset(
8688
this.lastLifecycleState = lastLifecycleState;
8789
this.description = description;
8890
this.currentVersion = currentVersion;
91+
this.columnLineage = columnLineage;
8992
this.facets = (facets == null) ? ImmutableMap.of() : facets;
9093
this.isDeleted = isDeleted;
9194
}

‎api/src/main/java/marquez/service/models/DbTable.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ public DbTable(
5353
lastLifecycleState,
5454
description,
5555
currentVersion,
56+
null,
5657
facets,
5758
isDeleted);
5859
}

‎api/src/main/java/marquez/service/models/Stream.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ public Stream(
5959
lastLifecycleState,
6060
description,
6161
currentVersion,
62+
null,
6263
facets,
6364
isDeleted);
6465
this.schemaLocation = schemaLocation;

0 commit comments

Comments
 (0)
⚡