fix(backend): allow excluding soft-deleted entities in relationship-queries; exclude soft-deleted members of groups (#10920)

- allow excluding soft-deleted entities in relationship-queries
- exclude soft-deleted members of groups
This commit is contained in:
ksrinath 2024-07-26 13:59:56 +05:30 committed by GitHub
parent 304fc4ebc2
commit 1717a300bc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 228 additions and 7 deletions

View File

@ -1869,7 +1869,9 @@ public class GmsGraphQLEngine {
"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))

View File

@ -5,6 +5,7 @@ import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.*;
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.EntityRelationshipsResult;
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 class EntityRelationshipsResultResolver
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 @@ public class EntityRelationshipsResultResolver
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<Urn> existentUrns;
if (context != null && _entityService != null && !includeSoftDelete) {
Set<Urn> allRelatedUrns =
entityRelationships.getRelationships().stream()
.map(EntityRelationship::getEntity)
.collect(Collectors.toSet());
existentUrns = _entityService.exists(context.getOperationContext(), allRelatedUrns, false);
} else {
existentUrns = null;
}
List<EntityRelationship> 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());

View File

@ -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
}
"""

View File

@ -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);
}
}

View File

@ -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
}

View File

@ -2147,6 +2147,43 @@ public abstract class EntityServiceTest<T_AD extends AspectDao, T_RS extends Ret
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<Pair<String, RecordTemplate>> 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<Urn> 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 {