From dfef1ae736543aba3e02e9e0b6bc4e1b649deb1f Mon Sep 17 00:00:00 2001 From: Jakub Dardzinski Date: Fri, 2 Feb 2024 09:39:38 +0100 Subject: [PATCH] Fill data in column lineage input nodes. Signed-off-by: Jakub Dardzinski --- .../main/java/marquez/db/models/InputFieldNodeData.java | 3 ++- .../main/java/marquez/service/ColumnLineageService.java | 7 ++++--- api/src/main/java/marquez/service/models/NodeData.java | 4 +++- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/api/src/main/java/marquez/db/models/InputFieldNodeData.java b/api/src/main/java/marquez/db/models/InputFieldNodeData.java index 8584227150..ac3f9c1bc7 100644 --- a/api/src/main/java/marquez/db/models/InputFieldNodeData.java +++ b/api/src/main/java/marquez/db/models/InputFieldNodeData.java @@ -11,11 +11,12 @@ import lombok.Getter; import lombok.NonNull; import lombok.ToString; +import marquez.service.models.NodeData; @Getter @AllArgsConstructor @ToString -public class InputFieldNodeData { +public class InputFieldNodeData implements NodeData { @NonNull String namespace; @NonNull String dataset; @Nullable UUID datasetVersion; diff --git a/api/src/main/java/marquez/service/ColumnLineageService.java b/api/src/main/java/marquez/service/ColumnLineageService.java index 206a5121c3..8cbe084ce4 100644 --- a/api/src/main/java/marquez/service/ColumnLineageService.java +++ b/api/src/main/java/marquez/service/ColumnLineageService.java @@ -68,10 +68,11 @@ private Lineage toLineage(Set lineageNodeData, boolean in NodeId nodeId = toNodeId(columnLineageNodeData, includeVersion); graphNodes.put(nodeId, Node.datasetField().data(columnLineageNodeData).id(nodeId)); columnLineageNodeData.getInputFields().stream() - .map(i -> toNodeId(i, includeVersion)) .forEach( - inputNodeId -> { - graphNodes.putIfAbsent(inputNodeId, Node.datasetField().id(inputNodeId)); + inputNode -> { + NodeId inputNodeId = toNodeId(inputNode, includeVersion); + graphNodes.putIfAbsent( + inputNodeId, Node.datasetField().id(inputNodeId).data(inputNode)); Optional.ofNullable(outEdges.get(inputNodeId)) .ifPresentOrElse( nodeEdges -> nodeEdges.add(nodeId), diff --git a/api/src/main/java/marquez/service/models/NodeData.java b/api/src/main/java/marquez/service/models/NodeData.java index c3cdcbdf13..b4f87bcf7b 100644 --- a/api/src/main/java/marquez/service/models/NodeData.java +++ b/api/src/main/java/marquez/service/models/NodeData.java @@ -8,11 +8,13 @@ import com.fasterxml.jackson.annotation.JsonSubTypes; import com.fasterxml.jackson.annotation.JsonTypeInfo; import marquez.db.models.ColumnLineageNodeData; +import marquez.db.models.InputFieldNodeData; @JsonTypeInfo(use = JsonTypeInfo.Id.DEDUCTION) @JsonSubTypes({ @JsonSubTypes.Type(DatasetData.class), @JsonSubTypes.Type(JobData.class), - @JsonSubTypes.Type(ColumnLineageNodeData.class) + @JsonSubTypes.Type(ColumnLineageNodeData.class), + @JsonSubTypes.Type(InputFieldNodeData.class) }) public interface NodeData {}