ENG-11241: Changes to fetch labels as part of entity joiner

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/dao/GatewayServiceEntityDao.java

@@ -6,6 +6,7 @@
 import io.reactivex.rxjava3.core.Scheduler;
 import io.reactivex.rxjava3.core.Single;
 import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import javax.inject.Inject;
 import javax.inject.Singleton;
@@ -120,15 +121,17 @@ private Single<Map<Entity, LabelResultSet>> buildLabelResultSetMap(
   }
 
   private LabelJoiner.LabelIdGetter<Entity> getEntityLabelsGetter(EntityRequest request) {
-    return entity -> {
-      Value labelAttributeValue =
-          entity.getAttributeOrDefault(
-              request.labelRequest().get().labelIdArrayAttributeRequest().attribute().id(), null);
-      if (labelAttributeValue == null) {
-        log.warn("Unable to fetch labels attribute for entity with id {}", entity.getId());
-        return Single.just(Collections.emptyList());
-      }
-      return Single.just(labelAttributeValue.getStringArrayList());
-    };
+    return entity -> Single.just(getLabelAttributeValue(request, entity));
+  }
+
+  private List<String> getLabelAttributeValue(EntityRequest request, Entity entity) {
+    Value labelAttributeValue =
+        entity.getAttributeOrDefault(
+            request.labelRequest().get().labelIdArrayAttributeRequest().attribute().id(), null);
+    if (labelAttributeValue == null) {
+      log.warn("Unable to fetch labels attribute for entity with id {}", entity.getId());
+      return Collections.emptyList();
+    }
+    return labelAttributeValue.getStringArrayList();
   }
 }

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/joiner/DefaultEntityJoinerBuilder.java

@@ -3,6 +3,7 @@
 import static com.google.common.collect.ImmutableList.copyOf;
 import static com.google.common.collect.Iterables.concat;
 import static com.google.common.collect.Tables.immutableCell;
+import static io.reactivex.rxjava3.core.Single.zip;
 import static java.util.Objects.requireNonNull;
 
 import com.google.common.collect.ImmutableSetMultimap;
@@ -35,6 +36,7 @@
 import lombok.extern.slf4j.Slf4j;
 import org.hypertrace.core.graphql.common.request.AttributeAssociation;
 import org.hypertrace.core.graphql.common.request.AttributeRequest;
+import org.hypertrace.core.graphql.common.request.AttributeRequestBuilder;
 import org.hypertrace.core.graphql.common.request.FilterRequestBuilder;
 import org.hypertrace.core.graphql.common.request.ResultSetRequest;
 import org.hypertrace.core.graphql.common.request.ResultSetRequestBuilder;
@@ -50,8 +52,9 @@
 import org.hypertrace.core.graphql.utils.schema.SelectionQuery;
 import org.hypertrace.graphql.entity.dao.EntityDao;
 import org.hypertrace.graphql.entity.request.EdgeSetGroupRequest;
+import org.hypertrace.graphql.entity.request.EntityLabelRequest;
+import org.hypertrace.graphql.entity.request.EntityLabelRequestBuilder;
 import org.hypertrace.graphql.entity.request.EntityRequest;
-import org.hypertrace.graphql.entity.request.LabelRequest;
 import org.hypertrace.graphql.entity.schema.Entity;
 import org.hypertrace.graphql.entity.schema.EntityJoinable;
 import org.hypertrace.graphql.entity.schema.EntityResultSet;
