fix(lineage): Fix Upstream + Downstream Count in presence of Soft-Deleted / Non-Existent references (#7374)

2025-12-25 00:48:45 +00:00 · 2023-02-20 14:00:14 -08:00 · 2023-02-20 14:00:14 -08:00 · 92cd2b2c1b
commit 92cd2b2c1b
parent bd11575d5f
10 changed files with 59 additions and 15 deletions
--- 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
@ -79,6 +79,7 @@ public class EntityLineageResultResolver implements DataFetcher<CompletableFutur
    result.setStart(entityLineageResult.getStart());
    result.setCount(entityLineageResult.getCount());
    result.setTotal(entityLineageResult.getTotal());
+    result.setFiltered(entityLineageResult.getFiltered());
    result.setRelationships(entityLineageResult.getRelationships()
        .stream()
        .map(this::mapEntityRelationship)
--- a/datahub-graphql-core/src/main/resources/entity.graphql
+++ b/datahub-graphql-core/src/main/resources/entity.graphql
@ -943,6 +943,11 @@ type EntityLineageResult {
    """
    total: Int

+    """
+    The number of results that were filtered out of the page (soft-deleted or non-existent)
+    """
+    filtered: Int
+
    """
    Relationships in the result set
    """
--- a/datahub-web-react/src/app/entity/EntityRegistry.tsx
+++ b/datahub-web-react/src/app/entity/EntityRegistry.tsx
@ -128,30 +128,28 @@ export default class EntityRegistry {
                ...entity.getLineageVizConfig?.(data),
                downstreamChildren: genericEntityProperties?.downstream?.relationships
                    ?.filter((relationship) => relationship.entity)
-                    // eslint-disable-next-line @typescript-eslint/dot-notation
-                    ?.filter((relationship) => !relationship.entity?.['status']?.removed)
                    ?.map((relationship) => ({
                        entity: relationship.entity as EntityInterface,
                        type: (relationship.entity as EntityInterface).type,
                    })),
-                downstreamRelationships: genericEntityProperties?.downstream?.relationships
-                    ?.filter((relationship) => relationship.entity)
-                    // eslint-disable-next-line @typescript-eslint/dot-notation
-                    ?.filter((relationship) => !relationship.entity?.['status']?.removed),
-                numDownstreamChildren: genericEntityProperties?.downstream?.total,
+                downstreamRelationships: genericEntityProperties?.downstream?.relationships?.filter(
+                    (relationship) => relationship.entity,
+                ),
+                numDownstreamChildren:
+                    (genericEntityProperties?.downstream?.total || 0) -
+                    (genericEntityProperties?.downstream?.filtered || 0),
                upstreamChildren: genericEntityProperties?.upstream?.relationships
                    ?.filter((relationship) => relationship.entity)
-                    // eslint-disable-next-line @typescript-eslint/dot-notation
-                    ?.filter((relationship) => !relationship.entity?.['status']?.removed)
                    ?.map((relationship) => ({
                        entity: relationship.entity as EntityInterface,
                        type: (relationship.entity as EntityInterface).type,
                    })),
-                upstreamRelationships: genericEntityProperties?.upstream?.relationships
-                    ?.filter((relationship) => relationship.entity)
-                    // eslint-disable-next-line @typescript-eslint/dot-notation
-                    ?.filter((relationship) => !relationship.entity?.['status']?.removed),
-                numUpstreamChildren: genericEntityProperties?.upstream?.total,
+                upstreamRelationships: genericEntityProperties?.upstream?.relationships?.filter(
+                    (relationship) => relationship.entity,
+                ),
+                numUpstreamChildren:
+                    (genericEntityProperties?.upstream?.total || 0) -
+                    (genericEntityProperties?.upstream?.filtered || 0),
                status: genericEntityProperties?.status,
                siblingPlatforms: genericEntityProperties?.siblingPlatforms,
                fineGrainedLineages: genericEntityProperties?.fineGrainedLineages,
--- a/datahub-web-react/src/app/lineage/utils/navigateToLineageUrl.ts
+++ b/datahub-web-react/src/app/lineage/utils/navigateToLineageUrl.ts
@ -44,7 +44,6 @@ export const navigateToLineageUrl = ({
    }
    const newSearchStringified = QueryString.stringify(newSearch, { arrayFormat: 'comma' });

-    console.log(location.pathname);
    history.push({
        pathname: location.pathname,
        search: newSearchStringified,
--- a/datahub-web-react/src/graphql/lineage.graphql
+++ b/datahub-web-react/src/graphql/lineage.graphql
@ -251,6 +251,7 @@ fragment fullLineageResults on EntityLineageResult {
    start
    count
    total
+    filtered
    relationships {
        type
        createdOn
@ -308,6 +309,7 @@ fragment leafLineageResults on EntityLineageResult {
    start
    count
    total
+    filtered
    relationships {
        type
        entity {
@ -321,6 +323,7 @@ fragment partialLineageResults on EntityLineageResult {
    start
    count
    total
+    filtered
 }

 query getEntityLineage(
--- a/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/EntityService.java
@ -76,6 +76,7 @@ import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
@ -1608,6 +1609,12 @@ public class EntityService {
    return new RollbackRunResult(removedAspects, rowsDeletedFromEntityDeletion);
  }

+  /**
+   * Returns true if the entity exists (has materialized aspects)
+   *
+   * @param urn the urn of the entity to check
+   * @return true if the entity exists, false otherwise
+   */
  public Boolean exists(Urn urn) {
    final Set<String> aspectsToFetch = getEntityAspectNames(urn);
    final List<EntityAspectIdentifier> dbKeys = aspectsToFetch.stream()
@ -1618,6 +1625,18 @@ public class EntityService {
    return aspects.values().stream().anyMatch(aspect -> aspect != null);
  }

+  /**
+   * Returns true if an entity is soft-deleted.
+   *
+   * @param urn the urn to check
+   * @return true is the entity is soft deleted, false otherwise.
+   */
+  public Boolean isSoftDeleted(@Nonnull final Urn urn) {
+    Objects.requireNonNull(urn, "urn is required");
+    final RecordTemplate statusAspect = getLatestAspect(urn, STATUS_ASPECT_NAME);
+    return statusAspect != null && ((Status) statusAspect).isRemoved();
+  }
+
  @Nullable
  public RollbackResult deleteAspect(String urn, String aspectName, @Nonnull Map<String, String> conditions, boolean hardDelete) {
    // Validate pre-conditions before running queries
--- a/metadata-io/src/main/java/com/linkedin/metadata/shared/ValidationUtils.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/shared/ValidationUtils.java
@ -119,8 +119,11 @@ public class ValidationUtils {

    final LineageRelationshipArray validatedRelationships = entityLineageResult.getRelationships().stream()
        .filter(relationship -> entityService.exists(relationship.getEntity()))
+        .filter(relationship -> !entityService.isSoftDeleted(relationship.getEntity()))
        .collect(Collectors.toCollection(LineageRelationshipArray::new));

+    validatedEntityLineageResult.setFiltered(
+        entityLineageResult.getRelationships().size() - validatedRelationships.size());
    validatedEntityLineageResult.setRelationships(validatedRelationships);

    return validatedEntityLineageResult;
--- 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
@ -93,6 +93,7 @@ public class SiblingGraphServiceTest {
    mockResult.setStart(0);
    mockResult.setTotal(200);
    mockResult.setCount(3);
+    mockResult.setFiltered(0);
    mockResult.setRelationships(relationships);

    when(_graphService.getLineage(
@ -137,6 +138,7 @@ public class SiblingGraphServiceTest {
    mockResult.setStart(0);
    mockResult.setTotal(200);
    mockResult.setCount(3);
+    mockResult.setFiltered(0);
    mockResult.setRelationships(relationships);

    when(_graphService.getLineage(
@ -261,6 +263,7 @@ public class SiblingGraphServiceTest {
    EntityLineageResult expectedResult = mockResult.clone();
    expectedResult.setTotal(3);
    expectedResult.setCount(2);
+    expectedResult.setFiltered(0);
    expectedResult.setRelationships(new LineageRelationshipArray(relationship1, relationship2));

    EntityLineageResult upstreamLineage = service.getLineage(datasetFourUrn, LineageDirection.UPSTREAM, 0, 100, 1);
@ -309,6 +312,7 @@ public class SiblingGraphServiceTest {
    expectedResult.setCount(3);
    expectedResult.setStart(0);
    expectedResult.setTotal(3);
+    expectedResult.setFiltered(0);
    expectedResult.setRelationships(expectedRelationships);

    mockResult.setStart(0);
@ -406,6 +410,7 @@ public class SiblingGraphServiceTest {
    expectedResult.setCount(2);
    expectedResult.setStart(0);
    expectedResult.setTotal(3);
+    expectedResult.setFiltered(0);
    expectedResult.setRelationships(expectedRelationships);

    mockResult.setStart(0);
@ -491,6 +496,7 @@ public class SiblingGraphServiceTest {
    expectedResult.setCount(1);
    expectedResult.setStart(0);
    expectedResult.setTotal(1);
+    expectedResult.setFiltered(0);
    expectedResult.setRelationships(expectedRelationships);

    mockResult.setStart(0);
--- a/metadata-models/src/main/pegasus/com/linkedin/metadata/graph/EntityLineageResult.pdl
+++ b/metadata-models/src/main/pegasus/com/linkedin/metadata/graph/EntityLineageResult.pdl
@ -19,6 +19,11 @@ record EntityLineageResult {
   */
  total: int

+  /**
+   * The number of results that were filtered out of the page (soft-deleted or non-existent)
+   */
+  filtered: optional int = 0
+
  /**
   * Relationships in the result set
   */
--- a/metadata-service/restli-api/src/main/snapshot/com.linkedin.lineage.relationships.snapshot.json
+++ b/metadata-service/restli-api/src/main/snapshot/com.linkedin.lineage.relationships.snapshot.json
@ -96,6 +96,11 @@
      "name" : "total",
      "type" : "int",
      "doc" : "Total number of results in the result set"
+    }, {
+      "name" : "filtered",
+      "type" : "int",
+      "doc" : "The number of results that were filtered out of the page (soft-deleted or non-existent)",
+      "optional" : true
    }, {
      "name" : "relationships",
      "type" : {