fix: add in filter and use action param for countAggregate (#118)

* fix: add in filter and use action param for countAggregate

* use finder for countAggregate and deprecate action

* use more generic metadata

dao-api/src/main/pegasus/com/linkedin/metadata/query/MapMetadata.pdl

@@ -0,0 +1,11 @@
+namespace com.linkedin.metadata.query
+
+/**
+ * The model for map metadata
+ */
+record MapMetadata {
+  /**
+   * Map for string keys to long values
+   */
+  longMap: map[string, long]
+}

dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java

@@ -1214,6 +1214,7 @@ private static String constructCountAggregateSQLQuery(@Nonnull IndexCriterionArr
       }
       whereClause.append(" t").append(i).append(".aspect = ?");
       if (criterion.getPathParams() != null) {
+        validateConditionAndValue(criterion);
         whereClause.append(" AND t")
             .append(i)
             .append(".path = ? AND t")
@@ -1222,7 +1223,7 @@ private static String constructCountAggregateSQLQuery(@Nonnull IndexCriterionArr
             .append(getGMAIndexPair(criterion).valueType)
             .append(" ")
             .append(getStringForOperator(criterion.getPathParams().getCondition()))
-            .append("?");
+            .append(getPlaceholderStringForValue(criterion.getPathParams().getValue()));
       }
     });
     whereClause.append(" AND tgroup.aspect = ? AND tgroup.path = ? ");

dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanLocalDAOTest.java

@@ -2183,6 +2183,22 @@ public void testCountAggregate() {
     result = dao.countAggregate(indexFilter3, indexGroupByCriterion1);
     assertEquals(result.size(), 1);
     assertEquals(result.get("valB").longValue(), 2);
+
+    // in filter
+    IndexValue indexValue5 = new IndexValue();
+    indexValue5.setArray(new StringArray("valA", "valB"));
+    IndexCriterion criterion5 = new IndexCriterion().setAspect(aspect1)
+        .setPathParams(new IndexPathParams().setPath("/value").setValue(indexValue5).setCondition(Condition.IN));
+
+    IndexCriterionArray indexCriterionArray4 = new IndexCriterionArray(Collections.singletonList(criterion5));
+    final IndexFilter indexFilter4 = new IndexFilter().setCriteria(indexCriterionArray4);
+
+    result = dao.countAggregate(indexFilter4, indexGroupByCriterion1);
+    List<FooUrn> test = dao.listUrns(indexFilter4, null, null, 10);
+    assertEquals(test.size(), 3);
+    assertEquals(result.size(), 2);
+    assertEquals(result.get("valB").longValue(), 2);
+    assertEquals(result.get("valA").longValue(), 1);
   }
 
   @Test

restli-resources/src/main/java/com/linkedin/metadata/restli/BaseEntityResource.java

@@ -2,6 +2,7 @@
 
 import com.linkedin.common.AuditStamp;
 import com.linkedin.common.urn.Urn;
+import com.linkedin.data.template.LongMap;
 import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.data.template.StringArray;
 import com.linkedin.data.template.UnionTemplate;
@@ -17,7 +18,9 @@
 import com.linkedin.metadata.query.IndexGroupByCriterion;
 import com.linkedin.metadata.query.IndexSortCriterion;
 import com.linkedin.metadata.query.ListResultMetadata;
+import com.linkedin.metadata.query.MapMetadata;
 import com.linkedin.parseq.Task;
+import com.linkedin.restli.common.EmptyRecord;
 import com.linkedin.restli.server.CollectionResult;
 import com.linkedin.restli.server.PagingContext;
 import com.linkedin.restli.server.annotations.Action;
@@ -551,17 +554,42 @@ public Task<ListResult<VALUE>> filter(
     });
   }
 
