diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityLineageResultResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityLineageResultResolver.java index d872ffad2783d..204e591b1da3e 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityLineageResultResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityLineageResultResolver.java @@ -25,6 +25,7 @@ import graphql.schema.DataFetchingEnvironment; import io.datahubproject.metadata.services.RestrictedService; import java.util.HashSet; +import java.util.Optional; import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.stream.Collectors; @@ -60,14 +61,16 @@ public CompletableFuture get(DataFetchingEnvironment enviro final LineageInput input = bindArgument(environment.getArgument("input"), LineageInput.class); final LineageDirection lineageDirection = input.getDirection(); - @Nullable final Integer start = input.getStart(); // Optional! - @Nullable final Integer count = input.getCount(); // Optional! - @Nullable final Boolean separateSiblings = input.getSeparateSiblings(); // Optional! - @Nullable final Long startTimeMillis = input.getStartTimeMillis(); // Optional! + // All inputs are optional + @Nullable final Integer start = input.getStart(); + @Nullable final Integer count = input.getCount(); + @Nullable final Boolean separateSiblings = input.getSeparateSiblings(); + @Nullable final Long startTimeMillis = input.getStartTimeMillis(); @Nullable final Long endTimeMillis = - ResolverUtils.getLineageEndTimeMillis( - input.getStartTimeMillis(), input.getEndTimeMillis()); // Optional! + ResolverUtils.getLineageEndTimeMillis(input.getStartTimeMillis(), input.getEndTimeMillis()); + final Boolean includeGhostEntities = + Optional.ofNullable(input.getIncludeGhostEntities()).orElse(false); com.linkedin.metadata.graph.LineageDirection resolvedDirection = com.linkedin.metadata.graph.LineageDirection.valueOf(lineageDirection.toString()); @@ -80,6 +83,8 @@ public CompletableFuture get(DataFetchingEnvironment enviro _siblingGraphService.getLineage( context .getOperationContext() + .withSearchFlags( + searchFlags -> searchFlags.setIncludeSoftDeleted(includeGhostEntities)) .withLineageFlags( flags -> flags @@ -91,6 +96,7 @@ public CompletableFuture get(DataFetchingEnvironment enviro count != null ? count : 100, 1, separateSiblings != null ? input.getSeparateSiblings() : false, + input.getIncludeGhostEntities(), new HashSet<>()); Set restrictedUrns = new HashSet<>(); diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql index fd112c9524ac9..07dcdf8cc9584 100644 --- a/datahub-graphql-core/src/main/resources/entity.graphql +++ b/datahub-graphql-core/src/main/resources/entity.graphql @@ -1258,6 +1258,11 @@ input LineageInput { An optional ending time to filter on """ endTimeMillis: Long + + """ + If enabled, include entities that do not exist or are soft deleted. + """ + includeGhostEntities: Boolean = false } """ diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/validation/ValidationUtils.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/validation/ValidationUtils.java index ddcc6b6599231..6ecac70e13c7e 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/validation/ValidationUtils.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/validation/ValidationUtils.java @@ -205,7 +205,8 @@ public static LineageSearchResult validateLineageSearchResult( public static EntityLineageResult validateEntityLineageResult( @Nonnull OperationContext opContext, @Nullable final EntityLineageResult entityLineageResult, - @Nonnull final EntityService entityService) { + @Nonnull final EntityService entityService, + boolean includeGhostEntities) { if (entityLineageResult == null) { return null; } @@ -223,8 +224,8 @@ public static EntityLineageResult validateEntityLineageResult( entityLineageResult.getRelationships(), LineageRelationship::getEntity, entityService, - true, - false) + !includeGhostEntities, + includeGhostEntities) .collect(Collectors.toCollection(LineageRelationshipArray::new)); validatedEntityLineageResult.setFiltered( @@ -280,6 +281,8 @@ private static Stream validateSearchUrns( boolean includeSoftDeleted) { if (enforceSQLExistence) { + // TODO: Always set includeSoftDeleted to true once 0.3.7 OSS merge occurs, as soft deleted + // results will be filtered by graph service Set existingUrns = entityService.exists( opContext, diff --git a/metadata-io/src/main/java/com/linkedin/metadata/graph/SiblingGraphService.java b/metadata-io/src/main/java/com/linkedin/metadata/graph/SiblingGraphService.java index f9287ab34cf19..993b5a457206d 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/graph/SiblingGraphService.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/graph/SiblingGraphService.java @@ -38,10 +38,29 @@ public EntityLineageResult getLineage( int offset, int count, int maxHops) { - return ValidationUtils.validateEntityLineageResult( + return getLineage(opContext, entityUrn, direction, offset, count, maxHops, false, false); + } + + @Nonnull + public EntityLineageResult getLineage( + @Nonnull OperationContext opContext, + @Nonnull Urn entityUrn, + @Nonnull LineageDirection direction, + int offset, + int count, + int maxHops, + boolean separateSiblings, + boolean includeGhostEntities) { + return getLineage( opContext, - getLineage(opContext, entityUrn, direction, offset, count, maxHops, false, new HashSet<>()), - _entityService); + entityUrn, + direction, + offset, + count, + maxHops, + separateSiblings, + includeGhostEntities, + new HashSet<>()); } /** @@ -60,12 +79,14 @@ public EntityLineageResult getLineage( int count, int maxHops, boolean separateSiblings, + boolean includeGhostEntities, @Nonnull Set visitedUrns) { if (separateSiblings) { return ValidationUtils.validateEntityLineageResult( opContext, _graphService.getLineage(opContext, entityUrn, direction, offset, count, maxHops), - _entityService); + _entityService, + includeGhostEntities); } if (maxHops > 1) { @@ -89,7 +110,7 @@ public EntityLineageResult getLineage( // remove your siblings from your lineage entityLineage = filterLineageResultFromSiblings( - opContext, entityUrn, allSiblingsInGroup, entityLineage, null); + opContext, entityUrn, allSiblingsInGroup, entityLineage, null, includeGhostEntities); // Update offset and count to fetch the correct number of edges from the next sibling node offset = Math.max(0, offset - entityLineage.getTotal()); @@ -109,8 +130,17 @@ public EntityLineageResult getLineage( siblingUrn, allSiblingsInGroup, getLineage( - opContext, siblingUrn, direction, offset, count, maxHops, false, visitedUrns), - entityLineage); + opContext, + siblingUrn, + direction, + offset, + count, + maxHops, + false, + includeGhostEntities, + visitedUrns), + entityLineage, + includeGhostEntities); // Update offset and count to fetch the correct number of edges from the next sibling node offset = Math.max(0, offset - nextEntityLineage.getTotal()); @@ -122,7 +152,8 @@ public EntityLineageResult getLineage( ; } - return ValidationUtils.validateEntityLineageResult(opContext, entityLineage, _entityService); + return ValidationUtils.validateEntityLineageResult( + opContext, entityLineage, _entityService, includeGhostEntities); } private int getFiltered(@Nullable EntityLineageResult entityLineageResult) { @@ -138,7 +169,8 @@ private EntityLineageResult filterLineageResultFromSiblings( @Nonnull final Urn urn, @Nonnull final Set allSiblingsInGroup, @Nonnull final EntityLineageResult entityLineageResult, - @Nullable final EntityLineageResult existingResult) { + @Nullable final EntityLineageResult existingResult, + boolean includeGhostEntities) { int numFiltered = 0; // 1) remove the source entities siblings from this entity's downstreams @@ -231,6 +263,6 @@ private EntityLineageResult filterLineageResultFromSiblings( combinedLineageResult.setFiltered( numFiltered + getFiltered(existingResult) + getFiltered(entityLineageResult)); return ValidationUtils.validateEntityLineageResult( - opContext, combinedLineageResult, _entityService); + opContext, combinedLineageResult, _entityService, includeGhostEntities); } } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/graph/sibling/SiblingGraphServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/graph/sibling/SiblingGraphServiceTest.java index 15165f59deb16..a61ea1f2562b0 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/graph/sibling/SiblingGraphServiceTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/graph/sibling/SiblingGraphServiceTest.java @@ -33,6 +33,7 @@ import javax.annotation.Nonnull; import org.mockito.Mockito; import org.testng.annotations.BeforeClass; +import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; public class SiblingGraphServiceTest { @@ -73,9 +74,6 @@ public class SiblingGraphServiceTest { @BeforeClass public void setup() { _mockEntityService = Mockito.mock(EntityService.class); - when(_mockEntityService.exists( - any(OperationContext.class), any(Collection.class), any(Boolean.class))) - .thenAnswer(args -> new HashSet<>(args.getArgument(1))); EntityRegistry entityRegistry = new ConfigEntityRegistry( Snapshot.class.getClassLoader().getResourceAsStream("entity-registry.yml")); @@ -84,34 +82,16 @@ public void setup() { _client = new SiblingGraphService(_mockEntityService, _graphService); } + @BeforeMethod + public void init() { + when(_mockEntityService.exists( + any(OperationContext.class), any(Collection.class), any(Boolean.class))) + .thenAnswer(args -> new HashSet<>(args.getArgument(1))); + } + @Test public void testNoSiblingMetadata() { - EntityLineageResult mockResult = new EntityLineageResult(); - LineageRelationshipArray relationships = new LineageRelationshipArray(); - LineageRelationship relationship1 = new LineageRelationship(); - relationship1.setDegree(0); - relationship1.setType(downstreamOf); - relationship1.setEntity(datasetOneUrn); - - LineageRelationship relationship2 = new LineageRelationship(); - relationship2.setDegree(0); - relationship2.setType(downstreamOf); - relationship2.setEntity(datasetTwoUrn); - - LineageRelationship relationship3 = new LineageRelationship(); - relationship3.setDegree(0); - relationship3.setType(downstreamOf); - relationship3.setEntity(datasetThreeUrn); - - relationships.add(relationship1); - relationships.add(relationship2); - relationships.add(relationship3); - - mockResult.setStart(0); - mockResult.setTotal(200); - mockResult.setCount(3); - mockResult.setFiltered(0); - mockResult.setRelationships(relationships); + EntityLineageResult mockResult = makeBasicMockResult(); when(_graphService.getLineage( any(OperationContext.class), @@ -137,35 +117,9 @@ public void testNoSiblingMetadata() { @Test public void testNoSiblingInResults() { - EntityLineageResult mockResult = new EntityLineageResult(); + EntityLineageResult mockResult = makeBasicMockResult(); EntityLineageResult siblingMockResult = new EntityLineageResult(); - LineageRelationshipArray relationships = new LineageRelationshipArray(); - LineageRelationship relationship1 = new LineageRelationship(); - relationship1.setDegree(0); - relationship1.setType(downstreamOf); - relationship1.setEntity(datasetOneUrn); - - LineageRelationship relationship2 = new LineageRelationship(); - relationship2.setDegree(0); - relationship2.setType(downstreamOf); - relationship2.setEntity(datasetTwoUrn); - - LineageRelationship relationship3 = new LineageRelationship(); - relationship3.setDegree(0); - relationship3.setType(downstreamOf); - relationship3.setEntity(datasetThreeUrn); - - relationships.add(relationship1); - relationships.add(relationship2); - relationships.add(relationship3); - - mockResult.setStart(0); - mockResult.setTotal(200); - mockResult.setCount(3); - mockResult.setFiltered(0); - mockResult.setRelationships(relationships); - when(_graphService.getLineage( any(OperationContext.class), eq(datasetFourUrn), @@ -229,34 +183,9 @@ public void testNoSiblingInResults() { @Test public void testSiblingInResult() throws Exception { - EntityLineageResult mockResult = new EntityLineageResult(); + EntityLineageResult mockResult = makeBasicMockResult(); EntityLineageResult siblingMockResult = new EntityLineageResult(); - LineageRelationshipArray relationships = new LineageRelationshipArray(); - LineageRelationship relationship1 = new LineageRelationship(); - relationship1.setDegree(0); - relationship1.setType(downstreamOf); - relationship1.setEntity(datasetOneUrn); - - LineageRelationship relationship2 = new LineageRelationship(); - relationship2.setDegree(0); - relationship2.setType(downstreamOf); - relationship2.setEntity(datasetTwoUrn); - - LineageRelationship relationship3 = new LineageRelationship(); - relationship3.setDegree(0); - relationship3.setType(downstreamOf); - relationship3.setEntity(datasetThreeUrn); - - relationships.add(relationship1); - relationships.add(relationship2); - relationships.add(relationship3); - - mockResult.setStart(0); - mockResult.setTotal(3); - mockResult.setCount(3); - mockResult.setRelationships(relationships); - siblingMockResult.setStart(0); siblingMockResult.setTotal(0); siblingMockResult.setCount(0); @@ -315,7 +244,9 @@ public void testSiblingInResult() throws Exception { expectedResult.setTotal(3); expectedResult.setCount(2); expectedResult.setFiltered(1); - expectedResult.setRelationships(new LineageRelationshipArray(relationship1, relationship2)); + expectedResult.setRelationships( + new LineageRelationshipArray( + makeBasicRelationship(datasetOneUrn), makeBasicRelationship(datasetTwoUrn))); EntityLineageResult upstreamLineage = service.getLineage(opContext, datasetFourUrn, LineageDirection.UPSTREAM, 0, 100, 1); @@ -335,25 +266,9 @@ public void testCombineSiblingResult() { LineageRelationshipArray siblingRelationships = new LineageRelationshipArray(); LineageRelationshipArray expectedRelationships = new LineageRelationshipArray(); - LineageRelationship relationship1 = new LineageRelationship(); - relationship1.setDegree(0); - relationship1.setType(downstreamOf); - relationship1.setEntity(datasetOneUrn); - - LineageRelationship relationship2 = new LineageRelationship(); - relationship2.setDegree(0); - relationship2.setType(downstreamOf); - relationship2.setEntity(datasetTwoUrn); - - LineageRelationship relationship3 = new LineageRelationship(); - relationship3.setDegree(0); - relationship3.setType(downstreamOf); - relationship3.setEntity(datasetThreeUrn); - - LineageRelationship relationship4 = new LineageRelationship(); - relationship4.setDegree(0); - relationship4.setType(downstreamOf); - relationship4.setEntity(datasetFiveUrn); + LineageRelationship relationship1 = makeBasicRelationship(datasetOneUrn); + LineageRelationship relationship2 = makeBasicRelationship(datasetTwoUrn); + LineageRelationship relationship4 = makeBasicRelationship(datasetFiveUrn); relationships.add(relationship1); @@ -449,25 +364,9 @@ public void testUpstreamOfSiblings() { LineageRelationshipArray siblingRelationships = new LineageRelationshipArray(); LineageRelationshipArray expectedRelationships = new LineageRelationshipArray(); - LineageRelationship relationship1 = new LineageRelationship(); - relationship1.setDegree(0); - relationship1.setType(downstreamOf); - relationship1.setEntity(datasetOneUrn); - - LineageRelationship relationship2 = new LineageRelationship(); - relationship2.setDegree(0); - relationship2.setType(downstreamOf); - relationship2.setEntity(datasetTwoUrn); - - LineageRelationship relationship3 = new LineageRelationship(); - relationship3.setDegree(0); - relationship3.setType(downstreamOf); - relationship3.setEntity(datasetThreeUrn); - - LineageRelationship relationship5 = new LineageRelationship(); - relationship5.setDegree(0); - relationship5.setType(downstreamOf); - relationship5.setEntity(datasetFiveUrn); + LineageRelationship relationship1 = makeBasicRelationship(datasetOneUrn); + LineageRelationship relationship2 = makeBasicRelationship(datasetTwoUrn); + LineageRelationship relationship5 = makeBasicRelationship(datasetFiveUrn); relationships.add(relationship1); @@ -607,11 +506,7 @@ public void testUpstreamOfSiblingSiblings() { LineageRelationshipArray relationships = new LineageRelationshipArray(); LineageRelationshipArray expectedRelationships = new LineageRelationshipArray(); - LineageRelationship relationship = new LineageRelationship(); - relationship.setDegree(0); - relationship.setType(downstreamOf); - relationship.setEntity(datasetFourUrn); - + LineageRelationship relationship = makeBasicRelationship(datasetFourUrn); relationships.add(relationship); expectedRelationships.add(relationship); @@ -722,25 +617,10 @@ public void testRelationshipWithSibling() throws CloneNotSupportedException { LineageRelationshipArray siblingRelationships = new LineageRelationshipArray(); LineageRelationshipArray expectedRelationships = new LineageRelationshipArray(); - LineageRelationship relationship1 = new LineageRelationship(); - relationship1.setDegree(0); - relationship1.setType(downstreamOf); - relationship1.setEntity(datasetOneUrn); - - LineageRelationship relationship2 = new LineageRelationship(); - relationship2.setDegree(0); - relationship2.setType(downstreamOf); - relationship2.setEntity(datasetTwoUrn); - - LineageRelationship relationship3 = new LineageRelationship(); - relationship3.setDegree(0); - relationship3.setType(downstreamOf); - relationship3.setEntity(datasetThreeUrn); - - LineageRelationship relationship5 = new LineageRelationship(); - relationship5.setDegree(0); - relationship5.setType(downstreamOf); - relationship5.setEntity(datasetFiveUrn); + LineageRelationship relationship1 = makeBasicRelationship(datasetOneUrn); + LineageRelationship relationship2 = makeBasicRelationship(datasetTwoUrn); + LineageRelationship relationship3 = makeBasicRelationship(datasetThreeUrn); + LineageRelationship relationship5 = makeBasicRelationship(datasetFiveUrn); relationships.add(relationship1); // relationship between entity and its sibling @@ -1006,7 +886,7 @@ public void testSiblingCombinations() throws URISyntaxException { // Tests for separateSiblings = true: primary sibling EntityLineageResult primaryDownstreamSeparated = service.getLineage( - opContext, primarySiblingUrn, LineageDirection.DOWNSTREAM, 0, 100, 1, true, Set.of()); + opContext, primarySiblingUrn, LineageDirection.DOWNSTREAM, 0, 100, 1, true, false); LineageRelationshipArray expectedRelationships = new LineageRelationshipArray(); expectedRelationships.add(relationship); @@ -1022,7 +902,7 @@ public void testSiblingCombinations() throws URISyntaxException { EntityLineageResult primaryUpstreamSeparated = service.getLineage( - opContext, primarySiblingUrn, LineageDirection.UPSTREAM, 0, 100, 1, true, Set.of()); + opContext, primarySiblingUrn, LineageDirection.UPSTREAM, 0, 100, 1, true, false); EntityLineageResult expectedResultPrimaryUpstreamSeparated = new EntityLineageResult(); expectedResultPrimaryUpstreamSeparated.setCount(2); expectedResultPrimaryUpstreamSeparated.setStart(0); @@ -1035,7 +915,7 @@ public void testSiblingCombinations() throws URISyntaxException { // Test for separateSiblings = true, secondary sibling EntityLineageResult secondarySiblingSeparated = service.getLineage( - opContext, alternateSiblingUrn, LineageDirection.DOWNSTREAM, 0, 100, 1, true, Set.of()); + opContext, alternateSiblingUrn, LineageDirection.DOWNSTREAM, 0, 100, 1, true, false); EntityLineageResult expectedResultSecondarySeparated = new EntityLineageResult(); expectedResultSecondarySeparated.setCount(numDownstreams); @@ -1048,7 +928,7 @@ public void testSiblingCombinations() throws URISyntaxException { EntityLineageResult secondaryUpstreamSeparated = service.getLineage( - opContext, alternateSiblingUrn, LineageDirection.UPSTREAM, 0, 100, 1, true, Set.of()); + opContext, alternateSiblingUrn, LineageDirection.UPSTREAM, 0, 100, 1, true, false); EntityLineageResult expectedResultSecondaryUpstreamSeparated = new EntityLineageResult(); expectedResultSecondaryUpstreamSeparated.setCount(3); expectedResultSecondaryUpstreamSeparated.setStart(0); @@ -1060,15 +940,7 @@ public void testSiblingCombinations() throws URISyntaxException { // Test for separateSiblings = false, primary sibling EntityLineageResult primarySiblingNonSeparated = - service.getLineage( - opContext, - primarySiblingUrn, - LineageDirection.DOWNSTREAM, - 0, - 100, - 1, - false, - new HashSet<>()); + service.getLineage(opContext, primarySiblingUrn, LineageDirection.DOWNSTREAM, 0, 100, 1); EntityLineageResult expectedResultPrimaryNonSeparated = new EntityLineageResult(); expectedResultPrimaryNonSeparated.setCount(numDownstreams); expectedResultPrimaryNonSeparated.setStart(0); @@ -1078,15 +950,7 @@ public void testSiblingCombinations() throws URISyntaxException { assertEquals(primarySiblingNonSeparated, expectedResultPrimaryNonSeparated); EntityLineageResult primarySiblingNonSeparatedUpstream = - service.getLineage( - opContext, - primarySiblingUrn, - LineageDirection.UPSTREAM, - 0, - 100, - 1, - false, - new HashSet<>()); + service.getLineage(opContext, primarySiblingUrn, LineageDirection.UPSTREAM, 0, 100, 1); EntityLineageResult expectedResultPrimaryUpstreamNonSeparated = new EntityLineageResult(); expectedResultPrimaryUpstreamNonSeparated.setCount(2); expectedResultPrimaryUpstreamNonSeparated.setStart(0); @@ -1097,30 +961,84 @@ public void testSiblingCombinations() throws URISyntaxException { // Test for separateSiblings = false, secondary sibling EntityLineageResult secondarySiblingNonSeparated = - service.getLineage( - opContext, - alternateSiblingUrn, - LineageDirection.DOWNSTREAM, - 0, - 100, - 1, - false, - new HashSet<>()); + service.getLineage(opContext, alternateSiblingUrn, LineageDirection.DOWNSTREAM, 0, 100, 1); assertEquals(secondarySiblingNonSeparated, expectedResultPrimaryNonSeparated); EntityLineageResult secondarySiblingNonSeparatedUpstream = - service.getLineage( - opContext, - alternateSiblingUrn, - LineageDirection.UPSTREAM, - 0, - 100, - 1, - false, - new HashSet<>()); + service.getLineage(opContext, alternateSiblingUrn, LineageDirection.UPSTREAM, 0, 100, 1); assertEquals(secondarySiblingNonSeparatedUpstream, expectedResultPrimaryUpstreamNonSeparated); } + @Test + public void testExcludeGhostEntities() { + when(_mockEntityService.exists(any(OperationContext.class), any(Collection.class), eq(false))) + .thenAnswer(args -> Set.of(datasetOneUrn)); + + EntityLineageResult mockGraphResult = makeBasicMockResult(); + + when(_graphService.getLineage( + any(OperationContext.class), + eq(datasetFourUrn), + eq(LineageDirection.UPSTREAM), + eq(0), + eq(100), + eq(1))) + .thenReturn(mockGraphResult); + + when(_mockEntityService.getLatestAspect( + any(OperationContext.class), eq(datasetFourUrn), eq(SIBLINGS_ASPECT_NAME))) + .thenReturn(null); + + SiblingGraphService service = _client; + + EntityLineageResult upstreamLineage = + service.getLineage(opContext, datasetFourUrn, LineageDirection.UPSTREAM, 0, 100, 1); + + EntityLineageResult mockResult = new EntityLineageResult(); + mockResult.setStart(0); + mockResult.setTotal(3); + mockResult.setCount(3); + mockResult.setFiltered(2); + LineageRelationshipArray relationshipsResult = new LineageRelationshipArray(); + relationshipsResult.add(makeBasicRelationship(datasetOneUrn)); + mockResult.setRelationships(relationshipsResult); + + // assert sibling graph service filters out entities that do not exist + assertEquals(upstreamLineage, mockResult); + } + + @Test + public void testIncludeGhostEntities() { + when(_mockEntityService.exists( + any(OperationContext.class), any(Collection.class), any(Boolean.class))) + .thenAnswer(args -> Set.of(datasetOneUrn)); + + EntityLineageResult mockResult = makeBasicMockResult(); + + when(_graphService.getLineage( + any(OperationContext.class), + eq(datasetFourUrn), + eq(LineageDirection.UPSTREAM), + eq(0), + eq(100), + eq(1))) + .thenReturn(mockResult); + + when(_mockEntityService.getLatestAspect( + any(OperationContext.class), eq(datasetFourUrn), eq(SIBLINGS_ASPECT_NAME))) + .thenReturn(null); + + SiblingGraphService service = _client; + + EntityLineageResult upstreamLineage = + service.getLineage( + opContext, datasetFourUrn, LineageDirection.UPSTREAM, 0, 100, 1, false, true); + + // assert sibling graph service is a pass through when there are no siblings and + // includeGhostEntities + assertEquals(upstreamLineage, mockResult); + } + static Urn createFromString(@Nonnull String rawUrn) { try { return Urn.createFromString(rawUrn); @@ -1128,4 +1046,29 @@ static Urn createFromString(@Nonnull String rawUrn) { return null; } } + + static LineageRelationship makeBasicRelationship(Urn urn) { + LineageRelationship relationship = new LineageRelationship(); + relationship.setDegree(0); + relationship.setType(downstreamOf); + relationship.setEntity(urn); + return relationship; + } + + static EntityLineageResult makeBasicMockResult() { + LineageRelationshipArray relationships = new LineageRelationshipArray(); + LineageRelationship relationship1 = makeBasicRelationship(datasetOneUrn); + LineageRelationship relationship2 = makeBasicRelationship(datasetTwoUrn); + LineageRelationship relationship3 = makeBasicRelationship(datasetThreeUrn); + relationships.addAll(List.of(relationship1, relationship2, relationship3)); + + EntityLineageResult mockResult = new EntityLineageResult(); + mockResult.setStart(0); + mockResult.setTotal(3); + mockResult.setCount(3); + mockResult.setFiltered(0); + mockResult.setRelationships(relationships); + + return mockResult; + } }