fix(graphql): revert Container.entities count-only batching (#18640) (#19104)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: gdatahub

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java

@@ -25,7 +25,6 @@
 import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.featureflags.FeatureFlags;
 import com.linkedin.datahub.graphql.generated.*;
-import com.linkedin.datahub.graphql.loaders.ContainerEntityCountsBatchLoader;
 import com.linkedin.datahub.graphql.loaders.DomainEntityCountsBatchLoader;
 import com.linkedin.datahub.graphql.plugins.SemanticSearchPlugin;
 import com.linkedin.datahub.graphql.resolvers.MeResolver;
@@ -961,9 +960,6 @@ public GraphQLEngine.Builder builder() {
     builder
         .addDataLoaders(loaderSuppliers(loadableTypes))
         .addDataLoader("Aspect", context -> createDataLoader(aspectType, context))
-        .addDataLoader(
-            ContainerEntityCountsBatchLoader.LOADER_NAME,
-            context -> ContainerEntityCountsBatchLoader.create(entityClient, context))
         .addDataLoader(
             DomainEntityCountsBatchLoader.LOADER_NAME,
             context -> DomainEntityCountsBatchLoader.create(entityClient, context))

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/loaders/ContainerEntityCountsBatchLoader.java

@@ -4,14 +4,11 @@
 import static com.linkedin.metadata.utils.CriterionUtils.buildCriterion;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.generated.Container;
 import com.linkedin.datahub.graphql.generated.ContainerEntitiesInput;
-import com.linkedin.datahub.graphql.generated.FacetFilterInput;
 import com.linkedin.datahub.graphql.generated.SearchResults;
-import com.linkedin.datahub.graphql.loaders.ContainerEntityCountsBatchLoader;
 import com.linkedin.datahub.graphql.types.mappers.UrnSearchResultsMapper;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.Constants;
@@ -23,21 +20,16 @@
 import com.linkedin.metadata.query.filter.Filter;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
-import graphql.schema.DataFetchingFieldSelectionSet;
-import graphql.schema.SelectedField;
 import java.util.Collections;
 import java.util.List;
-import java.util.Set;
 import java.util.concurrent.CompletableFuture;
-import javax.annotation.Nullable;
 import lombok.extern.slf4j.Slf4j;
-import org.dataloader.DataLoader;
 
 /** Retrieves a list of historical executions for a particular source. */
 @Slf4j
 public class ContainerEntitiesResolver implements DataFetcher<CompletableFuture<SearchResults>> {
 
-  public static final List<String> CONTAINABLE_ENTITY_NAMES =
+  static final List<String> CONTAINABLE_ENTITY_NAMES =
       ImmutableList.of(
           Constants.DATASET_ENTITY_NAME,
           Constants.CHART_ENTITY_NAME,
@@ -50,11 +42,6 @@ public class ContainerEntitiesResolver implements DataFetcher<CompletableFuture<
   private static final Integer DEFAULT_COUNT = 20;
   private static final ContainerEntitiesInput DEFAULT_ENTITIES_INPUT = new ContainerEntitiesInput();
 
-  // Selections answerable from an aggregation alone. Deliberately excludes start/count: those are
-  // echoed from the request rather than reported by the search layer on this path, so we only take
-  // it when the caller cannot observe the difference.
-  private static final Set<String> COUNT_ONLY_FIELDS = ImmutableSet.of("total", "__typename");
-
   static {
     DEFAULT_ENTITIES_INPUT.setQuery(DEFAULT_QUERY);
     DEFAULT_ENTITIES_INPUT.setStart(DEFAULT_START);
@@ -79,89 +66,16 @@ public CompletableFuture<SearchResults> get(final DataFetchingEnvironment enviro
             ? bindArgument(environment.getArgument(INPUT_ARG_NAME), ContainerEntitiesInput.class)
             : DEFAULT_ENTITIES_INPUT;
 
-    final String query = input.getQuery() != null ? input.getQuery() : DEFAULT_QUERY;
-    final int start = input.getStart() != null ? input.getStart() : DEFAULT_START;
-    final int count = input.getCount() != null ? input.getCount() : DEFAULT_COUNT;
-
-    // Fast path: every UI call site selects only `total`, so the hits the search would return are
-    // discarded. Serve those from a batched, request-scoped aggregation instead, so a page of N
-    // containers costs a fixed number of searches rather than N searches plus N primary-store
-    // existence checks. See ContainerEntityCountsBatchLoader.
-    if (canServeFromCounts(environment, query, input.getFilters())) {
-      return resolveCountFromLoader(environment, urn, start, count);
-    }
-
-    return resolveDirect(context, urn, input, query, start, count);
-  }
-
-  private static boolean canServeFromCounts(
-      final DataFetchingEnvironment environment,
-      final String query,
-      @Nullable final List<FacetFilterInput> filters) {
-    // The batched loader forces query "*" and applies no facet filters, so anything relying on a
-    // real query or filters must take the direct path to preserve exact behavior. `count` is not
-    // consulted: the gate below already establishes that no hits are read, and `total` is
-    // independent of paging.
-    if (!DEFAULT_QUERY.equals(query) || (filters != null && !filters.isEmpty())) {
-      return false;
-    }
-    return isCountOnlySelection(environment);
-  }
-
-  /**
-   * True when the caller reads nothing but the total. Uses the flattened selection set so fragment
-   * spreads are resolved, and treats an empty selection as ineligible rather than as trivially
-   * count-only.
-   */
-  private static boolean isCountOnlySelection(final DataFetchingEnvironment environment) {
-    final DataFetchingFieldSelectionSet selectionSet = environment.getSelectionSet();
-    if (selectionSet == null) {
-      return false;
-    }
-    final List<SelectedField> fields = selectionSet.getImmediateFields();
-    return !fields.isEmpty()
-        && fields.stream().allMatch(field -> COUNT_ONLY_FIELDS.contains(field.getName()));
-  }
-
-  private CompletableFuture<SearchResults> resolveCountFromLoader(
-      final DataFetchingEnvironment environment,
-      final String containerUrn,
-      final int start,
-      final int count) {
-    final DataLoader<String, Long> loader =
-        environment.getDataLoader(ContainerEntityCountsBatchLoader.LOADER_NAME);
-    return loader.load(containerUrn).thenApply(total -> toCountOnlyResults(total, start, count));
-  }
-
-  private static SearchResults toCountOnlyResults(
-      final Long total, final int start, final int count) {
-    final SearchResults results = new SearchResults();
-    results.setStart(start);
-    results.setCount(count);
-    results.setTotal(total != null ? total.intValue() : 0);
-    results.setSearchResults(Collections.emptyList());
-    return results;
-  }
+    final String query = input.getQuery() != null ? input.getQuery() : "*";
+    final int start = input.getStart() != null ? input.getStart() : 0;
+    final int count = input.getCount() != null ? input.getCount() : 20;
 
-  /** The original, unbatched behavior: one search per invocation. */
-  private CompletableFuture<SearchResults> resolveDirect(
-      final QueryContext context,
-      final String urn,
-      final ContainerEntitiesInput input,
-      final String query,
-      final int start,
-      final int count) {
     return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
           try {
 
-            final CriterionArray criteria = new CriterionArray();
             final Criterion filterCriterion =
                 buildCriterion(CONTAINER_FIELD_NAME + ".keyword", Condition.EQUAL, urn);
-            criteria.add(filterCriterion);
-            if (input.getFilters() != null) {
-              input.getFilters().forEach(filter -> criteria.add(criterionFromFilter(filter)));
-            }
 
             return UrnSearchResultsMapper.map(
                 context,
@@ -172,7 +86,9 @@ private CompletableFuture<SearchResults> resolveDirect(
                     new Filter()
                         .setOr(
                             new ConjunctiveCriterionArray(
-                                new ConjunctiveCriterion().setAnd(criteria))),
+                                new ConjunctiveCriterion()
+                                    .setAnd(
+                                        new CriterionArray(ImmutableList.of(filterCriterion))))),
                     start,
                     count,
                     Collections.emptyList()));