push limit and offset to EQS

gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/entity/EntityServiceEntityFetcher.java

@@ -23,8 +23,11 @@
 import org.hypertrace.gateway.service.explore.ExploreRequestContext;
 import org.hypertrace.gateway.service.v1.common.OrderByExpression;
 import org.hypertrace.gateway.service.v1.explore.ExploreRequest;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class EntityServiceEntityFetcher {
+  private static final Logger LOG = LoggerFactory.getLogger(EntityServiceEntityFetcher.class);
   private static final int DEFAULT_ENTITY_REQUEST_LIMIT = 10_000;
 
   private final AttributeMetadataProvider attributeMetadataProvider;
@@ -60,16 +63,19 @@ private EntityQueryRequest buildRequest(
 
     addGroupBys(exploreRequest, builder);
     addSelections(requestContext, exploreRequest, builder);
-
-    // Ideally, needs the limit and offset for group by, since the fetcher is only triggered when
-    // there is a group by, or a single aggregation selection. A single aggregated selection would
-    // always return a single result (i.e. limit 1)
-    builder.setLimit(DEFAULT_ENTITY_REQUEST_LIMIT);
-
-    // TODO: Push group by down to EQS
-    // If there is a group by, specify a large limit and track actual limit, offset and order by
-    // expression list, so we can compute these once the we get the results.
-    if (requestContext.hasGroupBy()) {
+    addLimitAndOffset(requestContext, exploreRequest, builder);
+
+    // TODO: Push order by down to EQS
+    // EQS (and document-store) currently doesn't support order by on functional expressions
+    // If there are order by expressions, specify a large limit and track actual limit, offset and
+    // order by
+    // expression list, so we can compute these once we get the results.
+    if (!requestContext.getOrderByExpressions().isEmpty()) {
+      // Ideally, needs the limit and offset for group by, since the fetcher is only triggered when
+      // there is a group by, or a single aggregation selection. A single aggregated selection would
+      // always return a single result (i.e. limit 1)
+      builder.setOffset(0);
+      builder.setLimit(DEFAULT_ENTITY_REQUEST_LIMIT);
       // Will need to do the ordering, limit and offset ourselves after we get the group by results
       requestContext.setOrderByExpressions(getRequestOrderByExpressions(exploreRequest));
     }
@@ -107,6 +113,33 @@ private void addSelections(
         });
   }
 
+  private void addLimitAndOffset(
+      ExploreRequestContext requestContext,
+      ExploreRequest exploreRequest,
+      EntityQueryRequest.Builder builder) {
+    // handle group by scenario with group limit set
+    if (requestContext.hasGroupBy()) {
+      int limit = exploreRequest.getLimit();
+      if (exploreRequest.getGroupLimit() > 0) {
+        // in group by scenario, set limit to minimum of limit or group-limit
+        limit = Math.min(exploreRequest.getLimit(), exploreRequest.getGroupLimit());
+      }
+      // don't exceed default group by limit
+      if (limit > DEFAULT_ENTITY_REQUEST_LIMIT) {
+        LOG.error(
+            "Trying to query for rows more than the default limit {} : {}",
+            DEFAULT_ENTITY_REQUEST_LIMIT,
+            exploreRequest);
+        throw new UnsupportedOperationException(
+            "Trying to query for rows more than the default limit " + exploreRequest);
+      }
+      builder.setLimit(limit);
+    } else {
+      builder.setLimit(exploreRequest.getLimit());
+      builder.setOffset(exploreRequest.getOffset());
+    }
+  }
+
   private Filter.Builder buildFilter(
       ExploreRequest exploreRequest, List<String> entityIdAttributeIds, Set<String> entityIds) {
     Builder filterBuilder =