Changes to get entities and services count associated with label

hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/fetcher/LabelFetcher.java

@@ -5,8 +5,7 @@
 import java.util.concurrent.CompletableFuture;
 import javax.inject.Inject;
 import org.hypertrace.core.graphql.common.fetcher.InjectableDataFetcher;
-import org.hypertrace.core.graphql.common.request.ContextualRequestBuilder;
-import org.hypertrace.graphql.label.dao.LabelDao;
+import org.hypertrace.graphql.label.joiner.EntityAndRuleJoinerBuilder;
 import org.hypertrace.graphql.label.schema.LabelResultSet;
 
 public class LabelFetcher extends InjectableDataFetcher<LabelResultSet> {
@@ -16,19 +15,18 @@ public LabelFetcher() {
   }
 
   static final class LabelFetcherImpl implements DataFetcher<CompletableFuture<LabelResultSet>> {
-    private final ContextualRequestBuilder requestBuilder;
-    private final LabelDao labelDao;
+    private final EntityAndRuleJoinerBuilder requestBuilder;
 
     @Inject
-    LabelFetcherImpl(ContextualRequestBuilder requestBuilder, LabelDao labelDao) {
+    LabelFetcherImpl(EntityAndRuleJoinerBuilder requestBuilder) {
       this.requestBuilder = requestBuilder;
-      this.labelDao = labelDao;
     }
 
     @Override
     public CompletableFuture<LabelResultSet> get(DataFetchingEnvironment environment) {
-      return this.labelDao
-          .getLabels(this.requestBuilder.build(environment.getContext()))
+      return this.requestBuilder
+          .build(environment.getContext(), environment.getSelectionSet())
+          .flatMap(request -> request.joinLabelsWithEntitiesAndRules())
           .toCompletionStage()
           .toCompletableFuture();
     }

hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/EntityAndRuleJoiner.java

@@ -0,0 +1,15 @@
+package org.hypertrace.graphql.label.joiner;
+
+import io.reactivex.rxjava3.core.Single;
+import org.hypertrace.graphql.label.schema.LabelResultSet;
+
+public interface EntityAndRuleJoiner {
+
+  /**
+   * Produces label result set after joining labels with associated entities and label application
+   * rules
+   *
+   * @return label result set
+   */
+  Single<LabelResultSet> joinLabelsWithEntitiesAndRules();
+}

hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/EntityAndRuleJoinerBuilder.java

@@ -0,0 +1,10 @@
+package org.hypertrace.graphql.label.joiner;
+
+import graphql.schema.DataFetchingFieldSelectionSet;
+import io.reactivex.rxjava3.core.Single;
+import org.hypertrace.core.graphql.context.GraphQlRequestContext;
+
+public interface EntityAndRuleJoinerBuilder {
+  Single<EntityAndRuleJoiner> build(
+      GraphQlRequestContext context, DataFetchingFieldSelectionSet selectionSet);
+}

hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelApplicationRulesRequest.java

@@ -0,0 +1,5 @@
+package org.hypertrace.graphql.label.joiner;
+
+public interface LabelApplicationRulesRequest {
+  int limit();
+}

hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabeledEntitiesRequest.java

@@ -0,0 +1,11 @@
+package org.hypertrace.graphql.label.joiner;
+
+import graphql.schema.DataFetchingFieldSelectionSet;
+
+public interface LabeledEntitiesRequest {
+  String entityType();
+
+  int limit();
+
+  DataFetchingFieldSelectionSet selectionSet();
+}

hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/Label.java

@@ -1,10 +1,30 @@
 package org.hypertrace.graphql.label.schema;
 
+import graphql.annotations.annotationTypes.GraphQLField;
 import graphql.annotations.annotationTypes.GraphQLName;
+import graphql.annotations.annotationTypes.GraphQLNonNull;
 import org.hypertrace.core.graphql.common.schema.id.Identifiable;
+import org.hypertrace.core.graphql.common.schema.results.arguments.page.LimitArgument;
+import org.hypertrace.graphql.label.schema.rule.LabelApplicationRuleResultSet;
 
 @GraphQLName(Label.TYPE_NAME)
-public interface Label extends Identifiable, LabelData, LabeledEntities {
+public interface Label extends Identifiable, LabelData {
   String TYPE_NAME = "Label";
   String ARGUMENT_NAME = "label";
+  String LABELED_ENTITIES_QUERY_NAME = "labeledEntities";
+  String ENTITY_TYPE_ARGUMENT_NAME = "type";
+  String LABEL_APPLICATION_RULES_QUERY_NAME = "labelApplicationRules";
+
+  @GraphQLField
+  @GraphQLNonNull
+  @GraphQLName(LABELED_ENTITIES_QUERY_NAME)
+  LabeledEntityResultSet labeledEntities(
+      @GraphQLNonNull @GraphQLName(ENTITY_TYPE_ARGUMENT_NAME) String entityType,
+      @GraphQLName(LimitArgument.ARGUMENT_NAME) int limit);
+
+  @GraphQLField
+  @GraphQLNonNull
+  @GraphQLName(LABEL_APPLICATION_RULES_QUERY_NAME)
+  LabelApplicationRuleResultSet labelApplicationRules(
+      @GraphQLName(LimitArgument.ARGUMENT_NAME) int limit);
 }

hypertrace-graphql-labels-schema-impl/build.gradle.kts

@@ -28,6 +28,8 @@ dependencies {
   implementation("org.hypertrace.core.graphql:hypertrace-core-graphql-rx-utils")
   implementation("org.hypertrace.core.graphql:hypertrace-core-graphql-deserialization")
 
+  implementation(project(":hypertrace-graphql-entity-schema"))
+
   implementation(project(":hypertrace-graphql-service-config"))
 }
 

hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/dao/LabelResponseConverter.java

@@ -4,14 +4,17 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.stream.Collectors;
 import lombok.Value;
 import lombok.experimental.Accessors;
+import org.hypertrace.core.graphql.common.utils.CollectorUtils;
+import org.hypertrace.graphql.entity.schema.Entity;
+import org.hypertrace.graphql.entity.schema.EntityResultSet;
 import org.hypertrace.graphql.label.schema.Label;
 import org.hypertrace.graphql.label.schema.LabelResultSet;
 import org.hypertrace.graphql.label.schema.LabeledEntity;
 import org.hypertrace.graphql.label.schema.LabeledEntityResultSet;
+import org.hypertrace.graphql.label.schema.rule.LabelApplicationRuleResultSet;
 import org.hypertrace.label.config.service.v1.CreateLabelResponse;
 import org.hypertrace.label.config.service.v1.GetLabelsResponse;
 import org.hypertrace.label.config.service.v1.UpdateLabelResponse;
@@ -48,7 +51,42 @@ private Label convertLabel(org.hypertrace.label.config.service.v1.Label label) {
         label.getData().getKey(),
         label.getData().hasColor() ? label.getData().getColor() : null,
         label.getData().hasDescription() ? label.getData().getDescription() : null,
-        Collections.emptyMap());
+        Collections.emptyMap(),
+        null);
+  }
+
+  public Single<Label> convertLabel(
+      Label label,
+      Map<String, EntityResultSet> entityResultMap,
+      LabelApplicationRuleResultSet labelApplicationRuleResultSet) {
+    Label convertedLabel =
+        new DefaultLabel(
+            label.id(),
+            label.key(),
+            label.color(),
+            label.description(),
+            getLabeledEntityResultSetMap(entityResultMap),
+            labelApplicationRuleResultSet);
+    return Single.just(convertedLabel);
+  }
+
+  private Map<String, LabeledEntityResultSet> getLabeledEntityResultSetMap(
+      Map<String, EntityResultSet> entityResultMap) {
+    return entityResultMap.entrySet().stream()
+        .map(entry -> Map.entry(entry.getKey(), getLabeledEntityResultSet(entry.getValue())))
+        .collect(CollectorUtils.immutableMapEntryCollector());
+  }
+
+  private LabeledEntityResultSet getLabeledEntityResultSet(EntityResultSet entityResultSet) {
+    List<LabeledEntity> labeledEntities = convertEntities(entityResultSet);
+    return new DefaultLabeledEntityResultSet(
+        labeledEntities, labeledEntities.size(), entityResultSet.total());
+  }
+
+  private List<LabeledEntity> convertEntities(EntityResultSet entityResultSet) {
+    return entityResultSet.results().stream()
+        .map(DefaultLabeledEntity::new)
+        .collect(Collectors.toList());
   }
 
   @Value
@@ -66,16 +104,17 @@ private static class DefaultLabel implements Label {
     String key;
     String color;
     String description;
-    Map<String, List<LabeledEntity>> labeledEntitiesMap;
+    Map<String, LabeledEntityResultSet> labeledEntityResultSetMap;
+    LabelApplicationRuleResultSet labelApplicationRuleResultSet;
 
     @Override
     public LabeledEntityResultSet labeledEntities(String entityType, int limit) {
-      List<LabeledEntity> labeledEntities =
-          Optional.ofNullable(labeledEntitiesMap.get(entityType)).orElse(Collections.emptyList());
-      List<LabeledEntity> labeledEntitiesWithLimit =
-          labeledEntities.size() > limit ? labeledEntities.subList(0, limit) : labeledEntities;
-      return new DefaultLabeledEntityResultSet(
-          labeledEntitiesWithLimit, labeledEntitiesWithLimit.size(), labeledEntities.size());
+      return labeledEntityResultSetMap.get(entityType);
+    }
+
+    @Override
+    public LabelApplicationRuleResultSet labelApplicationRules(int limit) {
+      return labelApplicationRuleResultSet;
     }
   }
 
@@ -90,12 +129,16 @@ private static class DefaultLabeledEntityResultSet implements LabeledEntityResul
   @Value
   @Accessors(fluent = true)
   private static class DefaultLabeledEntity implements LabeledEntity {
-    String id;
-    Map<String, Object> attributeValues;
+    Entity entity;
+
+    @Override
+    public String id() {
+      return entity.id();
+    }
 
     @Override
     public Object attribute(String key) {
-      return this.attributeValues.get(key);
+      return this.entity.attribute(key);
     }
   }
 }