-  /*
+  /**
+   * Gets a collection result with count aggregate metadata, which has the count of an aggregation
+   * specified by the aspect and field to group by.
+   *
+   * @param indexFilter {@link IndexFilter} that defines the filter conditions
+   * @param indexGroupByCriterion {@link IndexGroupByCriterion} that defines the aspect to group by
+   * @return {@link CollectionResult} containing metadata that has a map of the field values to their count
+   */
+  @Finder(FINDER_COUNT_AGGREGATE)
+  @Nonnull
+  public Task<CollectionResult<EmptyRecord, MapMetadata>> countAggregateFilter(
+      @QueryParam(PARAM_FILTER) @Optional @Nullable IndexFilter indexFilter,
+      @QueryParam(PARAM_GROUP) IndexGroupByCriterion indexGroupByCriterion
+  ) {
+    final IndexFilter filter = indexFilter == null ? getDefaultIndexFilter() : indexFilter;
+
+    return RestliUtils.toTask(() -> {
+      Map<String, Long> countAggregateMap = getLocalDAO().countAggregate(indexFilter, indexGroupByCriterion);
+      MapMetadata mapMetadata = new MapMetadata().setLongMap(new LongMap(countAggregateMap));
+      return new CollectionResult<EmptyRecord, MapMetadata>(new ArrayList<>(), mapMetadata);
+    });
+  }
+
+  /**
    * Gets the count of an aggregation specified by the aspect and field to group by.
    * @param indexFilter {@link IndexFilter} that defines the filter conditions
    * @param indexGroupByCriterion {@link IndexGroupByCriterion} that defines the aspect to group by
-   * @return map of the field to the count
+   * @return map of the field values to their count
+   *
+   * @deprecated Use {@link #countAggregateFilter(IndexFilter, IndexGroupByCriterion)} instead
    */
   @Action(name = ACTION_COUNT_AGGREGATE)
   @Nonnull
   public Task<Map<String, Long>> countAggregate(
-      @QueryParam(PARAM_FILTER) @Optional @Nullable IndexFilter indexFilter,
-      @QueryParam(PARAM_GROUP) IndexGroupByCriterion indexGroupByCriterion
+      @ActionParam(PARAM_FILTER) @Optional @Nullable IndexFilter indexFilter,
+      @ActionParam(PARAM_GROUP) IndexGroupByCriterion indexGroupByCriterion
   ) {
     final IndexFilter filter = indexFilter == null ? getDefaultIndexFilter() : indexFilter;
 

restli-resources/src/main/java/com/linkedin/metadata/restli/RestliConstants.java

@@ -5,6 +5,7 @@ private RestliConstants() { }
 
   public static final String FINDER_SEARCH = "search";
   public static final String FINDER_FILTER = "filter";
+  public static final String FINDER_COUNT_AGGREGATE = "countAggregate";
 
   public static final String ACTION_AUTOCOMPLETE = "autocomplete";
   public static final String ACTION_BACKFILL = "backfill";

restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntityResourceTest.java

@@ -2,6 +2,7 @@
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.linkedin.data.template.LongMap;
 import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.metadata.backfill.BackfillMode;
 import com.linkedin.metadata.dao.AspectKey;
@@ -15,11 +16,13 @@
 import com.linkedin.metadata.query.IndexFilter;
 import com.linkedin.metadata.query.IndexGroupByCriterion;
 import com.linkedin.metadata.query.IndexSortCriterion;
+import com.linkedin.metadata.query.MapMetadata;
 import com.linkedin.metadata.query.SortOrder;
 import com.linkedin.parseq.BaseEngineTest;
 import com.linkedin.restli.common.ComplexResourceKey;
 import com.linkedin.restli.common.EmptyRecord;
 import com.linkedin.restli.common.HttpStatus;
+import com.linkedin.restli.server.CollectionResult;
 import com.linkedin.restli.server.PagingContext;
 import com.linkedin.restli.server.ResourceContext;
 import com.linkedin.restli.server.RestLiServiceException;
@@ -646,4 +649,32 @@ public void testCountAggregate() {
 
     assertEquals(actual, mapResult);
   }
+
+  @Test
+  public void testCountAggregateFilter() {
+    FooUrn urn1 = makeFooUrn(1);
+    FooUrn urn2 = makeFooUrn(2);
+    AspectFoo foo1 = new AspectFoo().setValue("val1");
+    AspectFoo foo2 = new AspectFoo().setValue("val2");
+    AspectBar bar1 = new AspectBar().setValue("val1");
+    AspectBar bar2 = new AspectBar().setValue("val2");
+
+    UrnAspectEntry<FooUrn> entry1 = new UrnAspectEntry<>(urn1, Arrays.asList(foo1, bar1));
+    UrnAspectEntry<FooUrn> entry2 = new UrnAspectEntry<>(urn2, Arrays.asList(foo2, bar2));
+
+    IndexCriterion criterion = new IndexCriterion().setAspect(AspectFoo.class.getCanonicalName());
+    IndexCriterionArray criterionArray = new IndexCriterionArray(criterion);
+    IndexFilter indexFilter = new IndexFilter().setCriteria(criterionArray);
+    IndexGroupByCriterion indexGroupByCriterion = new IndexGroupByCriterion().setAspect(AspectFoo.class.getCanonicalName())
+        .setPath("/value");
+    Map<String, Long> mapResult = new HashMap<>();
+    mapResult.put("val1", 1L);
+    mapResult.put("val2", 1L);
+
+    when(_mockLocalDAO.countAggregate(indexFilter, indexGroupByCriterion)).thenReturn(mapResult);
+    CollectionResult<EmptyRecord, MapMetadata> actual =
+        runAndWait(_resource.countAggregateFilter(indexFilter, indexGroupByCriterion));
+
+    assertEquals(actual.getMetadata().getLongMap(), new LongMap(mapResult));
+  }
 }