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,8 @@
 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.LabelJoinerBuilder;
 import org.hypertrace.graphql.label.schema.LabelResultSet;
 
 public class LabelFetcher extends InjectableDataFetcher<LabelResultSet> {
@@ -16,19 +16,20 @@ public LabelFetcher() {
   }
 
   static final class LabelFetcherImpl implements DataFetcher<CompletableFuture<LabelResultSet>> {
-    private final ContextualRequestBuilder requestBuilder;
+    private final LabelJoinerBuilder requestBuilder;
     private final LabelDao labelDao;
 
     @Inject
-    LabelFetcherImpl(ContextualRequestBuilder requestBuilder, LabelDao labelDao) {
+    LabelFetcherImpl(LabelJoinerBuilder requestBuilder, LabelDao labelDao) {
       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.joinLabelsWithEntities())
           .toCompletionStage()
           .toCompletableFuture();
     }

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

@@ -17,6 +17,11 @@ public <T> Single<Map<T, LabelResultSet>> joinLabels(
             Collection<T> joinSources, LabelIdGetter<T> labelIdGetter) {
           return Single.just(Collections.emptyMap());
         }
+
+        @Override
+        public Single<LabelResultSet> joinLabelsWithEntities() {
+          throw new UnsupportedOperationException();
+        }
       };
 
   /**
@@ -30,6 +35,13 @@ public <T> Single<Map<T, LabelResultSet>> joinLabels(
   <T> Single<Map<T, LabelResultSet>> joinLabels(
       Collection<T> joinSources, LabelIdGetter<T> labelIdGetter);
 
+  /**
+   * Produces a map of label result set to source ids
+   *
+   * @return A map of each source to its matching label result set
+   */
+  Single<LabelResultSet> joinLabelsWithEntities();
+
   @FunctionalInterface
   interface LabelIdGetter<T> {
     Single<List<String>> getLabelIds(T source);

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

@@ -1,8 +1,12 @@
 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 LabelJoinerBuilder {
   Single<LabelJoiner> build(GraphQlRequestContext context);
+
+  Single<LabelJoiner> build(
+      GraphQlRequestContext context, DataFetchingFieldSelectionSet selectionSet);
 }

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-impl/build.gradle.kts

@@ -27,6 +27,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

@@ -8,6 +8,9 @@
 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;
@@ -51,6 +54,30 @@ private Label convertLabel(org.hypertrace.label.config.service.v1.Label label) {
         Collections.emptyMap());
   }
 
+  public Single<Label> convertLabel(Label label, Map<String, EntityResultSet> entityResultMap) {
+    Label convertedLabel =
+        new DefaultLabel(
+            label.id(),
+            label.key(),
+            label.color(),
+            label.description(),
+            getLabeledEntitiesMap(entityResultMap));
+    return Single.just(convertedLabel);
+  }
+
+  private Map<String, List<LabeledEntity>> getLabeledEntitiesMap(
+      Map<String, EntityResultSet> entityResultMap) {
+    return entityResultMap.entrySet().stream()
+        .map(entry -> Map.entry(entry.getKey(), convertEntities(entry.getValue())))
+        .collect(CollectorUtils.immutableMapEntryCollector());
+  }
+
+  private List<LabeledEntity> convertEntities(EntityResultSet entityResultSet) {
+    return entityResultSet.results().stream()
+        .map(DefaultLabeledEntity::new)
+        .collect(Collectors.toList());
+  }
+
   @Value
   @Accessors(fluent = true)
   private static class DefaultLabelResultSet implements LabelResultSet {
@@ -90,12 +117,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);
     }
   }
 }