diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java index b751849859be4..cba214365fdac 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java @@ -1869,7 +1869,9 @@ private void configureCorpGroupResolvers(final RuntimeWiring.Builder builder) { "CorpGroup", typeWiring -> typeWiring - .dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient)) + .dataFetcher( + "relationships", + new EntityRelationshipsResultResolver(graphClient, entityService)) .dataFetcher("privileges", new EntityPrivilegesResolver(entityClient)) .dataFetcher( "aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry)) diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityRelationshipsResultResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityRelationshipsResultResolver.java index f775853dd5956..fd72edb2972e3 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityRelationshipsResultResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityRelationshipsResultResolver.java @@ -5,6 +5,7 @@ import com.linkedin.common.EntityRelationship; import com.linkedin.common.EntityRelationships; +import com.linkedin.common.urn.Urn; import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils; import com.linkedin.datahub.graphql.generated.Entity; @@ -12,11 +13,13 @@ import com.linkedin.datahub.graphql.generated.RelationshipsInput; import com.linkedin.datahub.graphql.types.common.mappers.AuditStampMapper; import com.linkedin.datahub.graphql.types.common.mappers.UrnToEntityMapper; +import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.graph.GraphClient; import com.linkedin.metadata.query.filter.RelationshipDirection; import graphql.schema.DataFetcher; import graphql.schema.DataFetchingEnvironment; import java.util.List; +import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.stream.Collectors; import javax.annotation.Nullable; @@ -29,8 +32,16 @@ public class EntityRelationshipsResultResolver private final GraphClient _graphClient; + private final EntityService _entityService; + public EntityRelationshipsResultResolver(final GraphClient graphClient) { + this(graphClient, null); + } + + public EntityRelationshipsResultResolver( + final GraphClient graphClient, final EntityService entityService) { _graphClient = graphClient; + _entityService = entityService; } @Override @@ -47,13 +58,16 @@ public CompletableFuture get(DataFetchingEnvironment final Integer count = input.getCount(); // Optional! final RelationshipDirection resolvedDirection = RelationshipDirection.valueOf(relationshipDirection.toString()); + final boolean includeSoftDelete = input.getIncludeSoftDelete(); + return GraphQLConcurrencyUtils.supplyAsync( () -> mapEntityRelationships( context, fetchEntityRelationships( urn, relationshipTypes, resolvedDirection, start, count, context.getActorUrn()), - resolvedDirection), + resolvedDirection, + includeSoftDelete), this.getClass().getSimpleName(), "get"); } @@ -72,13 +86,28 @@ private EntityRelationships fetchEntityRelationships( private EntityRelationshipsResult mapEntityRelationships( @Nullable final QueryContext context, final EntityRelationships entityRelationships, - final RelationshipDirection relationshipDirection) { + final RelationshipDirection relationshipDirection, + final boolean includeSoftDelete) { final EntityRelationshipsResult result = new EntityRelationshipsResult(); + final Set existentUrns; + if (context != null && _entityService != null && !includeSoftDelete) { + Set allRelatedUrns = + entityRelationships.getRelationships().stream() + .map(EntityRelationship::getEntity) + .collect(Collectors.toSet()); + existentUrns = _entityService.exists(context.getOperationContext(), allRelatedUrns, false); + } else { + existentUrns = null; + } + List viewable = entityRelationships.getRelationships().stream() .filter( - rel -> context == null || canView(context.getOperationContext(), rel.getEntity())) + rel -> + (existentUrns == null || existentUrns.contains(rel.getEntity())) + && (context == null + || canView(context.getOperationContext(), rel.getEntity()))) .collect(Collectors.toList()); result.setStart(entityRelationships.getStart()); diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql index 246ace2fc0f5f..c18550d2d407f 100644 --- a/datahub-graphql-core/src/main/resources/entity.graphql +++ b/datahub-graphql-core/src/main/resources/entity.graphql @@ -1267,6 +1267,11 @@ input RelationshipsInput { The number of results to be returned """ count: Int + + """ + Whether to include soft-deleted, related, entities + """ + includeSoftDelete: Boolean = true } """ diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/load/EntityRelationshipsResultResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/load/EntityRelationshipsResultResolverTest.java new file mode 100644 index 0000000000000..d2799278c1238 --- /dev/null +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/load/EntityRelationshipsResultResolverTest.java @@ -0,0 +1,124 @@ +package com.linkedin.datahub.graphql.resolvers.load; + +import static com.linkedin.datahub.graphql.TestUtils.getMockAllowContext; +import static org.mockito.ArgumentMatchers.*; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; + +import com.linkedin.common.EntityRelationship; +import com.linkedin.common.EntityRelationshipArray; +import com.linkedin.common.EntityRelationships; +import com.linkedin.common.urn.Urn; +import com.linkedin.datahub.graphql.QueryContext; +import com.linkedin.datahub.graphql.generated.*; +import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.graph.GraphClient; +import graphql.schema.DataFetchingEnvironment; +import java.net.URISyntaxException; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ExecutionException; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +public class EntityRelationshipsResultResolverTest { + private final Urn existentUser = Urn.createFromString("urn:li:corpuser:johndoe"); + private final Urn softDeletedUser = Urn.createFromString("urn:li:corpuser:deletedUser"); + + private CorpUser existentEntity; + private CorpUser softDeletedEntity; + + private EntityService _entityService; + private GraphClient _graphClient; + + private EntityRelationshipsResultResolver resolver; + private RelationshipsInput input; + private DataFetchingEnvironment mockEnv; + + public EntityRelationshipsResultResolverTest() throws URISyntaxException {} + + @BeforeMethod + public void setupTest() { + _entityService = mock(EntityService.class); + _graphClient = mock(GraphClient.class); + resolver = new EntityRelationshipsResultResolver(_graphClient, _entityService); + + mockEnv = mock(DataFetchingEnvironment.class); + QueryContext context = getMockAllowContext(); + when(mockEnv.getContext()).thenReturn(context); + + CorpGroup source = new CorpGroup(); + source.setUrn("urn:li:corpGroup:group1"); + when(mockEnv.getSource()).thenReturn(source); + + when(_entityService.exists(any(), eq(Set.of(existentUser, softDeletedUser)), eq(true))) + .thenReturn(Set.of(existentUser, softDeletedUser)); + when(_entityService.exists(any(), eq(Set.of(existentUser, softDeletedUser)), eq(false))) + .thenReturn(Set.of(existentUser)); + + input = new RelationshipsInput(); + input.setStart(0); + input.setCount(10); + input.setDirection(RelationshipDirection.INCOMING); + input.setTypes(List.of("SomeType")); + + EntityRelationships entityRelationships = + new EntityRelationships() + .setStart(0) + .setCount(2) + .setTotal(2) + .setRelationships( + new EntityRelationshipArray( + new EntityRelationship().setEntity(existentUser).setType("SomeType"), + new EntityRelationship().setEntity(softDeletedUser).setType("SomeType"))); + + // always expected INCOMING, and "SomeType" in all tests + when(_graphClient.getRelatedEntities( + eq(source.getUrn()), + eq(input.getTypes()), + same(com.linkedin.metadata.query.filter.RelationshipDirection.INCOMING), + eq(input.getStart()), + eq(input.getCount()), + any())) + .thenReturn(entityRelationships); + + when(mockEnv.getArgument(eq("input"))).thenReturn(input); + + existentEntity = new CorpUser(); + existentEntity.setUrn(existentUser.toString()); + existentEntity.setType(EntityType.CORP_USER); + + softDeletedEntity = new CorpUser(); + softDeletedEntity.setUrn(softDeletedUser.toString()); + softDeletedEntity.setType(EntityType.CORP_USER); + } + + @Test + public void testIncludeSoftDeleted() throws ExecutionException, InterruptedException { + EntityRelationshipsResult expected = new EntityRelationshipsResult(); + expected.setRelationships( + List.of(resultRelationship(existentEntity), resultRelationship(softDeletedEntity))); + expected.setStart(0); + expected.setCount(2); + expected.setTotal(2); + assertEquals(resolver.get(mockEnv).get().toString(), expected.toString()); + } + + @Test + public void testExcludeSoftDeleted() throws ExecutionException, InterruptedException { + input.setIncludeSoftDelete(false); + EntityRelationshipsResult expected = new EntityRelationshipsResult(); + expected.setRelationships(List.of(resultRelationship(existentEntity))); + expected.setStart(0); + expected.setCount(1); + expected.setTotal(1); + assertEquals(resolver.get(mockEnv).get().toString(), expected.toString()); + } + + private com.linkedin.datahub.graphql.generated.EntityRelationship resultRelationship( + Entity entity) { + return new com.linkedin.datahub.graphql.generated.EntityRelationship( + "SomeType", RelationshipDirection.INCOMING, entity, null); + } +} diff --git a/datahub-web-react/src/graphql/group.graphql b/datahub-web-react/src/graphql/group.graphql index ee04489540f9c..60da627fd254d 100644 --- a/datahub-web-react/src/graphql/group.graphql +++ b/datahub-web-react/src/graphql/group.graphql @@ -48,6 +48,7 @@ query getGroup($urn: String!, $membersCount: Int!) { direction: INCOMING start: 0 count: $membersCount + includeSoftDelete: false } ) { start @@ -86,6 +87,7 @@ query getAllGroupMembers($urn: String!, $start: Int!, $count: Int!) { direction: INCOMING start: $start count: $count + includeSoftDelete: false } ) { start @@ -121,7 +123,15 @@ query getAllGroupMembers($urn: String!, $start: Int!, $count: Int!) { query getGroupMembers($urn: String!, $start: Int!, $count: Int!) { corpGroup(urn: $urn) { - relationships(input: { types: ["IsMemberOfGroup"], direction: INCOMING, start: $start, count: $count }) { + relationships( + input: { + types: ["IsMemberOfGroup"] + direction: INCOMING + start: $start + count: $count + includeSoftDelete: false + } + ) { start count total @@ -155,7 +165,15 @@ query getGroupMembers($urn: String!, $start: Int!, $count: Int!) { query getNativeGroupMembers($urn: String!, $start: Int!, $count: Int!) { corpGroup(urn: $urn) { - relationships(input: { types: ["IsMemberOfNativeGroup"], direction: INCOMING, start: $start, count: $count }) { + relationships( + input: { + types: ["IsMemberOfNativeGroup"] + direction: INCOMING + start: $start + count: $count + includeSoftDelete: false + } + ) { start count total @@ -209,7 +227,13 @@ query listGroups($input: ListGroupsInput!) { pictureLink } memberCount: relationships( - input: { types: ["IsMemberOfGroup", "IsMemberOfNativeGroup"], direction: INCOMING, start: 0, count: 1 } + input: { + types: ["IsMemberOfGroup", "IsMemberOfNativeGroup"] + direction: INCOMING + start: 0 + count: 1 + includeSoftDelete: false + } ) { total } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java index 33598be8fc72b..acdbd7855f7b0 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceTest.java @@ -2147,6 +2147,43 @@ public void testCreateChangeTypeProposal() { opContext, secondCreateProposal, TEST_AUDIT_STAMP, false)); } + @Test + public void testExists() throws Exception { + Urn existentUrn = UrnUtils.getUrn("urn:li:corpuser:exists"); + Urn softDeletedUrn = UrnUtils.getUrn("urn:li:corpuser:softDeleted"); + Urn nonExistentUrn = UrnUtils.getUrn("urn:li:corpuser:nonExistent"); + Urn noStatusUrn = UrnUtils.getUrn("urn:li:corpuser:noStatus"); + + List> pairToIngest = new ArrayList<>(); + SystemMetadata metadata = AspectGenerationUtils.createSystemMetadata(); + + // to ensure existence + CorpUserInfo userInfoAspect = AspectGenerationUtils.createCorpUserInfo("email@test.com"); + pairToIngest.add(getAspectRecordPair(userInfoAspect, CorpUserInfo.class)); + + _entityServiceImpl.ingestAspects( + opContext, noStatusUrn, pairToIngest, TEST_AUDIT_STAMP, metadata); + + Status statusExistsAspect = new Status().setRemoved(false); + pairToIngest.add(getAspectRecordPair(statusExistsAspect, Status.class)); + + _entityServiceImpl.ingestAspects( + opContext, existentUrn, pairToIngest, TEST_AUDIT_STAMP, metadata); + + Status statusRemovedAspect = new Status().setRemoved(true); + pairToIngest.set(1, getAspectRecordPair(statusRemovedAspect, Status.class)); + + _entityServiceImpl.ingestAspects( + opContext, softDeletedUrn, pairToIngest, TEST_AUDIT_STAMP, metadata); + + Set inputUrns = Set.of(existentUrn, softDeletedUrn, nonExistentUrn, noStatusUrn); + assertEquals( + _entityServiceImpl.exists(opContext, inputUrns, false), Set.of(existentUrn, noStatusUrn)); + assertEquals( + _entityServiceImpl.exists(opContext, inputUrns, true), + Set.of(existentUrn, noStatusUrn, softDeletedUrn)); + } + @Nonnull protected com.linkedin.entity.Entity createCorpUserEntity(Urn entityUrn, String email) throws Exception {