From a093d239972ac6c7a377e423a5bd3d388ef00ae8 Mon Sep 17 00:00:00 2001
From: saxenakshitiz <kshitiz@traceable.ai>
Date: Fri, 24 Dec 2021 10:57:11 +0530
Subject: [PATCH 1/3] Changes to get entities and services count associated
 with label

---
 .../graphql/label/fetcher/LabelFetcher.java   |  11 +-
 .../graphql/label/joiner/LabelJoiner.java     |  12 +
 .../label/joiner/LabelJoinerBuilder.java      |   4 +
 .../label/joiner/LabeledEntitiesRequest.java  |  11 +
 .../build.gradle.kts                          |   2 +
 .../label/dao/LabelResponseConverter.java     |  37 ++-
 .../joiner/DefaultLabelJoinerBuilder.java     | 234 +++++++++++++++++-
 7 files changed, 298 insertions(+), 13 deletions(-)
 create mode 100644 hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabeledEntitiesRequest.java

diff --git a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/fetcher/LabelFetcher.java b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/fetcher/LabelFetcher.java
index 604fbb18..c1d7d5da 100644
--- a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/fetcher/LabelFetcher.java
+++ b/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();
     }
diff --git a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoiner.java b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoiner.java
index e9f413e3..9c4a2cfb 100644
--- a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoiner.java
+++ b/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);
diff --git a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoinerBuilder.java b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoinerBuilder.java
index acb554d2..3f5f9533 100644
--- a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoinerBuilder.java
+++ b/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);
 }
diff --git a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabeledEntitiesRequest.java b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabeledEntitiesRequest.java
new file mode 100644
index 00000000..9f3fd0eb
--- /dev/null
+++ b/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();
+}
diff --git a/hypertrace-graphql-labels-schema-impl/build.gradle.kts b/hypertrace-graphql-labels-schema-impl/build.gradle.kts
index 641f70ec..99128c02 100644
--- a/hypertrace-graphql-labels-schema-impl/build.gradle.kts
+++ b/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"))
 }
 
diff --git a/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/dao/LabelResponseConverter.java b/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/dao/LabelResponseConverter.java
index 703ead87..08f0775c 100644
--- a/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/dao/LabelResponseConverter.java
+++ b/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);
     }
   }
 }
diff --git a/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultLabelJoinerBuilder.java b/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultLabelJoinerBuilder.java
index 9c8ba845..43a7dc27 100644
--- a/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultLabelJoinerBuilder.java
+++ b/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultLabelJoinerBuilder.java
@@ -1,58 +1,181 @@
 package org.hypertrace.graphql.label.joiner;
 
+import static io.reactivex.rxjava3.core.Single.zip;
 import static java.util.function.Function.identity;
+import static org.hypertrace.graphql.label.schema.LabeledEntities.ENTITY_TYPE_ARGUMENT_NAME;
+import static org.hypertrace.graphql.label.schema.LabeledEntities.LABELED_ENTITIES_QUERY_NAME;
 
+import graphql.schema.DataFetchingFieldSelectionSet;
 import io.reactivex.rxjava3.core.Observable;
 import io.reactivex.rxjava3.core.Single;
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
 import java.util.stream.Collectors;
 import javax.inject.Inject;
 import lombok.AllArgsConstructor;
+import lombok.Value;
+import lombok.experimental.Accessors;
 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.ContextualRequest;
 import org.hypertrace.core.graphql.common.request.ContextualRequestBuilder;
+import org.hypertrace.core.graphql.common.request.FilterRequestBuilder;
+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;
+import org.hypertrace.core.graphql.common.schema.attributes.AttributeScope;
+import org.hypertrace.core.graphql.common.schema.results.ResultSet;
+import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterArgument;
+import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterOperatorType;
+import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterType;
+import org.hypertrace.core.graphql.common.schema.results.arguments.page.LimitArgument;
 import org.hypertrace.core.graphql.common.utils.CollectorUtils;
 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.dao.EntityDao;
+import org.hypertrace.graphql.entity.request.EdgeSetGroupRequest;
+import org.hypertrace.graphql.entity.request.EntityLabelRequest;
+import org.hypertrace.graphql.entity.request.EntityRequest;
+import org.hypertrace.graphql.entity.schema.EntityResultSet;
 import org.hypertrace.graphql.label.dao.LabelDao;
 import org.hypertrace.graphql.label.dao.LabelResponseConverter;
 import org.hypertrace.graphql.label.schema.Label;
 import org.hypertrace.graphql.label.schema.LabelResultSet;
