Skip to content

Commit

Permalink
fix(backend): allow excluding soft-deleted entities in relationship-q…
Browse files Browse the repository at this point in the history
…ueries; exclude soft-deleted members of groups (#10920)

- allow excluding soft-deleted entities in relationship-queries
- exclude soft-deleted members of groups
  • Loading branch information
ksrinath authored Jul 26, 2024
1 parent 304fc4e commit 1717a30
Show file tree
Hide file tree
Showing 6 changed files with 228 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,18 +5,21 @@

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;
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;
Expand All @@ -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
Expand All @@ -47,13 +58,16 @@ public CompletableFuture<EntityRelationshipsResult> 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");
}
Expand All @@ -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<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());
Expand Down
5 changes: 5 additions & 0 deletions datahub-graphql-core/src/main/resources/entity.graphql
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

"""
Expand Down
Original file line number Diff line number Diff line change
@@ -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);
}
}
30 changes: 27 additions & 3 deletions datahub-web-react/src/graphql/group.graphql
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ query getGroup($urn: String!, $membersCount: Int!) {
direction: INCOMING
start: 0
count: $membersCount
includeSoftDelete: false
}
) {
start
Expand Down Expand Up @@ -86,6 +87,7 @@ query getAllGroupMembers($urn: String!, $start: Int!, $count: Int!) {
direction: INCOMING
start: $start
count: $count
includeSoftDelete: false
}
) {
start
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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<Pair<String, RecordTemplate>> pairToIngest = new ArrayList<>();
SystemMetadata metadata = AspectGenerationUtils.createSystemMetadata();

// to ensure existence
CorpUserInfo userInfoAspect = AspectGenerationUtils.createCorpUserInfo("[email protected]");
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 {
Expand Down

0 comments on commit 1717a30

Please sign in to comment.