@@ -62,15 +65,16 @@
 
 @Slf4j
 class DefaultEntityJoinerBuilder implements EntityJoinerBuilder {
-
   private static final int ZERO_OFFSET = 0;
 
   private final EntityDao entityDao;
   private final GraphQlSelectionFinder selectionFinder;
   private final ArgumentDeserializer argumentDeserializer;
   private final ResultSetRequestBuilder resultSetRequestBuilder;
   private final FilterRequestBuilder filterRequestBuilder;
+  private final AttributeRequestBuilder attributeRequestBuilder;
   private final Scheduler boundedIoScheduler;
+  private final EntityLabelRequestBuilder entityLabelRequestBuilder;
 
   @Inject
   DefaultEntityJoinerBuilder(
@@ -79,14 +83,18 @@ class DefaultEntityJoinerBuilder implements EntityJoinerBuilder {
       ArgumentDeserializer argumentDeserializer,
       ResultSetRequestBuilder resultSetRequestBuilder,
       FilterRequestBuilder filterRequestBuilder,
-      @BoundedIoScheduler Scheduler boundedIoScheduler) {
+      AttributeRequestBuilder attributeRequestBuilder,
+      @BoundedIoScheduler Scheduler boundedIoScheduler,
+      EntityLabelRequestBuilder entityLabelRequestBuilder) {
 
     this.entityDao = entityDao;
     this.selectionFinder = selectionFinder;
     this.argumentDeserializer = argumentDeserializer;
     this.resultSetRequestBuilder = resultSetRequestBuilder;
     this.filterRequestBuilder = filterRequestBuilder;
+    this.attributeRequestBuilder = attributeRequestBuilder;
     this.boundedIoScheduler = boundedIoScheduler;
+    this.entityLabelRequestBuilder = entityLabelRequestBuilder;
   }
 
   @Override
@@ -241,19 +249,30 @@ Single<EntityRequest> buildEntityRequestForType(
       return buildIdFilter(context, entityType, entityIdsToFilter)
           .flatMap(
               filterArguments ->
-                  resultSetRequestBuilder.build(
-                      context,
-                      entityType,
-                      entityIdsToFilter.size(),
-                      ZERO_OFFSET,
-                      new InstantTimeRange(),
-                      List
-                          .<AttributeAssociation<AggregatableOrderArgument>>
-                              of(), // Order does not matter
-                      filterArguments,
-                      this.entityFieldsByType.get(entityType).stream(),
-                      Optional.empty()))
-          .map(resultSetRequest -> new DefaultEntityRequest(entityType, resultSetRequest));
+                  buildEntityRequest(
+                      context, entityType, entityIdsToFilter.size(), filterArguments));
+    }
+
+    private Single<EntityRequest> buildEntityRequest(
+        GraphQlRequestContext context,
+        String entityType,
+        int entityIdsToFilterSize,
+        List<AttributeAssociation<FilterArgument>> filterArguments) {
+      return zip(
+          resultSetRequestBuilder.build(
+              context,
+              entityType,
+              entityIdsToFilterSize,
+              ZERO_OFFSET,
+              new InstantTimeRange(),
+              List.<AttributeAssociation<AggregatableOrderArgument>>of(), // Order does not matter
+              filterArguments,
+              this.entityFieldsByType.get(entityType).stream(),
+              Optional.empty()),
+          entityLabelRequestBuilder.buildLabelRequestIfPresentInAnyEntity(
+              context, entityType, this.entityFieldsByType.get(entityType)),
+          (resultSetRequest, optionalLabelRequest) ->
+              new DefaultEntityRequest(entityType, resultSetRequest, optionalLabelRequest));
     }
 
     private Single<List<AttributeAssociation<FilterArgument>>> buildIdFilter(
@@ -273,7 +292,7 @@ private static class DefaultEntityRequest implements EntityRequest {
     EdgeSetGroupRequest incomingEdgeRequests = new EmptyEdgeSetGroupRequest();
     EdgeSetGroupRequest outgoingEdgeRequests = new EmptyEdgeSetGroupRequest();
     boolean includeInactive = true; // When joining we want the entity regardless of time range
-    Optional<LabelRequest> labelRequest = Optional.empty();
+    Optional<EntityLabelRequest> labelRequest;
   }
 
   @Value

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/DefaultEntityLabelRequestBuilder.java

@@ -0,0 +1,79 @@
+package org.hypertrace.graphql.entity.request;
+
+import graphql.schema.DataFetchingFieldSelectionSet;
+import graphql.schema.SelectedField;
+import io.reactivex.rxjava3.core.Single;
+import java.util.Collection;
+import java.util.List;
+import java.util.Optional;
+import javax.inject.Inject;
+import lombok.Value;
+import lombok.experimental.Accessors;
+import org.hypertrace.core.graphql.common.request.AttributeRequest;
+import org.hypertrace.core.graphql.common.request.AttributeRequestBuilder;
+import org.hypertrace.core.graphql.common.schema.results.ResultSet;
+import org.hypertrace.core.graphql.context.GraphQlRequestContext;
+import org.hypertrace.core.graphql.utils.schema.GraphQlSelectionFinder;
+import org.hypertrace.core.graphql.utils.schema.SelectionQuery;
+import org.hypertrace.graphql.entity.schema.Entity;
+
+class DefaultEntityLabelRequestBuilder implements EntityLabelRequestBuilder {
+
+  private static final String LABELS_ATTRIBUTE_KEY = "labels";
+
+  private static final SelectionQuery RESULT_SET_LABEL_FIELD_QUERY =
+      SelectionQuery.builder()
+          .selectionPath(List.of(ResultSet.RESULT_SET_RESULTS_NAME, Entity.LABELS_KEY))
+          .build();
+
+  private static final SelectionQuery LABEL_FIELD_QUERY =
+      SelectionQuery.namedChild(Entity.LABELS_KEY);
+
+  private final AttributeRequestBuilder attributeRequestBuilder;
+  private final GraphQlSelectionFinder selectionFinder;
+
+  @Inject
+  DefaultEntityLabelRequestBuilder(
+      AttributeRequestBuilder attributeRequestBuilder, GraphQlSelectionFinder selectionFinder) {
+    this.attributeRequestBuilder = attributeRequestBuilder;
+    this.selectionFinder = selectionFinder;
+  }
+
+  @Override
+  public Single<Optional<EntityLabelRequest>> buildLabelRequestIfPresentInResultSet(
+      GraphQlRequestContext context, String scope, DataFetchingFieldSelectionSet selectionSet) {
+    if (isLabelFieldRequested(selectionSet, RESULT_SET_LABEL_FIELD_QUERY)) {
+      return buildRequest(context, scope);
+    }
+    return Single.just(Optional.empty());
+  }
+
+  @Override
+  public Single<Optional<EntityLabelRequest>> buildLabelRequestIfPresentInAnyEntity(
+      GraphQlRequestContext context, String scope, Collection<SelectedField> entityFields) {
+    if (entityFields.stream()
+        .anyMatch(field -> isLabelFieldRequested(field.getSelectionSet(), LABEL_FIELD_QUERY))) {
+      return buildRequest(context, scope);
+    }
+    return Single.just(Optional.empty());
+  }
+
+  private boolean isLabelFieldRequested(
+      DataFetchingFieldSelectionSet selectionSet, SelectionQuery labelFieldQuery) {
+    return this.selectionFinder.findSelections(selectionSet, labelFieldQuery).findAny().isPresent();
+  }
+
+  private Single<Optional<EntityLabelRequest>> buildRequest(
+      GraphQlRequestContext context, String scope) {
+    return this.attributeRequestBuilder
+        .buildForKey(context, scope, LABELS_ATTRIBUTE_KEY)
+        .map(DefaultLabelRequest::new)
+        .map(Optional::of);
+  }
+
+  @Value
+  @Accessors(fluent = true)
+  private static class DefaultLabelRequest implements EntityLabelRequest {
+    AttributeRequest labelIdArrayAttributeRequest;
+  }
+}

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/DefaultEntityRequestBuilder.java

@@ -14,8 +14,6 @@
 import javax.inject.Inject;
 import lombok.Value;
 import lombok.experimental.Accessors;
-import org.hypertrace.core.graphql.common.request.AttributeRequest;
-import org.hypertrace.core.graphql.common.request.AttributeRequestBuilder;
 import org.hypertrace.core.graphql.common.request.ResultSetRequest;
 import org.hypertrace.core.graphql.common.request.ResultSetRequestBuilder;
 import org.hypertrace.core.graphql.common.schema.arguments.TimeRangeArgument;
@@ -25,7 +23,6 @@
 import org.hypertrace.core.graphql.deserialization.ArgumentDeserializer;
 import org.hypertrace.core.graphql.utils.schema.GraphQlSelectionFinder;
 import org.hypertrace.core.graphql.utils.schema.SelectionQuery;
-import org.hypertrace.graphql.entity.schema.Entity;
 import org.hypertrace.graphql.entity.schema.EntityType;
 import org.hypertrace.graphql.entity.schema.argument.EntityScopeArgument;
 import org.hypertrace.graphql.entity.schema.argument.EntityTypeArgument;
@@ -36,14 +33,12 @@
 
 class DefaultEntityRequestBuilder implements EntityRequestBuilder {
 
-  private static final String LABELS_KEY_NAME = "labels";
-
   private final ResultSetRequestBuilder resultSetRequestBuilder;
   private final MetricRequestBuilder metricRequestBuilder;
   private final ArgumentDeserializer argumentDeserializer;
   private final GraphQlSelectionFinder selectionFinder;
   private final EdgeRequestBuilder edgeRequestBuilder;
-  private final AttributeRequestBuilder attributeRequestBuilder;
+  private final EntityLabelRequestBuilder entityLabelRequestBuilder;
 
   @Inject
   DefaultEntityRequestBuilder(
@@ -52,13 +47,13 @@ class DefaultEntityRequestBuilder implements EntityRequestBuilder {
       ArgumentDeserializer argumentDeserializer,
       GraphQlSelectionFinder selectionFinder,
       EdgeRequestBuilder edgeRequestBuilder,
-      AttributeRequestBuilder attributeRequestBuilder) {
+      EntityLabelRequestBuilder entityLabelRequestBuilder) {
     this.resultSetRequestBuilder = resultSetRequestBuilder;
     this.metricRequestBuilder = metricRequestBuilder;
     this.argumentDeserializer = argumentDeserializer;
     this.selectionFinder = selectionFinder;
     this.edgeRequestBuilder = edgeRequestBuilder;
-    this.attributeRequestBuilder = attributeRequestBuilder;
+    this.entityLabelRequestBuilder = entityLabelRequestBuilder;
   }
 
   @Override
@@ -110,12 +105,13 @@ private Single<EntityRequest> build(
             this.timeRange(arguments),
             this.space(arguments),
             this.getOutgoingEdges(selectionSet)),
-        buildLabelRequest(context, scope, selectionSet),
+        this.entityLabelRequestBuilder.buildLabelRequestIfPresentInResultSet(
+            context, scope, selectionSet),
         (resultSetRequest,
             metricRequestList,
             incomingEdges,
             outgoingEdges,
-            labelsAttributeRequest) ->
+            optionalLabelsAttributeRequest) ->
             new DefaultEntityRequest(
                 scope,
                 resultSetRequest,
@@ -124,28 +120,7 @@ private Single<EntityRequest> build(
                 outgoingEdges,
                 includeInactive,
                 fetchTotal,
-                labelsAttributeRequest));
-  }
-
-  private Single<Optional<LabelRequest>> buildLabelRequest(
-      GraphQlRequestContext context, String scope, DataFetchingFieldSelectionSet selectionSet) {
-    boolean canFetchLabels =
-        this.selectionFinder
-                .findSelections(
-                    selectionSet,
-                    SelectionQuery.builder()
-                        .selectionPath(
-                            List.of(ResultSet.RESULT_SET_RESULTS_NAME, Entity.LABELS_KEY))
-                        .build())
-                .count()
-            > 0;
-
-    return canFetchLabels
-        ? attributeRequestBuilder
-            .buildForKey(context, scope, LABELS_KEY_NAME)
-            .map(DefaultLabelRequest::new)
-            .map(Optional::of)
-        : Single.just(Optional.empty());
+                optionalLabelsAttributeRequest));
   }
 
   private Stream<SelectedField> getResultSets(DataFetchingFieldSelectionSet selectionSet) {
@@ -189,12 +164,6 @@ private static class DefaultEntityRequest implements EntityRequest {
     EdgeSetGroupRequest outgoingEdgeRequests;
     boolean includeInactive;
     boolean fetchTotal;
-    Optional<LabelRequest> labelRequest;
-  }
-
-  @Value
-  @Accessors(fluent = true)
-  private static class DefaultLabelRequest implements LabelRequest {
-    AttributeRequest labelIdArrayAttributeRequest;
+    Optional<EntityLabelRequest> labelRequest;
   }
 }

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/LabelRequest.java → hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityLabelRequest.java

