Added support to query labels on entity schema (#101)

Author: saxenakshitiz

hypertrace-graphql-entity-schema/build.gradle.kts

@@ -29,6 +29,8 @@ dependencies {
   implementation("org.hypertrace.core.graphql:hypertrace-core-graphql-deserialization")
   implementation("org.hypertrace.core.graphql:hypertrace-core-graphql-rx-utils")
 
+  implementation(project(":hypertrace-graphql-labels-schema"))
+
   testImplementation("org.junit.jupiter:junit-jupiter")
   testImplementation("org.mockito:mockito-core")
   testImplementation("org.mockito:mockito-junit-jupiter")

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

@@ -28,6 +28,7 @@
 import org.hypertrace.gateway.service.v1.common.Value;
 import org.hypertrace.gateway.service.v1.entity.Entity;
 import org.hypertrace.graphql.entity.health.BaselineDao;
+import org.hypertrace.graphql.label.joiner.LabelJoinerBuilder;
 import org.hypertrace.graphql.metric.request.MetricAggregationRequest;
 import org.hypertrace.graphql.metric.request.MetricRequest;
 import org.hypertrace.graphql.metric.request.MetricSeriesRequest;
@@ -46,6 +47,7 @@ protected void configure() {
     requireBinding(GrpcChannelRegistry.class);
     requireBinding(GrpcContextBuilder.class);
     requireBinding(Key.get(Scheduler.class, BoundedIoScheduler.class));
+    requireBinding(LabelJoinerBuilder.class);
 
     requireBinding(
         Key.get(new TypeLiteral<Converter<Collection<AttributeRequest>, Set<Expression>>>() {}));

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

@@ -26,6 +26,7 @@
 import org.hypertrace.graphql.entity.schema.Entity;
 import org.hypertrace.graphql.entity.schema.EntityResultSet;
 import org.hypertrace.graphql.entity.schema.EntityType;
+import org.hypertrace.graphql.label.schema.LabelResultSet;
 import org.hypertrace.graphql.metric.request.MetricRequest;
 import org.hypertrace.graphql.metric.schema.MetricContainer;
 
@@ -58,17 +59,23 @@ class GatewayServiceEntityConverter {
   }
 
   Single<EntityResultSet> convert(
-      EntityRequest request, EntitiesResponse response, BaselineEntitiesResponse baselineResponse) {
+      EntityRequest request,
+      EntitiesResponse response,
+      BaselineEntitiesResponse baselineResponse,
+      Map<org.hypertrace.gateway.service.v1.entity.Entity, LabelResultSet> labelResultSetMap) {
     return this.edgeLookupConverter
         .convert(request, response)
-        .flatMap(edgeLookup -> this.convert(request, response, baselineResponse, edgeLookup));
+        .flatMap(
+            edgeLookup ->
+                this.convert(request, response, baselineResponse, edgeLookup, labelResultSetMap));
   }
 
   private Single<EntityResultSet> convert(
       EntityRequest request,
       EntitiesResponse response,
       BaselineEntitiesResponse baselineResponse,
-      EdgeLookup edgeLookup) {
+      EdgeLookup edgeLookup,
+      Map<org.hypertrace.gateway.service.v1.entity.Entity, LabelResultSet> labelResultSetMap) {
     Map<String, BaselineEntity> baselineEntityMap = getBaselineEntityMap(baselineResponse);
     return Observable.fromIterable(response.getEntityList())
         .flatMapSingle(
@@ -78,7 +85,8 @@ private Single<EntityResultSet> convert(
                     entity,
                     getBaselineEntity(baselineEntityMap, entity.getId()),
                     edgeLookup.getIncoming().row(entity),
-                    edgeLookup.getOutgoing().row(entity)))
+                    edgeLookup.getOutgoing().row(entity),
+                    labelResultSetMap.get(entity)))
         .toList()
         .map(
             entities ->
@@ -103,7 +111,8 @@ private Single<Entity> convertEntity(
       org.hypertrace.gateway.service.v1.entity.Entity platformEntity,
       BaselineEntity baselineEntity,
       Map<String, EdgeResultSet> incomingEdges,
-      Map<String, EdgeResultSet> outgoingEdges) {
+      Map<String, EdgeResultSet> outgoingEdges,
+      LabelResultSet labelResultSet) {
     return zip(
         this.attributeMapConverter.convert(
             entityRequest.resultSetRequest().attributes(), platformEntity.getAttributeMap()),
@@ -118,7 +127,8 @@ private Single<Entity> convertEntity(
                 attrMap,
                 containerMap,
                 incomingEdges,
-                outgoingEdges));
+                outgoingEdges,
+                labelResultSet));
   }
 
   @lombok.Value
@@ -130,6 +140,7 @@ private static class ConvertedEntity implements Entity {
     Map<String, MetricContainer> metricContainers;
     Map<String, EdgeResultSet> incomingEdges;
     Map<String, EdgeResultSet> outgoingEdges;
+    LabelResultSet labels;
 
     @Override
     public Object attribute(String key) {

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

@@ -5,30 +5,41 @@
 import io.grpc.CallCredentials;
 import io.reactivex.rxjava3.core.Scheduler;
 import io.reactivex.rxjava3.core.Single;
+import java.util.Collections;
+import java.util.Map;
 import javax.inject.Inject;
 import javax.inject.Singleton;
+import lombok.extern.slf4j.Slf4j;
 import org.hypertrace.core.graphql.context.GraphQlRequestContext;
 import org.hypertrace.core.graphql.rx.BoundedIoScheduler;
 import org.hypertrace.core.graphql.spi.config.GraphQlServiceConfig;
 import org.hypertrace.core.graphql.utils.grpc.GrpcChannelRegistry;
 import org.hypertrace.core.graphql.utils.grpc.GrpcContextBuilder;
 import org.hypertrace.gateway.service.GatewayServiceGrpc;
 import org.hypertrace.gateway.service.GatewayServiceGrpc.GatewayServiceFutureStub;
+import org.hypertrace.gateway.service.v1.common.Value;
 import org.hypertrace.gateway.service.v1.entity.EntitiesRequest;
 import org.hypertrace.gateway.service.v1.entity.EntitiesResponse;
+import org.hypertrace.gateway.service.v1.entity.Entity;
 import org.hypertrace.graphql.entity.health.BaselineDao;
 import org.hypertrace.graphql.entity.request.EntityRequest;
 import org.hypertrace.graphql.entity.schema.EntityResultSet;
+import org.hypertrace.graphql.label.joiner.LabelJoiner;
+import org.hypertrace.graphql.label.joiner.LabelJoinerBuilder;
+import org.hypertrace.graphql.label.schema.LabelResultSet;
 
+@Slf4j
 @Singleton
 class GatewayServiceEntityDao implements EntityDao {
+
   private final GatewayServiceFutureStub gatewayServiceStub;
   private final GrpcContextBuilder grpcContextBuilder;
   private final GatewayServiceEntityRequestBuilder requestBuilder;
   private final GatewayServiceEntityConverter entityConverter;
   private final BaselineDao baselineDao;
   private final GraphQlServiceConfig serviceConfig;
   private final Scheduler boundedIoScheduler;
+  private final LabelJoinerBuilder labelJoinerBuilder;
 
   @Inject
   GatewayServiceEntityDao(
@@ -39,11 +50,13 @@ class GatewayServiceEntityDao implements EntityDao {
       GatewayServiceEntityRequestBuilder requestBuilder,
       GatewayServiceEntityConverter entityConverter,
       BaselineDao baselineDao,
+      LabelJoinerBuilder labelJoinerBuilder,
       @BoundedIoScheduler Scheduler boundedIoScheduler) {
     this.grpcContextBuilder = grpcContextBuilder;
     this.requestBuilder = requestBuilder;
     this.entityConverter = entityConverter;
     this.baselineDao = baselineDao;
+    this.labelJoinerBuilder = labelJoinerBuilder;
     this.serviceConfig = serviceConfig;
     this.boundedIoScheduler = boundedIoScheduler;
 
@@ -56,27 +69,32 @@ class GatewayServiceEntityDao implements EntityDao {
 
   @Override
   public Single<EntityResultSet> getEntities(EntityRequest request) {
+    GraphQlRequestContext context = request.resultSetRequest().context();
     return this.requestBuilder
         .buildRequest(request)
         .subscribeOn(this.boundedIoScheduler)
-        .flatMap(
-            serverRequest ->
-                this.makeRequest(request.resultSetRequest().context(), serverRequest)
-                    .flatMap(
-                        serverResponse ->
-                            baselineDao
-                                .getBaselines(
-                                    request.resultSetRequest().context(),
-                                    serverRequest,
-                                    serverResponse,
-                                    request)
-                                .flatMap(
-                                    baselineResponse ->
-                                        this.entityConverter.convert(
-                                            request, serverResponse, baselineResponse))));
+        .flatMap(serverRequest -> this.fetchAndMapEntities(context, request, serverRequest));
+  }
+
+  private Single<EntityResultSet> fetchAndMapEntities(
+      GraphQlRequestContext context, EntityRequest request, EntitiesRequest serverRequest) {
+    return this.makeEntityRequest(context, serverRequest)
+        .flatMap(serverResponse -> this.getEntityResultSet(request, serverRequest, serverResponse));
+  }
+
+  private Single<EntityResultSet> getEntityResultSet(
+      EntityRequest request, EntitiesRequest serverRequest, EntitiesResponse serverResponse) {
+    GraphQlRequestContext context = request.resultSetRequest().context();
+    return Single.zip(
+            baselineDao.getBaselines(context, serverRequest, serverResponse, request),
+            buildLabelResultSetMap(context, request, serverResponse),
+            (baselineResponse, labelResultSetMap) ->
+                this.entityConverter.convert(
+                    request, serverResponse, baselineResponse, labelResultSetMap))
+        .flatMap(entityResultSet -> entityResultSet);
   }
 
-  private Single<EntitiesResponse> makeRequest(
+  private Single<EntitiesResponse> makeEntityRequest(
       GraphQlRequestContext context, EntitiesRequest request) {
     return Single.fromFuture(
         this.grpcContextBuilder
@@ -88,4 +106,29 @@ private Single<EntitiesResponse> makeRequest(
                             serviceConfig.getGatewayServiceTimeout().toMillis(), MILLISECONDS)
                         .getEntities(request)));
   }
+
+  private Single<Map<Entity, LabelResultSet>> buildLabelResultSetMap(
+      GraphQlRequestContext context, EntityRequest request, EntitiesResponse entitiesResponse) {
+    return request
+        .labelRequest()
+        .map(labelRequest -> labelJoinerBuilder.build(context))
+        .orElse(Single.just(LabelJoiner.NO_OP_JOINER))
+        .flatMap(
+            joiner ->
+                joiner.joinLabels(
+                    entitiesResponse.getEntityList(), getEntityLabelsGetter(request)));
+  }
+
+  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());
+    };
+  }
 }

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

@@ -5,6 +5,7 @@
 
 import io.reactivex.rxjava3.core.Single;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Set;
 import javax.inject.Inject;
@@ -67,19 +68,26 @@ Single<EntitiesRequest> buildRequest(EntityRequest entityRequest) {
                 .collect(flatten(MetricRequest::seriesRequests))),
         this.interactionRequestBuilder.build(entityRequest.incomingEdgeRequests()),
         this.interactionRequestBuilder.build(entityRequest.outgoingEdgeRequests()),
+        this.selectionConverter.convert(
+            entityRequest
+                .labelRequest()
+                .map(request -> Set.of(request.labelIdArrayAttributeRequest()))
+                .orElse(Collections.emptySet())),
         (selections,
             orderBys,
             filter,
             aggregations,
             series,
             incomingInteractions,
-            outgoingInteractions) ->
+            outgoingInteractions,
+            labelSelections) ->
             EntitiesRequest.newBuilder()
                 .setEntityType(entityRequest.entityType())
                 .setStartTimeMillis(resultSetRequest.timeRange().startTime().toEpochMilli())
                 .setEndTimeMillis(resultSetRequest.timeRange().endTime().toEpochMilli())
                 .addAllSelection(selections)
                 .addAllSelection(aggregations)
+                .addAllSelection(labelSelections)
                 .addAllTimeAggregation(series)
                 .setIncomingInteractions(incomingInteractions)
                 .setOutgoingInteractions(outgoingInteractions)

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

@@ -51,6 +51,7 @@
 import org.hypertrace.graphql.entity.dao.EntityDao;
 import org.hypertrace.graphql.entity.request.EdgeSetGroupRequest;
 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;
@@ -272,6 +273,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();
   }
 
   @Value

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

@@ -14,6 +14,8 @@
 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;
@@ -23,6 +25,7 @@
 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;
@@ -33,24 +36,29 @@
 
 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;
 
   @Inject
   DefaultEntityRequestBuilder(
       ResultSetRequestBuilder resultSetRequestBuilder,
       MetricRequestBuilder metricRequestBuilder,
       ArgumentDeserializer argumentDeserializer,
       GraphQlSelectionFinder selectionFinder,
-      EdgeRequestBuilder edgeRequestBuilder) {
+      EdgeRequestBuilder edgeRequestBuilder,
+      AttributeRequestBuilder attributeRequestBuilder) {
     this.resultSetRequestBuilder = resultSetRequestBuilder;
     this.metricRequestBuilder = metricRequestBuilder;
     this.argumentDeserializer = argumentDeserializer;
     this.selectionFinder = selectionFinder;
     this.edgeRequestBuilder = edgeRequestBuilder;
+    this.attributeRequestBuilder = attributeRequestBuilder;
   }
 
   @Override
@@ -88,6 +96,17 @@ private Single<EntityRequest> build(
                 .count()
             > 0;
 
+    boolean canFetchLabels =
+        this.selectionFinder
+                .findSelections(
+                    selectionSet,
+                    SelectionQuery.builder()
+                        .selectionPath(
+                            List.of(ResultSet.RESULT_SET_RESULTS_NAME, Entity.LABELS_KEY))
+                        .build())
+                .count()
+            > 0;
+
     return zip(
         this.resultSetRequestBuilder.build(
             context, scope, arguments, selectionSet, AggregatableOrderArgument.class),
@@ -102,15 +121,23 @@ private Single<EntityRequest> build(
             this.timeRange(arguments),
             this.space(arguments),
             this.getOutgoingEdges(selectionSet)),
-        (resultSetRequest, metricRequestList, incomingEdges, outgoingEdges) ->
+        attributeRequestBuilder.buildForKey(context, scope, LABELS_KEY_NAME),
+        (resultSetRequest,
+            metricRequestList,
+            incomingEdges,
+            outgoingEdges,
+            labelsAttributeRequest) ->
             new DefaultEntityRequest(
                 scope,
                 resultSetRequest,
                 metricRequestList,
                 incomingEdges,
                 outgoingEdges,
                 includeInactive,
-                fetchTotal));
+                fetchTotal,
+                canFetchLabels
+                    ? Optional.of(new DefaultLabelRequest(labelsAttributeRequest))
+                    : Optional.empty()));
   }
 
   private Stream<SelectedField> getResultSets(DataFetchingFieldSelectionSet selectionSet) {
@@ -154,5 +181,12 @@ 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;
   }
 }