+import org.hypertrace.graphql.metric.request.MetricAggregationRequest;
+import org.hypertrace.graphql.metric.request.MetricRequest;
+import org.hypertrace.graphql.metric.schema.argument.AggregatableOrderArgument;
 
 @Slf4j
 class DefaultLabelJoinerBuilder implements LabelJoinerBuilder {
 
+  private static final Integer DEFAULT_LIMIT = 10;
+  private static final int ZERO_OFFSET = 0;
+
+  private static final SelectionQuery RESULT_SET_LABELED_ENTITIES_SELECTION_QUERY =
+      SelectionQuery.builder()
+          .selectionPath(List.of(ResultSet.RESULT_SET_RESULTS_NAME, LABELED_ENTITIES_QUERY_NAME))
+          .build();
+
   private final ContextualRequestBuilder requestBuilder;
+  private final ResultSetRequestBuilder resultSetRequestBuilder;
+  private final FilterRequestBuilder filterRequestBuilder;
   private final LabelDao labelDao;
   private final LabelResponseConverter responseConverter;
+  private final GraphQlSelectionFinder selectionFinder;
+  private final EntityDao entityDao;
 
   @Inject
   DefaultLabelJoinerBuilder(
       ContextualRequestBuilder requestBuilder,
+      ResultSetRequestBuilder resultSetRequestBuilder,
+      FilterRequestBuilder filterRequestBuilder,
       LabelDao labelDao,
-      LabelResponseConverter responseConverter) {
+      EntityDao entityDao,
+      LabelResponseConverter responseConverter,
+      GraphQlSelectionFinder selectionFinder) {
     this.requestBuilder = requestBuilder;
+    this.resultSetRequestBuilder = resultSetRequestBuilder;
+    this.filterRequestBuilder = filterRequestBuilder;
     this.labelDao = labelDao;
+    this.entityDao = entityDao;
     this.responseConverter = responseConverter;
+    this.selectionFinder = selectionFinder;
   }
 
   @Override
   public Single<LabelJoiner> build(GraphQlRequestContext context) {
-    return Single.just(new DefaultLabelJoiner(context));
+    return buildContextualRequest(context)
+        .map(request -> new DefaultLabelJoiner(context, request, null));
+  }
+
+  @Override
+  public Single<LabelJoiner> build(
+      GraphQlRequestContext context, DataFetchingFieldSelectionSet selectionSet) {
+    return zip(
+        buildContextualRequest(context),
+        buildLabeledEntitiesRequestIfPresentInResultSet(selectionSet),
+        (contextualRequest, labeledEntitiesRequests) ->
+            new DefaultLabelJoiner(context, contextualRequest, labeledEntitiesRequests));
+  }
+
+  private Single<ContextualRequest> buildContextualRequest(GraphQlRequestContext context) {
+    return Single.just(this.requestBuilder.build(context));
+  }
+
+  private Single<List<LabeledEntitiesRequest>> buildLabeledEntitiesRequestIfPresentInResultSet(
+      DataFetchingFieldSelectionSet selectionSet) {
+    return Single.just(
+        this.selectionFinder
+            .findSelections(selectionSet, RESULT_SET_LABELED_ENTITIES_SELECTION_QUERY)
+            .map(
+                selectedField ->
+                    new DefaultLabeledEntitiesRequest(
+                        (String) selectedField.getArguments().get(ENTITY_TYPE_ARGUMENT_NAME),
+                        (Integer)
+                            selectedField
+                                .getArguments()
+                                .getOrDefault(LimitArgument.ARGUMENT_NAME, DEFAULT_LIMIT),
+                        selectedField.getSelectionSet()))
+            .collect(Collectors.toList()));
   }
 
   @AllArgsConstructor
   private class DefaultLabelJoiner implements LabelJoiner {
     private final GraphQlRequestContext context;
+    private final ContextualRequest resultSetRequest;
+    private final List<LabeledEntitiesRequest> labeledEntitiesRequests;
 
     @Override
     public <T> Single<Map<T, LabelResultSet>> joinLabels(
         Collection<T> joinSources, LabelIdGetter<T> labelIdGetter) {
       return labelDao
-          .getLabels(requestBuilder.build(context))
+          .getLabels(resultSetRequest)
           .flatMap(labelResultSet -> buildLabelMap(labelResultSet.results()))
-          .flatMap(labelMap -> getLabelResultSetMap(joinSources, labelIdGetter, labelMap));
+          .flatMap(labelsMap -> getLabelResultSetMap(joinSources, labelIdGetter, labelsMap));
+    }
+
+    @Override
+    public Single<LabelResultSet> joinLabelsWithEntities() {
+      return labelDao.getLabels(resultSetRequest).flatMap(this::getUpdatedLabelResultSet);
+    }
+
+    private Single<LabelResultSet> getUpdatedLabelResultSet(LabelResultSet labelResultSet) {
+      return getUpdatedLabels(labelResultSet).flatMap(responseConverter::convert);
+    }
+
+    private Single<List<Label>> getUpdatedLabels(LabelResultSet labelResultSet) {
+      return Observable.fromIterable(labelResultSet.results())
+          .flatMapSingle(this::getUpdatedLabel)
+          .collect(Collectors.toList());
+    }
+
+    private Single<Label> getUpdatedLabel(Label label) {
+      return getEntityResultSetMap(label)
+          .flatMap(entityResultMap -> responseConverter.convertLabel(label, entityResultMap));
+    }
+
+    private Single<Map<String, EntityResultSet>> getEntityResultSetMap(Label label) {
+      String labelId = label.id();
+      return Observable.fromIterable(labeledEntitiesRequests)
+          .flatMapSingle(request -> getEntityResultSetMapEntry(labelId, request))
+          .collect(CollectorUtils.immutableMapEntryCollector());
+    }
+
+    private Single<Map.Entry<String, EntityResultSet>> getEntityResultSetMapEntry(
+        String labelId, LabeledEntitiesRequest request) {
+      return buildEntityRequestForType(
+              context, request.entityType(), labelId, request.limit(), request.selectionSet())
+          .flatMap(entityDao::getEntities)
+          .map(resultSet -> Map.entry(request.entityType(), resultSet));
     }
 
     private Single<Map<String, Label>> buildLabelMap(List<Label> results) {
@@ -72,7 +195,7 @@ private <T> Single<Map.Entry<T, LabelResultSet>> buildLabelResultSetMapEntry(
       return labelIdGetter
           .getLabelIds(source)
           .flatMap(labelIds -> filterLabels(labelIds, labelMap))
-          .flatMap(labels -> responseConverter.convert(labels))
+          .flatMap(responseConverter::convert)
           .map(labelResultSet -> Map.entry(source, labelResultSet));
     }
 
@@ -91,5 +214,106 @@ private Label getLabelFromMap(String labelId, Map<String, Label> labelMap) {
       }
       return label;
     }
+
+    private Single<EntityRequest> buildEntityRequestForType(
+        GraphQlRequestContext context,
+        String entityType,
+        String labelId,
+        int limit,
+        DataFetchingFieldSelectionSet selectionSet) {
+      return buildLabelIdFilter(context, entityType, labelId)
+          .flatMap(
+              filterArguments ->
+                  buildEntityRequest(context, entityType, limit, filterArguments, selectionSet));
+    }
+
+    private Single<EntityRequest> buildEntityRequest(
+        GraphQlRequestContext context,
+        String entityType,
+        int entityIdsToFilterSize,
+        List<AttributeAssociation<FilterArgument>> filterArguments,
+        DataFetchingFieldSelectionSet selectionSet) {
+      boolean fetchTotal =
+          selectionFinder
+                  .findSelections(
+                      selectionSet, SelectionQuery.namedChild(ResultSet.RESULT_SET_TOTAL_NAME))
+                  .count()
+              > 0;
+      return resultSetRequestBuilder
+          .build(
+              context,
+              entityType,
+              entityIdsToFilterSize,
+              ZERO_OFFSET,
+              new InstantTimeRange(),
+              List.<AttributeAssociation<AggregatableOrderArgument>>of(), // Order does not matter
+              filterArguments,
+              selectionSet.getFields().stream(),
+              Optional.empty())
+          .map(request -> new DefaultEntityRequest(entityType, request, fetchTotal));
+    }
+
+    private Single<List<AttributeAssociation<FilterArgument>>> buildLabelIdFilter(
+        GraphQlRequestContext context, String entityScope, String labelId) {
+      return filterRequestBuilder.build(context, entityScope, Set.of(new LabelIdFilter(labelId)));
+    }
+  }
+
+  @Value
+  @Accessors(fluent = true)
+  private static class DefaultEntityRequest implements EntityRequest {
+    String entityType;
+    ResultSetRequest<AggregatableOrderArgument> resultSetRequest;
+    List<MetricRequest> metricRequests = Collections.emptyList(); // Only support attributes for now
+    boolean fetchTotal;
+    EdgeSetGroupRequest incomingEdgeRequests = new EmptyEdgeSetGroupRequest();
+    EdgeSetGroupRequest outgoingEdgeRequests = new EmptyEdgeSetGroupRequest();
+    boolean includeInactive = true; // When joining we want the entity regardless of time range
+    Optional<EntityLabelRequest> labelRequest = Optional.empty();
+  }
+
+  @Value
+  @Accessors(fluent = true)
+  private static class EmptyEdgeSetGroupRequest implements EdgeSetGroupRequest {
+    Set<String> entityTypes = Collections.emptySet();
+    Collection<AttributeRequest> attributeRequests = Collections.emptyList();
+    Collection<MetricAggregationRequest> metricAggregationRequests = Collections.emptyList();
+    AttributeRequest neighborIdAttribute = null;
+    AttributeRequest neighborTypeAttribute = null;
+
+    @Override
+    public Single<EntityRequest> buildNeighborRequest(
+        String entityType, Collection<String> neighborIds) {
+      return Single.error(
+          new UnsupportedOperationException(
+              "Does not support fetching neighbors for joined entities"));
+    }
+  }
+
+  @Value
+  @Accessors(fluent = true)
+  private static class InstantTimeRange implements TimeRangeArgument {
+    // Bug in gateway requires this time range be larger, even though we don't care about it
+    Instant startTime = Instant.now().minus(15, ChronoUnit.MINUTES);
+    Instant endTime = Instant.now();
+  }
+
+  @Value
+  @Accessors(fluent = true)
+  private static class LabelIdFilter implements FilterArgument {
+    FilterType type = FilterType.ATTRIBUTE;
+    String key = "labels";
+    FilterOperatorType operator = FilterOperatorType.EQUALS;
+    String value;
+    AttributeScope idType = null;
+    String idScope = null;
+  }
+
+  @Value
+  @Accessors(fluent = true)
+  private static class DefaultLabeledEntitiesRequest implements LabeledEntitiesRequest {
+    String entityType;
+    int limit;
+    DataFetchingFieldSelectionSet selectionSet;
   }
 }