@@ -2,7 +2,7 @@
 
 import org.hypertrace.core.graphql.common.request.AttributeRequest;
 
-public interface LabelRequest {
+public interface EntityLabelRequest {
 
   AttributeRequest labelIdArrayAttributeRequest();
 }

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityLabelRequestBuilder.java

@@ -0,0 +1,17 @@
+package org.hypertrace.graphql.entity.request;
+
+import graphql.schema.DataFetchingFieldSelectionSet;
+import graphql.schema.SelectedField;
+import io.reactivex.rxjava3.core.Single;
+import java.util.Collection;
+import java.util.Optional;
+import org.hypertrace.core.graphql.context.GraphQlRequestContext;
+
+public interface EntityLabelRequestBuilder {
+
+  Single<Optional<EntityLabelRequest>> buildLabelRequestIfPresentInResultSet(
+      GraphQlRequestContext context, String scope, DataFetchingFieldSelectionSet selectionSet);
+
+  Single<Optional<EntityLabelRequest>> buildLabelRequestIfPresentInAnyEntity(
+      GraphQlRequestContext context, String scope, Collection<SelectedField> entityField);
+}

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityRequest.java

@@ -21,5 +21,5 @@ public interface EntityRequest {
 
   boolean fetchTotal();
 
-  Optional<LabelRequest> labelRequest();
+  Optional<EntityLabelRequest> labelRequest();
 }

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityRequestModule.java

@@ -14,6 +14,7 @@ public class EntityRequestModule extends AbstractModule {
   @Override
   protected void configure() {
     bind(EntityRequestBuilder.class).to(DefaultEntityRequestBuilder.class);
+    bind(EntityLabelRequestBuilder.class).to(DefaultEntityLabelRequestBuilder.class);
     requireBinding(ResultSetRequestBuilder.class);
     requireBinding(ArgumentDeserializer.class);
     requireBinding(MetricRequestBuilder.class);

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/NeighborEntitiesRequestBuilder.java

@@ -179,6 +179,6 @@ private static class NeighborEntityRequest implements EntityRequest {
     EdgeSetGroupRequest outgoingEdgeRequests;
     boolean includeInactive;
     boolean fetchTotal;
-    Optional<LabelRequest> labelRequest;
+    Optional<EntityLabelRequest> labelRequest;
   }
 }