diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/entity/EntityServiceEntityFetcher.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/entity/EntityServiceEntityFetcher.java index 64ca8fe3..f169763b 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/entity/EntityServiceEntityFetcher.java +++ b/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 entityIdAttributeIds, Set entityIds) { Builder filterBuilder =