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

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))

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityRelationshipsResultResolver.java

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

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

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

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
             }

metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java

@@ -2176,7 +2176,8 @@ public Set<Urn> exists(
               urn ->
                   // key aspect is always returned, make sure to only consider the status aspect
                   statusResult.getOrDefault(urn, List.of()).stream()
-                      .filter(aspect -> STATUS_ASPECT_NAME.equals(aspect.schema().getName()))
+                      .filter(
+                          aspect -> STATUS_ASPECT_NAME.equalsIgnoreCase(aspect.schema().getName()))
                       .noneMatch(aspect -> ((Status) aspect).isRemoved()))
           .collect(Collectors.toSet());
     }

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<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 {