Fill data in column lineage input nodes.

Signed-off-by: Jakub Dardzinski <kuba0221@gmail.com>
MarquezProject · Feb 2, 2024 · dfef1ae · dfef1ae
1 parent 0ac3a83
commit dfef1ae
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git 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
@@ -68,10 +68,11 @@ private Lineage toLineage(Set<ColumnLineageNodeData> 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
@@ -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 {}