From 799a86cd02850d8476d505664f551afdcc1c4421 Mon Sep 17 00:00:00 2001
From: saxenakshitiz <kshitiz@traceable.ai>
Date: Mon, 27 Dec 2021 18:12:50 +0530
Subject: [PATCH 2/3] Address review comments. Also exposed rules count

---
 .../graphql/label/fetcher/LabelFetcher.java   |  11 +-
 .../label/joiner/EntityAndRuleJoiner.java     |  15 +
 .../joiner/EntityAndRuleJoinerBuilder.java    |  10 +
 .../joiner/LabelApplicationRulesRequest.java  |   5 +
 .../graphql/label/joiner/LabelJoiner.java     |  12 -
 .../label/joiner/LabelJoinerBuilder.java      |   4 -
 .../graphql/label/schema/Label.java           |   2 +-
 .../label/schema/LabelApplicationRules.java   |  19 +
 .../label/dao/LabelResponseConverter.java     |  38 +-
 .../DefaultEntityAndRuleJoinerBuilder.java    | 374 ++++++++++++++++++
 .../joiner/DefaultLabelJoinerBuilder.java     | 234 +----------
 .../label/joiner/LabelJoinerModule.java       |   1 +
 12 files changed, 459 insertions(+), 266 deletions(-)
 create mode 100644 hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/EntityAndRuleJoiner.java
 create mode 100644 hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/EntityAndRuleJoinerBuilder.java
 create mode 100644 hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelApplicationRulesRequest.java
 create mode 100644 hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/LabelApplicationRules.java
 create mode 100644 hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultEntityAndRuleJoinerBuilder.java

