From 424d357cfdab8e517e5f9624cda9e1399518a3de Mon Sep 17 00:00:00 2001 From: Dexter Lee Date: Mon, 7 Mar 2022 09:36:49 -0800 Subject: [PATCH] Fix lineage (#4326) --- .../metadata/graph/LineageRegistry.java | 8 +++---- .../metadata/graph/GraphServiceTestBase.java | 22 +++++++++++++++++++ 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/metadata-io/src/main/java/com/linkedin/metadata/graph/LineageRegistry.java b/metadata-io/src/main/java/com/linkedin/metadata/graph/LineageRegistry.java index 8eb626bda5..16066e2480 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/graph/LineageRegistry.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/graph/LineageRegistry.java @@ -52,14 +52,14 @@ public class LineageRegistry { // A produces B : A -> downstream (produces, OUTGOING), B -> upstream (produces, INCOMING) for (LineageEdge edge : lineageEdges) { if (edge.isUpstream()) { - upstreamPerEntity.computeIfAbsent(edge.sourceEntity, (k) -> new HashSet<>()) + upstreamPerEntity.computeIfAbsent(edge.sourceEntity.toLowerCase(), (k) -> new HashSet<>()) .add(new EdgeInfo(edge.type, RelationshipDirection.OUTGOING)); - downstreamPerEntity.computeIfAbsent(edge.destEntity, (k) -> new HashSet<>()) + downstreamPerEntity.computeIfAbsent(edge.destEntity.toLowerCase(), (k) -> new HashSet<>()) .add(new EdgeInfo(edge.type, RelationshipDirection.INCOMING)); } else { - downstreamPerEntity.computeIfAbsent(edge.sourceEntity, (k) -> new HashSet<>()) + downstreamPerEntity.computeIfAbsent(edge.sourceEntity.toLowerCase(), (k) -> new HashSet<>()) .add(new EdgeInfo(edge.type, RelationshipDirection.OUTGOING)); - upstreamPerEntity.computeIfAbsent(edge.destEntity, (k) -> new HashSet<>()) + upstreamPerEntity.computeIfAbsent(edge.destEntity.toLowerCase(), (k) -> new HashSet<>()) .add(new EdgeInfo(edge.type, RelationshipDirection.INCOMING)); } } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/graph/GraphServiceTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/graph/GraphServiceTestBase.java index 235e3b8391..2b6c40a2f0 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/graph/GraphServiceTestBase.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/graph/GraphServiceTestBase.java @@ -426,6 +426,28 @@ abstract public class GraphServiceTestBase { downstreamLineage = service.getLineage(datasetThreeUrn, LineageDirection.DOWNSTREAM, 0, 1000, 1); assertEquals(downstreamLineage.getTotal().intValue(), 0); assertEquals(downstreamLineage.getRelationships().size(), 0); + + upstreamLineage = service.getLineage(dataJobOneUrn, LineageDirection.UPSTREAM, 0, 1000, 1); + assertEquals(upstreamLineage.getTotal().intValue(), 2); + assertEquals(upstreamLineage.getRelationships().size(), 2); + relationships = upstreamLineage.getRelationships().stream().collect(Collectors.toMap(LineageRelationship::getEntity, + Function.identity())); + assertTrue(relationships.containsKey(datasetOneUrn)); + assertEquals(relationships.get(datasetOneUrn).getType(), consumes); + assertTrue(relationships.containsKey(datasetTwoUrn)); + assertEquals(relationships.get(datasetTwoUrn).getType(), consumes); + + downstreamLineage = service.getLineage(dataJobOneUrn, LineageDirection.DOWNSTREAM, 0, 1000, 1); + assertEquals(downstreamLineage.getTotal().intValue(), 3); + assertEquals(downstreamLineage.getRelationships().size(), 3); + relationships = downstreamLineage.getRelationships().stream().collect(Collectors.toMap(LineageRelationship::getEntity, + Function.identity())); + assertTrue(relationships.containsKey(datasetThreeUrn)); + assertEquals(relationships.get(datasetThreeUrn).getType(), produces); + assertTrue(relationships.containsKey(datasetFourUrn)); + assertEquals(relationships.get(datasetFourUrn).getType(), produces); + assertTrue(relationships.containsKey(dataJobTwoUrn)); + assertEquals(relationships.get(dataJobTwoUrn).getType(), downstreamOf); } @DataProvider(name = "FindRelatedEntitiesSourceEntityFilterTests")