diff --git a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/fetcher/LabelFetcher.java b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/fetcher/LabelFetcher.java
index c1d7d5da..3381db12 100644
--- a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/fetcher/LabelFetcher.java
+++ b/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.graphql.label.dao.LabelDao;
-import org.hypertrace.graphql.label.joiner.LabelJoinerBuilder;
+import org.hypertrace.graphql.label.joiner.EntityAndRuleJoinerBuilder;
 import org.hypertrace.graphql.label.schema.LabelResultSet;
 
 public class LabelFetcher extends InjectableDataFetcher<LabelResultSet> {
@@ -16,20 +15,18 @@ public LabelFetcher() {
   }
 
   static final class LabelFetcherImpl implements DataFetcher<CompletableFuture<LabelResultSet>> {
-    private final LabelJoinerBuilder requestBuilder;
-    private final LabelDao labelDao;
+    private final EntityAndRuleJoinerBuilder requestBuilder;
 
     @Inject
-    LabelFetcherImpl(LabelJoinerBuilder requestBuilder, LabelDao labelDao) {
+    LabelFetcherImpl(EntityAndRuleJoinerBuilder requestBuilder) {
       this.requestBuilder = requestBuilder;
-      this.labelDao = labelDao;
     }
 
     @Override
     public CompletableFuture<LabelResultSet> get(DataFetchingEnvironment environment) {
       return this.requestBuilder
           .build(environment.getContext(), environment.getSelectionSet())
-          .flatMap(request -> request.joinLabelsWithEntities())
+          .flatMap(request -> request.joinLabelsWithEntitiesAndRules())
           .toCompletionStage()
           .toCompletableFuture();
     }
diff --git a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/EntityAndRuleJoiner.java b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/EntityAndRuleJoiner.java
new file mode 100644
index 00000000..03e2be2c
--- /dev/null
+++ b/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();
+}
diff --git a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/EntityAndRuleJoinerBuilder.java b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/EntityAndRuleJoinerBuilder.java
new file mode 100644
index 00000000..8a3fc42d
--- /dev/null
+++ b/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);
+}
diff --git a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelApplicationRulesRequest.java b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelApplicationRulesRequest.java
new file mode 100644
index 00000000..77989f94
--- /dev/null
+++ b/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();
+}
diff --git a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoiner.java b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoiner.java
index 9c4a2cfb..e9f413e3 100644
--- a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoiner.java
+++ b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoiner.java
@@ -17,11 +17,6 @@ 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();
-        }
       };
 
   /**
@@ -35,13 +30,6 @@ public Single<LabelResultSet> joinLabelsWithEntities() {
   <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);
diff --git a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoinerBuilder.java b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoinerBuilder.java
index 3f5f9533..acb554d2 100644
--- a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoinerBuilder.java
+++ b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoinerBuilder.java
@@ -1,12 +1,8 @@
 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);
 }
diff --git a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/Label.java b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/Label.java
index 059333e2..fc2728d6 100644
--- a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/Label.java
+++ b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/Label.java
@@ -4,7 +4,7 @@
 import org.hypertrace.core.graphql.common.schema.id.Identifiable;
 
 @GraphQLName(Label.TYPE_NAME)
-public interface Label extends Identifiable, LabelData, LabeledEntities {
+public interface Label extends Identifiable, LabelData, LabeledEntities, LabelApplicationRules {
   String TYPE_NAME = "Label";
   String ARGUMENT_NAME = "label";
 }
diff --git a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/LabelApplicationRules.java b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/LabelApplicationRules.java
new file mode 100644
index 00000000..44891e94
--- /dev/null
+++ b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/LabelApplicationRules.java
@@ -0,0 +1,19 @@
+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.results.arguments.page.LimitArgument;
+import org.hypertrace.graphql.label.schema.rule.LabelApplicationRuleResultSet;
+
+@GraphQLName(LabelApplicationRules.TYPE_NAME)
+public interface LabelApplicationRules {
+  String TYPE_NAME = "LabelApplicationRules";
+  String LABEL_APPLICATION_RULES_QUERY_NAME = "labelApplicationRules";
+
+  @GraphQLField
+  @GraphQLNonNull
+  @GraphQLName(LABEL_APPLICATION_RULES_QUERY_NAME)
+  LabelApplicationRuleResultSet labelApplicationRules(
+      @GraphQLName(LimitArgument.ARGUMENT_NAME) int limit);
+}
diff --git a/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/dao/LabelResponseConverter.java b/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/dao/LabelResponseConverter.java
index 08f0775c..e0322fae 100644
--- a/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/dao/LabelResponseConverter.java
+++ b/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/dao/LabelResponseConverter.java
@@ -4,7 +4,6 @@
 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;
@@ -15,6 +14,7 @@
 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;
@@ -51,27 +51,38 @@ 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) {
+  public Single<Label> convertLabel(
+      Label label,
+      Map<String, EntityResultSet> entityResultMap,
+      LabelApplicationRuleResultSet labelApplicationRuleResultSet) {
     Label convertedLabel =
         new DefaultLabel(
             label.id(),
             label.key(),
             label.color(),
             label.description(),
-            getLabeledEntitiesMap(entityResultMap));
+            getLabeledEntityResultSetMap(entityResultMap),
+            labelApplicationRuleResultSet);
     return Single.just(convertedLabel);
   }
 
-  private Map<String, List<LabeledEntity>> getLabeledEntitiesMap(
+  private Map<String, LabeledEntityResultSet> getLabeledEntityResultSetMap(
       Map<String, EntityResultSet> entityResultMap) {
     return entityResultMap.entrySet().stream()
-        .map(entry -> Map.entry(entry.getKey(), convertEntities(entry.getValue())))
+        .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)
@@ -93,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;
     }
   }
 
diff --git a/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultEntityAndRuleJoinerBuilder.java b/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultEntityAndRuleJoinerBuilder.java
new file mode 100644
index 00000000..80381128
--- /dev/null
+++ b/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultEntityAndRuleJoinerBuilder.java
@@ -0,0 +1,374 @@
+package org.hypertrace.graphql.label.joiner;
+
+import static io.reactivex.rxjava3.core.Single.zip;
+import static org.hypertrace.graphql.label.schema.LabelApplicationRules.LABEL_APPLICATION_RULES_QUERY_NAME;
+import static org.hypertrace.graphql.label.schema.LabeledEntities.LABELED_ENTITIES_QUERY_NAME;
+
+import graphql.schema.DataFetchingFieldSelectionSet;
+import graphql.schema.SelectedField;
+import io.reactivex.rxjava3.core.Observable;
+import io.reactivex.rxjava3.core.Single;
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+import javax.inject.Inject;
+import lombok.AllArgsConstructor;
+import lombok.Value;
+import lombok.experimental.Accessors;
+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.ContextualRequest;
+import org.hypertrace.core.graphql.common.request.ContextualRequestBuilder;
+import org.hypertrace.core.graphql.common.request.FilterRequestBuilder;
+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;
+import org.hypertrace.core.graphql.common.schema.attributes.AttributeScope;
+import org.hypertrace.core.graphql.common.schema.results.ResultSet;
+import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterArgument;
+import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterOperatorType;
+import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterType;
+import org.hypertrace.core.graphql.common.schema.results.arguments.page.LimitArgument;
+import org.hypertrace.core.graphql.common.utils.CollectorUtils;
+import org.hypertrace.core.graphql.context.GraphQlRequestContext;
+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.dao.EntityDao;
+import org.hypertrace.graphql.entity.request.EdgeSetGroupRequest;
+import org.hypertrace.graphql.entity.request.EntityLabelRequest;
+import org.hypertrace.graphql.entity.request.EntityRequest;
+import org.hypertrace.graphql.entity.schema.EntityResultSet;
+import org.hypertrace.graphql.entity.schema.argument.EntityTypeStringArgument;
+import org.hypertrace.graphql.label.dao.LabelApplicationRuleDao;
+import org.hypertrace.graphql.label.dao.LabelDao;
+import org.hypertrace.graphql.label.dao.LabelResponseConverter;
+import org.hypertrace.graphql.label.schema.Label;
+import org.hypertrace.graphql.label.schema.LabelResultSet;
+import org.hypertrace.graphql.label.schema.rule.LabelApplicationRule;
+import org.hypertrace.graphql.label.schema.rule.LabelApplicationRuleResultSet;
+import org.hypertrace.graphql.metric.request.MetricAggregationRequest;
+import org.hypertrace.graphql.metric.request.MetricRequest;
+import org.hypertrace.graphql.metric.schema.argument.AggregatableOrderArgument;
+
+@Slf4j
+class DefaultEntityAndRuleJoinerBuilder implements EntityAndRuleJoinerBuilder {
+
+  private static final Integer DEFAULT_LIMIT = 10;
+  private static final int ZERO_OFFSET = 0;
+
+  private static final SelectionQuery RESULT_SET_LABELED_ENTITIES_SELECTION_QUERY =
+      SelectionQuery.builder()
+          .selectionPath(List.of(ResultSet.RESULT_SET_RESULTS_NAME, LABELED_ENTITIES_QUERY_NAME))
+          .build();
+
+  private static final SelectionQuery RESULT_SET_LABEL_APPLICATION_RULES_SELECTION_QUERY =
+      SelectionQuery.builder()
+          .selectionPath(
+              List.of(ResultSet.RESULT_SET_RESULTS_NAME, LABEL_APPLICATION_RULES_QUERY_NAME))
+          .build();
+
+  private final LabelApplicationRuleResultSet EMPTY_LABEL_APPLICATION_RULE_RESULT_SET =
+      new DefaultLabelApplicationRuleResultSet(Collections.emptyList(), 0, 0);
+
+  private final ContextualRequestBuilder requestBuilder;
+  private final ResultSetRequestBuilder resultSetRequestBuilder;
+  private final FilterRequestBuilder filterRequestBuilder;
+  private final LabelDao labelDao;
+  private final LabelApplicationRuleDao labelApplicationRuleDao;
+  private final EntityDao entityDao;
+  private final LabelResponseConverter responseConverter;
+  private final GraphQlSelectionFinder selectionFinder;
+  private final ArgumentDeserializer argumentDeserializer;
+
+  @Inject
+  DefaultEntityAndRuleJoinerBuilder(
+      ContextualRequestBuilder requestBuilder,
+      ResultSetRequestBuilder resultSetRequestBuilder,
+      FilterRequestBuilder filterRequestBuilder,
+      LabelDao labelDao,
+      LabelApplicationRuleDao labelApplicationRuleDao,
+      EntityDao entityDao,
+      LabelResponseConverter responseConverter,
+      GraphQlSelectionFinder selectionFinder,
+      ArgumentDeserializer argumentDeserializer) {
+    this.requestBuilder = requestBuilder;
+    this.resultSetRequestBuilder = resultSetRequestBuilder;
+    this.filterRequestBuilder = filterRequestBuilder;
+    this.labelDao = labelDao;
+    this.labelApplicationRuleDao = labelApplicationRuleDao;
+    this.entityDao = entityDao;
+    this.responseConverter = responseConverter;
+    this.selectionFinder = selectionFinder;
+    this.argumentDeserializer = argumentDeserializer;
+  }
+
+  @Override
+  public Single<EntityAndRuleJoiner> build(
+      GraphQlRequestContext context, DataFetchingFieldSelectionSet selectionSet) {
+    return zip(
+        buildContextualRequest(context),
+        buildLabeledEntitiesRequestIfPresentInResultSet(selectionSet),
+        buildLabelApplicationRulesRequestIfPresentInResultSet(selectionSet),
+        (contextualRequest, labeledEntitiesRequests, labelApplicationRulesRequest) ->
+            new DefaultEntityAndRuleJoiner(
+                context, contextualRequest, labeledEntitiesRequests, labelApplicationRulesRequest));
+  }
+
+  private Single<ContextualRequest> buildContextualRequest(GraphQlRequestContext context) {
+    return Single.just(this.requestBuilder.build(context));
+  }
+
+  private Single<List<LabeledEntitiesRequest>> buildLabeledEntitiesRequestIfPresentInResultSet(
+      DataFetchingFieldSelectionSet selectionSet) {
+    return Single.just(
+        this.selectionFinder
+            .findSelections(selectionSet, RESULT_SET_LABELED_ENTITIES_SELECTION_QUERY)
+            .map(
+                selectedField ->
+                    new DefaultLabeledEntitiesRequest(
+                        getEntityType(selectedField),
+                        getLimit(selectedField),
+                        selectedField.getSelectionSet()))
+            .collect(Collectors.toList()));
+  }
+
+  private Single<Optional<LabelApplicationRulesRequest>>
+      buildLabelApplicationRulesRequestIfPresentInResultSet(
+          DataFetchingFieldSelectionSet selectionSet) {
+    return Single.just(
+        this.selectionFinder
+            .findSelections(selectionSet, RESULT_SET_LABEL_APPLICATION_RULES_SELECTION_QUERY)
+            .findAny()
+            .map(
+                selectedField -> new DefaultLabelApplicationRulesRequest(getLimit(selectedField))));
+  }
+
+  private String getEntityType(SelectedField selectedField) {
+    return this.argumentDeserializer
+        .deserializePrimitive(selectedField.getArguments(), EntityTypeStringArgument.class)
+        .orElseThrow();
+  }
+
+  private int getLimit(SelectedField selectedField) {
+    return this.argumentDeserializer
+        .deserializePrimitive(selectedField.getArguments(), LimitArgument.class)
+        .orElse(DEFAULT_LIMIT);
+  }
+
+  @AllArgsConstructor
+  private class DefaultEntityAndRuleJoiner implements EntityAndRuleJoiner {
+    private final GraphQlRequestContext context;
+    private final ContextualRequest contextualRequest;
+    private final List<LabeledEntitiesRequest> labeledEntitiesRequests;
+    private final Optional<LabelApplicationRulesRequest> labelApplicationRulesRequest;
+
+    @Override
+    public Single<LabelResultSet> joinLabelsWithEntitiesAndRules() {
+      return zip(
+              labelDao.getLabels(contextualRequest),
+              getLabelApplicationRuleResultSet(),
+              this::getUpdatedLabelResultSet)
+          .flatMap(labelResultSet -> labelResultSet);
+    }
+
+    private Single<LabelApplicationRuleResultSet> getLabelApplicationRuleResultSet() {
+      return labelApplicationRulesRequest
+          .map(request -> labelApplicationRuleDao.getLabelApplicationRules(contextualRequest))
+          .orElse(Single.just(EMPTY_LABEL_APPLICATION_RULE_RESULT_SET));
+    }
+
+    private Single<LabelResultSet> getUpdatedLabelResultSet(
+        LabelResultSet labelResultSet,
+        LabelApplicationRuleResultSet labelApplicationRuleResultSet) {
+      return getUpdatedLabels(labelResultSet, labelApplicationRuleResultSet)
+          .flatMap(responseConverter::convert);
+    }
+
+    private Single<List<Label>> getUpdatedLabels(
+        LabelResultSet labelResultSet,
+        LabelApplicationRuleResultSet labelApplicationRuleResultSet) {
+      return Observable.fromIterable(labelResultSet.results())
+          .flatMapSingle(label -> this.getUpdatedLabel(label, labelApplicationRuleResultSet))
+          .collect(Collectors.toList());
+    }
+
+    private Single<Label> getUpdatedLabel(
+        Label label, LabelApplicationRuleResultSet labelApplicationRuleResultSet) {
+      return zip(
+              getEntityResultSetMap(label),
+              filterLabelApplicationRuleResultSet(label, labelApplicationRuleResultSet),
+              (entityResultMap, filteredLabelApplicationRuleResultSet) ->
+                  responseConverter.convertLabel(
+                      label, entityResultMap, filteredLabelApplicationRuleResultSet))
+          .flatMap(updatedLabel -> updatedLabel);
+    }
+
+    private Single<LabelApplicationRuleResultSet> filterLabelApplicationRuleResultSet(
+        Label label, LabelApplicationRuleResultSet resultSet) {
+      if (resultSet.equals(EMPTY_LABEL_APPLICATION_RULE_RESULT_SET)) {
+        return Single.just(resultSet);
+      }
+      String labelId = label.id();
+      List<LabelApplicationRule> labelApplicationRules =
+          resultSet.results().stream()
+              .filter(rule -> this.hasMatchingLabelId(rule, labelId))
+              .collect(Collectors.toList());
+      int limit = labelApplicationRulesRequest.get().limit();
+      List<LabelApplicationRule> labelApplicationRuleSubList =
+          labelApplicationRules.size() > limit
+              ? labelApplicationRules.subList(ZERO_OFFSET, limit)
+              : labelApplicationRules;
+      return Single.just(
+          new DefaultLabelApplicationRuleResultSet(
+              labelApplicationRuleSubList,
+              labelApplicationRules.size(),
+              labelApplicationRuleSubList.size()));
+    }
+
+    private boolean hasMatchingLabelId(LabelApplicationRule rule, String labelId) {
+      if (rule.labelApplicationRuleData().action().staticLabels() != null) {
+        return rule.labelApplicationRuleData().action().staticLabels().ids().contains(labelId);
+      }
+      return false;
+    }
+
+    private Single<Map<String, EntityResultSet>> getEntityResultSetMap(Label label) {
+      String labelId = label.id();
+      return Observable.fromIterable(labeledEntitiesRequests)
+          .flatMapSingle(request -> getEntityResultSetMapEntry(labelId, request))
+          .collect(CollectorUtils.immutableMapEntryCollector());
+    }
+
+    private Single<Map.Entry<String, EntityResultSet>> getEntityResultSetMapEntry(
+        String labelId, LabeledEntitiesRequest request) {
+      return buildEntityRequestForType(
+              context, request.entityType(), labelId, request.limit(), request.selectionSet())
+          .flatMap(entityDao::getEntities)
+          .map(resultSet -> Map.entry(request.entityType(), resultSet));
+    }
+
+    private Single<EntityRequest> buildEntityRequestForType(
+        GraphQlRequestContext context,
+        String entityType,
+        String labelId,
+        int limit,
+        DataFetchingFieldSelectionSet selectionSet) {
+      return buildLabelIdFilter(context, entityType, labelId)
+          .flatMap(
+              filterArguments ->
+                  buildEntityRequest(context, entityType, limit, filterArguments, selectionSet));
+    }
+
+    private Single<EntityRequest> buildEntityRequest(
+        GraphQlRequestContext context,
+        String entityType,
+        int entityIdsToFilterSize,
+        List<AttributeAssociation<FilterArgument>> filterArguments,
+        DataFetchingFieldSelectionSet selectionSet) {
+      boolean fetchTotal =
+          selectionFinder
+                  .findSelections(
+                      selectionSet, SelectionQuery.namedChild(ResultSet.RESULT_SET_TOTAL_NAME))
+                  .count()
+              > 0;
+      return resultSetRequestBuilder
+          .build(
+              context,
+              entityType,
+              entityIdsToFilterSize,
+              ZERO_OFFSET,
+              new InstantTimeRange(),
+              List.<AttributeAssociation<AggregatableOrderArgument>>of(), // Order does not matter
+              filterArguments,
+              selectionSet.getFields().stream(),
+              Optional.empty())
+          .map(request -> new DefaultEntityRequest(entityType, request, fetchTotal));
+    }
+
+    private Single<List<AttributeAssociation<FilterArgument>>> buildLabelIdFilter(
+        GraphQlRequestContext context, String entityScope, String labelId) {
+      return filterRequestBuilder.build(context, entityScope, Set.of(new LabelIdFilter(labelId)));
+    }
+  }
+
+  @Value
+  @Accessors(fluent = true)
+  private static class DefaultEntityRequest implements EntityRequest {
+    String entityType;
+    ResultSetRequest<AggregatableOrderArgument> resultSetRequest;
+    List<MetricRequest> metricRequests = Collections.emptyList(); // Only support attributes for now
+    boolean fetchTotal;
+    EdgeSetGroupRequest incomingEdgeRequests = new EmptyEdgeSetGroupRequest();
+    EdgeSetGroupRequest outgoingEdgeRequests = new EmptyEdgeSetGroupRequest();
+    boolean includeInactive = true; // When joining we want the entity regardless of time range
+    Optional<EntityLabelRequest> labelRequest = Optional.empty();
+  }
+
+  @Value
+  @Accessors(fluent = true)
+  private static class EmptyEdgeSetGroupRequest implements EdgeSetGroupRequest {
+    Set<String> entityTypes = Collections.emptySet();
+    Collection<AttributeRequest> attributeRequests = Collections.emptyList();
+    Collection<MetricAggregationRequest> metricAggregationRequests = Collections.emptyList();
+    AttributeRequest neighborIdAttribute = null;
+    AttributeRequest neighborTypeAttribute = null;
+
+    @Override
+    public Single<EntityRequest> buildNeighborRequest(
+        String entityType, Collection<String> neighborIds) {
+      return Single.error(
+          new UnsupportedOperationException(
+              "Does not support fetching neighbors for joined entities"));
+    }
+  }
+
+  @Value
+  @Accessors(fluent = true)
+  private static class InstantTimeRange implements TimeRangeArgument {
+    // Bug in gateway requires this time range be larger, even though we don't care about it
+    Instant startTime = Instant.now().minus(15, ChronoUnit.MINUTES);
+    Instant endTime = Instant.now();
+  }
+
+  @Value
+  @Accessors(fluent = true)
+  private static class LabelIdFilter implements FilterArgument {
+    FilterType type = FilterType.ATTRIBUTE;
+    String key = "labels";
+    FilterOperatorType operator = FilterOperatorType.EQUALS;
+    String value;
+    AttributeScope idType = null;
+    String idScope = null;
+  }
+
+  @Value
+  @Accessors(fluent = true)
+  private static class DefaultLabeledEntitiesRequest implements LabeledEntitiesRequest {
+    String entityType;
+    int limit;
+    DataFetchingFieldSelectionSet selectionSet;
+  }
+
+  @Value
+  @Accessors(fluent = true)
+  private static class DefaultLabelApplicationRulesRequest implements LabelApplicationRulesRequest {
+    int limit;
+  }
+
+  @Value
+  @Accessors(fluent = true)
+  private static class DefaultLabelApplicationRuleResultSet
+      implements LabelApplicationRuleResultSet {
+    List<LabelApplicationRule> results;
+    long total;
+    long count;
+  }
+}
diff --git a/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultLabelJoinerBuilder.java b/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultLabelJoinerBuilder.java
index 43a7dc27..9c8ba845 100644
--- a/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultLabelJoinerBuilder.java
+++ b/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultLabelJoinerBuilder.java
@@ -1,181 +1,58 @@
 package org.hypertrace.graphql.label.joiner;
 
-import static io.reactivex.rxjava3.core.Single.zip;
 import static java.util.function.Function.identity;
-import static org.hypertrace.graphql.label.schema.LabeledEntities.ENTITY_TYPE_ARGUMENT_NAME;
-import static org.hypertrace.graphql.label.schema.LabeledEntities.LABELED_ENTITIES_QUERY_NAME;
 
-import graphql.schema.DataFetchingFieldSelectionSet;
 import io.reactivex.rxjava3.core.Observable;
 import io.reactivex.rxjava3.core.Single;
-import java.time.Instant;
-import java.time.temporal.ChronoUnit;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Optional;
-import java.util.Set;
 import java.util.stream.Collectors;
 import javax.inject.Inject;
 import lombok.AllArgsConstructor;
-import lombok.Value;
-import lombok.experimental.Accessors;
 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.ContextualRequest;
 import org.hypertrace.core.graphql.common.request.ContextualRequestBuilder;
-import org.hypertrace.core.graphql.common.request.FilterRequestBuilder;
-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;
-import org.hypertrace.core.graphql.common.schema.attributes.AttributeScope;
-import org.hypertrace.core.graphql.common.schema.results.ResultSet;
-import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterArgument;
-import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterOperatorType;
-import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterType;
-import org.hypertrace.core.graphql.common.schema.results.arguments.page.LimitArgument;
 import org.hypertrace.core.graphql.common.utils.CollectorUtils;
 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.dao.EntityDao;
-import org.hypertrace.graphql.entity.request.EdgeSetGroupRequest;
-import org.hypertrace.graphql.entity.request.EntityLabelRequest;
-import org.hypertrace.graphql.entity.request.EntityRequest;
-import org.hypertrace.graphql.entity.schema.EntityResultSet;
 import org.hypertrace.graphql.label.dao.LabelDao;
 import org.hypertrace.graphql.label.dao.LabelResponseConverter;
 import org.hypertrace.graphql.label.schema.Label;
 import org.hypertrace.graphql.label.schema.LabelResultSet;
-import org.hypertrace.graphql.metric.request.MetricAggregationRequest;
-import org.hypertrace.graphql.metric.request.MetricRequest;
-import org.hypertrace.graphql.metric.schema.argument.AggregatableOrderArgument;
 
 @Slf4j
 class DefaultLabelJoinerBuilder implements LabelJoinerBuilder {
 
-  private static final Integer DEFAULT_LIMIT = 10;
-  private static final int ZERO_OFFSET = 0;
-
-  private static final SelectionQuery RESULT_SET_LABELED_ENTITIES_SELECTION_QUERY =
-      SelectionQuery.builder()
-          .selectionPath(List.of(ResultSet.RESULT_SET_RESULTS_NAME, LABELED_ENTITIES_QUERY_NAME))
-          .build();
-
   private final ContextualRequestBuilder requestBuilder;
-  private final ResultSetRequestBuilder resultSetRequestBuilder;
-  private final FilterRequestBuilder filterRequestBuilder;
   private final LabelDao labelDao;
   private final LabelResponseConverter responseConverter;
-  private final GraphQlSelectionFinder selectionFinder;
-  private final EntityDao entityDao;
 
   @Inject
   DefaultLabelJoinerBuilder(
       ContextualRequestBuilder requestBuilder,
-      ResultSetRequestBuilder resultSetRequestBuilder,
-      FilterRequestBuilder filterRequestBuilder,
       LabelDao labelDao,
-      EntityDao entityDao,
-      LabelResponseConverter responseConverter,
-      GraphQlSelectionFinder selectionFinder) {
+      LabelResponseConverter responseConverter) {
     this.requestBuilder = requestBuilder;
-    this.resultSetRequestBuilder = resultSetRequestBuilder;
-    this.filterRequestBuilder = filterRequestBuilder;
     this.labelDao = labelDao;
-    this.entityDao = entityDao;
     this.responseConverter = responseConverter;
-    this.selectionFinder = selectionFinder;
   }
 
   @Override
   public Single<LabelJoiner> build(GraphQlRequestContext context) {
-    return buildContextualRequest(context)
-        .map(request -> new DefaultLabelJoiner(context, request, null));
-  }
-
-  @Override
-  public Single<LabelJoiner> build(
-      GraphQlRequestContext context, DataFetchingFieldSelectionSet selectionSet) {
-    return zip(
-        buildContextualRequest(context),
-        buildLabeledEntitiesRequestIfPresentInResultSet(selectionSet),
-        (contextualRequest, labeledEntitiesRequests) ->
-            new DefaultLabelJoiner(context, contextualRequest, labeledEntitiesRequests));
-  }
-
-  private Single<ContextualRequest> buildContextualRequest(GraphQlRequestContext context) {
-    return Single.just(this.requestBuilder.build(context));
-  }
-
-  private Single<List<LabeledEntitiesRequest>> buildLabeledEntitiesRequestIfPresentInResultSet(
-      DataFetchingFieldSelectionSet selectionSet) {
-    return Single.just(
-        this.selectionFinder
-            .findSelections(selectionSet, RESULT_SET_LABELED_ENTITIES_SELECTION_QUERY)
-            .map(
-                selectedField ->
-                    new DefaultLabeledEntitiesRequest(
-                        (String) selectedField.getArguments().get(ENTITY_TYPE_ARGUMENT_NAME),
-                        (Integer)
-                            selectedField
-                                .getArguments()
-                                .getOrDefault(LimitArgument.ARGUMENT_NAME, DEFAULT_LIMIT),
-                        selectedField.getSelectionSet()))
-            .collect(Collectors.toList()));
+    return Single.just(new DefaultLabelJoiner(context));
   }
 
   @AllArgsConstructor
   private class DefaultLabelJoiner implements LabelJoiner {
     private final GraphQlRequestContext context;
-    private final ContextualRequest resultSetRequest;
-    private final List<LabeledEntitiesRequest> labeledEntitiesRequests;
 
     @Override
     public <T> Single<Map<T, LabelResultSet>> joinLabels(
         Collection<T> joinSources, LabelIdGetter<T> labelIdGetter) {
       return labelDao
-          .getLabels(resultSetRequest)
+          .getLabels(requestBuilder.build(context))
           .flatMap(labelResultSet -> buildLabelMap(labelResultSet.results()))
-          .flatMap(labelsMap -> getLabelResultSetMap(joinSources, labelIdGetter, labelsMap));
-    }
-
-    @Override
-    public Single<LabelResultSet> joinLabelsWithEntities() {
-      return labelDao.getLabels(resultSetRequest).flatMap(this::getUpdatedLabelResultSet);
-    }
-
-    private Single<LabelResultSet> getUpdatedLabelResultSet(LabelResultSet labelResultSet) {
-      return getUpdatedLabels(labelResultSet).flatMap(responseConverter::convert);
-    }
-
-    private Single<List<Label>> getUpdatedLabels(LabelResultSet labelResultSet) {
-      return Observable.fromIterable(labelResultSet.results())
-          .flatMapSingle(this::getUpdatedLabel)
-          .collect(Collectors.toList());
-    }
-
-    private Single<Label> getUpdatedLabel(Label label) {
-      return getEntityResultSetMap(label)
-          .flatMap(entityResultMap -> responseConverter.convertLabel(label, entityResultMap));
-    }
-
-    private Single<Map<String, EntityResultSet>> getEntityResultSetMap(Label label) {
-      String labelId = label.id();
-      return Observable.fromIterable(labeledEntitiesRequests)
-          .flatMapSingle(request -> getEntityResultSetMapEntry(labelId, request))
-          .collect(CollectorUtils.immutableMapEntryCollector());
-    }
-
-    private Single<Map.Entry<String, EntityResultSet>> getEntityResultSetMapEntry(
-        String labelId, LabeledEntitiesRequest request) {
-      return buildEntityRequestForType(
-              context, request.entityType(), labelId, request.limit(), request.selectionSet())
-          .flatMap(entityDao::getEntities)
-          .map(resultSet -> Map.entry(request.entityType(), resultSet));
+          .flatMap(labelMap -> getLabelResultSetMap(joinSources, labelIdGetter, labelMap));
     }
 
     private Single<Map<String, Label>> buildLabelMap(List<Label> results) {
@@ -195,7 +72,7 @@ private <T> Single<Map.Entry<T, LabelResultSet>> buildLabelResultSetMapEntry(
       return labelIdGetter
           .getLabelIds(source)
           .flatMap(labelIds -> filterLabels(labelIds, labelMap))
-          .flatMap(responseConverter::convert)
+          .flatMap(labels -> responseConverter.convert(labels))
           .map(labelResultSet -> Map.entry(source, labelResultSet));
     }
 
@@ -214,106 +91,5 @@ private Label getLabelFromMap(String labelId, Map<String, Label> labelMap) {
       }
       return label;
     }
-
-    private Single<EntityRequest> buildEntityRequestForType(
-        GraphQlRequestContext context,
-        String entityType,
-        String labelId,
-        int limit,
-        DataFetchingFieldSelectionSet selectionSet) {
-      return buildLabelIdFilter(context, entityType, labelId)
-          .flatMap(
-              filterArguments ->
-                  buildEntityRequest(context, entityType, limit, filterArguments, selectionSet));
-    }
-
-    private Single<EntityRequest> buildEntityRequest(
-        GraphQlRequestContext context,
-        String entityType,
-        int entityIdsToFilterSize,
-        List<AttributeAssociation<FilterArgument>> filterArguments,
-        DataFetchingFieldSelectionSet selectionSet) {
-      boolean fetchTotal =
-          selectionFinder
-                  .findSelections(
-                      selectionSet, SelectionQuery.namedChild(ResultSet.RESULT_SET_TOTAL_NAME))
-                  .count()
-              > 0;
-      return resultSetRequestBuilder
-          .build(
-              context,
-              entityType,
-              entityIdsToFilterSize,
-              ZERO_OFFSET,
-              new InstantTimeRange(),
-              List.<AttributeAssociation<AggregatableOrderArgument>>of(), // Order does not matter
-              filterArguments,
-              selectionSet.getFields().stream(),
-              Optional.empty())
-          .map(request -> new DefaultEntityRequest(entityType, request, fetchTotal));
-    }
-
-    private Single<List<AttributeAssociation<FilterArgument>>> buildLabelIdFilter(
-        GraphQlRequestContext context, String entityScope, String labelId) {
-      return filterRequestBuilder.build(context, entityScope, Set.of(new LabelIdFilter(labelId)));
-    }
-  }
-
-  @Value
-  @Accessors(fluent = true)
-  private static class DefaultEntityRequest implements EntityRequest {
-    String entityType;
-    ResultSetRequest<AggregatableOrderArgument> resultSetRequest;
-    List<MetricRequest> metricRequests = Collections.emptyList(); // Only support attributes for now
-    boolean fetchTotal;
-    EdgeSetGroupRequest incomingEdgeRequests = new EmptyEdgeSetGroupRequest();
-    EdgeSetGroupRequest outgoingEdgeRequests = new EmptyEdgeSetGroupRequest();
-    boolean includeInactive = true; // When joining we want the entity regardless of time range
-    Optional<EntityLabelRequest> labelRequest = Optional.empty();
-  }
-
-  @Value
-  @Accessors(fluent = true)
-  private static class EmptyEdgeSetGroupRequest implements EdgeSetGroupRequest {
-    Set<String> entityTypes = Collections.emptySet();
-    Collection<AttributeRequest> attributeRequests = Collections.emptyList();
-    Collection<MetricAggregationRequest> metricAggregationRequests = Collections.emptyList();
-    AttributeRequest neighborIdAttribute = null;
-    AttributeRequest neighborTypeAttribute = null;
-
-    @Override
-    public Single<EntityRequest> buildNeighborRequest(
-        String entityType, Collection<String> neighborIds) {
-      return Single.error(
-          new UnsupportedOperationException(
-              "Does not support fetching neighbors for joined entities"));
-    }
-  }
-
-  @Value
-  @Accessors(fluent = true)
-  private static class InstantTimeRange implements TimeRangeArgument {
-    // Bug in gateway requires this time range be larger, even though we don't care about it
-    Instant startTime = Instant.now().minus(15, ChronoUnit.MINUTES);
-    Instant endTime = Instant.now();
-  }
-
-  @Value
-  @Accessors(fluent = true)
-  private static class LabelIdFilter implements FilterArgument {
-    FilterType type = FilterType.ATTRIBUTE;
-    String key = "labels";
-    FilterOperatorType operator = FilterOperatorType.EQUALS;
-    String value;
-    AttributeScope idType = null;
-    String idScope = null;
-  }
-
-  @Value
-  @Accessors(fluent = true)
-  private static class DefaultLabeledEntitiesRequest implements LabeledEntitiesRequest {
-    String entityType;
-    int limit;
-    DataFetchingFieldSelectionSet selectionSet;
   }
 }
diff --git a/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoinerModule.java b/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoinerModule.java
index 33d7ad70..8680175e 100644
--- a/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoinerModule.java
+++ b/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/LabelJoinerModule.java
@@ -8,6 +8,7 @@ public class LabelJoinerModule extends AbstractModule {
   @Override
   protected void configure() {
     bind(LabelJoinerBuilder.class).to(DefaultLabelJoinerBuilder.class);
+    bind(EntityAndRuleJoinerBuilder.class).to(DefaultEntityAndRuleJoinerBuilder.class);
 
     requireBinding(LabelDao.class);
   }

From 021112aeb2aac24bb04044436b53775efc05e3d3 Mon Sep 17 00:00:00 2001
From: saxenakshitiz <kshitiz@traceable.ai>
Date: Tue, 28 Dec 2021 01:31:15 +0530
Subject: [PATCH 3/3] Address review comments

---
 .../graphql/label/schema/Label.java           | 22 ++++++++++++++++++-
 .../label/schema/LabelApplicationRules.java   | 19 ----------------
 .../graphql/label/schema/LabeledEntities.java | 20 -----------------
 .../DefaultEntityAndRuleJoinerBuilder.java    |  4 ++--
 4 files changed, 23 insertions(+), 42 deletions(-)
 delete mode 100644 hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/LabelApplicationRules.java
 delete mode 100644 hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/LabeledEntities.java

diff --git a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/Label.java b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/Label.java
index fc2728d6..da6996f3 100644
--- a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/Label.java
+++ b/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, LabelApplicationRules {
+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);
 }
diff --git a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/LabelApplicationRules.java b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/LabelApplicationRules.java
deleted file mode 100644
index 44891e94..00000000
--- a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/LabelApplicationRules.java
+++ /dev/null
@@ -1,19 +0,0 @@
-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.results.arguments.page.LimitArgument;
-import org.hypertrace.graphql.label.schema.rule.LabelApplicationRuleResultSet;
-
-@GraphQLName(LabelApplicationRules.TYPE_NAME)
-public interface LabelApplicationRules {
-  String TYPE_NAME = "LabelApplicationRules";
-  String LABEL_APPLICATION_RULES_QUERY_NAME = "labelApplicationRules";
-
-  @GraphQLField
-  @GraphQLNonNull
-  @GraphQLName(LABEL_APPLICATION_RULES_QUERY_NAME)
-  LabelApplicationRuleResultSet labelApplicationRules(
-      @GraphQLName(LimitArgument.ARGUMENT_NAME) int limit);
-}
diff --git a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/LabeledEntities.java b/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/LabeledEntities.java
deleted file mode 100644
index 8d3e78c3..00000000
--- a/hypertrace-graphql-labels-schema-api/src/main/java/org/hypertrace/graphql/label/schema/LabeledEntities.java
+++ /dev/null
@@ -1,20 +0,0 @@
-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.results.arguments.page.LimitArgument;
-
-@GraphQLName(LabeledEntities.TYPE_NAME)
-public interface LabeledEntities {
-  String TYPE_NAME = "LabeledEntities";
-  String LABELED_ENTITIES_QUERY_NAME = "labeledEntities";
-  String ENTITY_TYPE_ARGUMENT_NAME = "type";
-
-  @GraphQLField
-  @GraphQLNonNull
-  @GraphQLName(LABELED_ENTITIES_QUERY_NAME)
-  LabeledEntityResultSet labeledEntities(
-      @GraphQLNonNull @GraphQLName(ENTITY_TYPE_ARGUMENT_NAME) String entityType,
-      @GraphQLName(LimitArgument.ARGUMENT_NAME) int limit);
-}
diff --git a/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultEntityAndRuleJoinerBuilder.java b/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultEntityAndRuleJoinerBuilder.java
index 80381128..a569781c 100644
--- a/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultEntityAndRuleJoinerBuilder.java
+++ b/hypertrace-graphql-labels-schema-impl/src/main/java/org/hypertrace/graphql/label/joiner/DefaultEntityAndRuleJoinerBuilder.java
@@ -1,8 +1,8 @@
 package org.hypertrace.graphql.label.joiner;
 
 import static io.reactivex.rxjava3.core.Single.zip;
-import static org.hypertrace.graphql.label.schema.LabelApplicationRules.LABEL_APPLICATION_RULES_QUERY_NAME;
-import static org.hypertrace.graphql.label.schema.LabeledEntities.LABELED_ENTITIES_QUERY_NAME;
+import static org.hypertrace.graphql.label.schema.Label.LABELED_ENTITIES_QUERY_NAME;
+import static org.hypertrace.graphql.label.schema.Label.LABEL_APPLICATION_RULES_QUERY_NAME;
 
 import graphql.schema.DataFetchingFieldSelectionSet;
 import graphql.schema.SelectedField;