refactor(explore): expression context in handlers

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

@@ -9,25 +9,24 @@
 import org.hypertrace.gateway.service.v1.explore.ExploreRequest;
 
 public class ExploreRequestContext extends QueryRequestContext {
-  private final ExploreRequest exploreRequest;
+  private final ExploreRequest request;
 
   private boolean hasGroupBy = false;
   private List<OrderByExpression> orderByExpressions;
 
-  public ExploreRequestContext(RequestContext grpcRequestContext, ExploreRequest exploreRequest) {
-    super(
-        grpcRequestContext, exploreRequest.getStartTimeMillis(), exploreRequest.getEndTimeMillis());
+  public ExploreRequestContext(RequestContext grpcRequestContext, ExploreRequest request) {
+    super(grpcRequestContext, request.getStartTimeMillis(), request.getEndTimeMillis());
 
-    this.exploreRequest = exploreRequest;
-    this.orderByExpressions = exploreRequest.getOrderByList();
+    this.request = request;
+    this.orderByExpressions = request.getOrderByList();
   }
 
-  public ExploreRequest getExploreRequest() {
-    return this.exploreRequest;
+  public ExploreRequest getRequest() {
+    return this.request;
   }
 
   public String getContext() {
-    return this.exploreRequest.getContext();
+    return this.request.getContext();
   }
 
   /**
@@ -77,19 +76,19 @@ public int getRowLimitAfterRest() {
   }
 
   public int getOffset() {
-    return this.exploreRequest.getOffset();
+    return this.request.getOffset();
   }
 
   public boolean getIncludeRestGroup() {
-    return this.exploreRequest.getIncludeRestGroup();
+    return this.request.getIncludeRestGroup();
   }
 
   private int getGroupByLimit() {
     // If a request has no group limit, default to row limit
     if (this.providedGroupLimitUnset()) {
-      return this.exploreRequest.getLimit();
+      return this.request.getLimit();
     }
-    return this.exploreRequest.getGroupLimit();
+    return this.request.getGroupLimit();
   }
 
   /**
@@ -100,21 +99,21 @@ private int getBackwardsCompatibleRowLimit() {
     // Row limit previously was the number of groups for a grouped + time series request. If group
     // limit isn't set, assume the old format and expand the row limit to default limit
     if (this.providedGroupLimitUnset() && this.hasGroupBy() && this.hasTimeAggregations()) {
-      return Math.max(this.exploreRequest.getLimit(), DEFAULT_QUERY_SERVICE_GROUP_BY_LIMIT);
+      return Math.max(this.request.getLimit(), DEFAULT_QUERY_SERVICE_GROUP_BY_LIMIT);
     }
     // Previously, a rest group would not count against the limit
     if (this.providedGroupLimitUnset() && this.hasGroupBy() && this.getIncludeRestGroup()) {
-      return this.exploreRequest.getLimit() + 1;
+      return this.request.getLimit() + 1;
     }
 
-    return this.exploreRequest.getLimit();
+    return this.request.getLimit();
   }
 
   private boolean providedGroupLimitUnset() {
-    return this.exploreRequest.getGroupLimit() == 0;
+    return this.request.getGroupLimit() == 0;
   }
 
   private boolean hasTimeAggregations() {
-    return !this.exploreRequest.getTimeAggregationList().isEmpty();
+    return !this.request.getTimeAggregationList().isEmpty();
   }
 }

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

@@ -72,9 +72,6 @@ public ExploreResponse explore(RequestContext requestContext, ExploreRequest req
 
     final Instant start = Instant.now();
     try {
-      ExploreRequestContext exploreRequestContext =
-          new ExploreRequestContext(requestContext.getGrpcContext(), request);
-
       // Add extra filters based on the scope.
       request =
           ExploreRequest.newBuilder(request)
@@ -83,21 +80,28 @@ public ExploreResponse explore(RequestContext requestContext, ExploreRequest req
                       request.getContext(),
                       request.getFilter(),
                       attributeMetadataProvider,
-                      exploreRequestContext))
+                      requestContext))
               .build();
-      ExploreRequestContext newExploreRequestContext =
-          new ExploreRequestContext(requestContext.getGrpcContext(), request);
-
       Map<String, AttributeMetadata> attributeMetadataMap =
-          attributeMetadataProvider.getAttributesMetadata(
-              newExploreRequestContext, request.getContext());
+          attributeMetadataProvider.getAttributesMetadata(requestContext, request.getContext());
+
       exploreRequestValidator.validate(request, attributeMetadataMap);
 
-      IRequestHandler requestHandler = getRequestHandler(request, attributeMetadataMap);
+      ExploreRequestContext exploreRequestContext =
+          new ExploreRequestContext(requestContext.getGrpcContext(), request);
+      ExpressionContext expressionContext =
+          new ExpressionContext(
+              attributeMetadataMap,
+              request.getFilter(),
+              request.getSelectionList(),
+              request.getTimeAggregationList(),
+              request.getOrderByList(),
+              request.getGroupByList());
+
+      IRequestHandler requestHandler = getRequestHandler(request, expressionContext);
 
       ExploreResponse.Builder responseBuilder =
-          requestHandler.handleRequest(newExploreRequestContext, request);
-
+          requestHandler.handleRequest(exploreRequestContext, expressionContext);
       return responseBuilder.build();
     } finally {
       queryExecutionTimer.record(
@@ -111,18 +115,10 @@ private boolean isContextAnEntityType(ExploreRequest request) {
   }
 
   private IRequestHandler getRequestHandler(
-      ExploreRequest request, Map<String, AttributeMetadata> attributeMetadataMap) {
+      ExploreRequest request, ExpressionContext expressionContext) {
     if (isContextAnEntityType(request)
         && !hasTimeAggregations(request)
         && !request.getGroupByList().isEmpty()) {
-      ExpressionContext expressionContext =
-          new ExpressionContext(
-              attributeMetadataMap,
-              request.getFilter(),
-              request.getSelectionList(),
-              request.getTimeAggregationList(),
-              request.getOrderByList(),
-              request.getGroupByList());
       Optional<String> source =
           ExpressionContext.getSingleSourceForAllAttributes(expressionContext);
       if ((source.isPresent() && EDS.toString().equals(source.get())) || !hasTimeRange(request)) {

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

@@ -1,6 +1,6 @@
 package org.hypertrace.gateway.service.explore;
 
-import org.hypertrace.gateway.service.v1.explore.ExploreRequest;
+import org.hypertrace.gateway.service.common.ExpressionContext;
 import org.hypertrace.gateway.service.v1.explore.ExploreResponse;
 
 /**
@@ -17,5 +17,5 @@
  */
 interface IRequestHandler {
   ExploreResponse.Builder handleRequest(
-      ExploreRequestContext requestContext, ExploreRequest request);
+      ExploreRequestContext requestContext, ExpressionContext expressionContext);
 }

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

@@ -18,6 +18,7 @@
 import org.hypertrace.core.query.service.api.Row;
 import org.hypertrace.core.query.service.api.Value;
 import org.hypertrace.gateway.service.common.AttributeMetadataProvider;
+import org.hypertrace.gateway.service.common.ExpressionContext;
 import org.hypertrace.gateway.service.common.converters.QueryAndGatewayDtoConverter;
 import org.hypertrace.gateway.service.common.util.AttributeMetadataUtil;
 import org.hypertrace.gateway.service.common.util.DataCollectionUtil;
@@ -44,14 +45,14 @@ public RequestHandler(
       QueryServiceClient queryServiceClient, AttributeMetadataProvider attributeMetadataProvider) {
     this.queryServiceClient = queryServiceClient;
     this.attributeMetadataProvider = attributeMetadataProvider;
-    this.theRestGroupRequestHandler = new TheRestGroupRequestHandler(this);
+    this.theRestGroupRequestHandler =
+        new TheRestGroupRequestHandler(this, attributeMetadataProvider);
   }
 
   @Override
   public ExploreResponse.Builder handleRequest(
-      ExploreRequestContext requestContext, ExploreRequest request) {
-    QueryRequest queryRequest =
-        buildQueryRequest(requestContext, request, attributeMetadataProvider);
+      ExploreRequestContext requestContext, ExpressionContext expressionContext) {
+    QueryRequest queryRequest = buildQueryRequest(requestContext, attributeMetadataProvider);
 
     Iterator<ResultSetChunk> resultSetChunkIterator = executeQuery(requestContext, queryRequest);
 
@@ -60,11 +61,10 @@ public ExploreResponse.Builder handleRequest(
   }
 
   QueryRequest buildQueryRequest(
-      ExploreRequestContext requestContext,
-      ExploreRequest request,
-      AttributeMetadataProvider attributeMetadataProvider) {
+      ExploreRequestContext requestContext, AttributeMetadataProvider attributeMetadataProvider) {
     // Track if we have Group By so we can determine if we need to do Order By, Limit and Offset
     // ourselves.
+    ExploreRequest request = requestContext.getRequest();
     if (!request.getGroupByList().isEmpty()) {
       requestContext.setHasGroupBy(true);
     }
@@ -221,7 +221,7 @@ private ExploreResponse.Builder handleQueryServiceResponse(
 
     if (requestContext.hasGroupBy() && requestContext.getIncludeRestGroup()) {
       theRestGroupRequestHandler.getRowsForTheRestGroup(
-          context, requestContext.getExploreRequest(), builder);
+          context, requestContext.getRequest(), builder);
     }
 
     return builder;
@@ -273,8 +273,7 @@ void handleQueryServiceResponseSingleColumn(
         attributeMetadataProvider.getAttributesMetadata(
             requestContext, requestContext.getContext());
     Map<String, AttributeMetadata> resultKeyToAttributeMetadataMap =
-        this.remapAttributeMetadataByResultName(
-            requestContext.getExploreRequest(), attributeMetadataMap);
+        this.remapAttributeMetadataByResultName(requestContext.getRequest(), attributeMetadataMap);
     org.hypertrace.gateway.service.v1.common.Value gwValue;
     if (function != null) { // Function expression value
       gwValue =

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

@@ -6,6 +6,9 @@
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
+import org.hypertrace.core.attribute.service.v1.AttributeMetadata;
+import org.hypertrace.gateway.service.common.AttributeMetadataProvider;
+import org.hypertrace.gateway.service.common.ExpressionContext;
 import org.hypertrace.gateway.service.common.util.ExpressionReader;
 import org.hypertrace.gateway.service.v1.common.Expression;
 import org.hypertrace.gateway.service.v1.common.Filter;
@@ -30,34 +33,49 @@
 public class TheRestGroupRequestHandler {
   private static final String OTHER_COLUMN_VALUE = "__Other";
   private final RequestHandlerWithSorting requestHandler;
+  private final AttributeMetadataProvider attributeMetadataProvider;
 
-  TheRestGroupRequestHandler(RequestHandlerWithSorting requestHandler) {
+  TheRestGroupRequestHandler(
+      RequestHandlerWithSorting requestHandler,
+      AttributeMetadataProvider attributeMetadataProvider) {
     this.requestHandler = requestHandler;
+    this.attributeMetadataProvider = attributeMetadataProvider;
   }
 
   public void getRowsForTheRestGroup(
-      ExploreRequestContext context,
+      ExploreRequestContext requestContext,
       ExploreRequest originalRequest,
       ExploreResponse.Builder originalResponse) {
     // Return if there was no data in the original request
     if (originalResponse.getRowBuilderList().isEmpty()) {
       return;
     }
 
+    Map<String, AttributeMetadata> attributeMetadataMap =
+        attributeMetadataProvider.getAttributesMetadata(
+            requestContext, originalRequest.getContext());
+
     ExploreRequest theRestRequest = createRequest(originalRequest, originalResponse);
     ExploreRequestContext theRestRequestContext =
-        new ExploreRequestContext(context.getGrpcContext(), theRestRequest);
-
+        new ExploreRequestContext(requestContext.getGrpcContext(), theRestRequest);
+    ExpressionContext theRestExpressionContext =
+        new ExpressionContext(
+            attributeMetadataMap,
+            theRestRequest.getFilter(),
+            theRestRequest.getSelectionList(),
+            theRestRequest.getTimeAggregationList(),
+            theRestRequest.getOrderByList(),
+            theRestRequest.getGroupByList());
     ExploreResponse.Builder theRestGroupResponse =
-        requestHandler.handleRequest(theRestRequestContext, theRestRequest);
+        requestHandler.handleRequest(theRestRequestContext, theRestExpressionContext);
 
     List<OrderByExpression> orderByExpressions =
         requestHandler.getRequestOrderByExpressions(theRestRequest);
     mergeAndSort(
         originalResponse,
         theRestGroupResponse,
         orderByExpressions,
-        context.getRowLimitAfterRest(), // check how many rows expected from original request
+        requestContext.getRowLimitAfterRest(), // check how many rows expected from original request
         originalRequest.getOffset(),
         requestHandler,
         originalRequest.getGroupByList());

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

@@ -39,11 +39,10 @@ public class TimeAggregationsRequestHandler extends RequestHandler {
 
   @Override
   QueryRequest buildQueryRequest(
-      ExploreRequestContext requestContext,
-      ExploreRequest request,
-      AttributeMetadataProvider attributeMetadataProvider) {
+      ExploreRequestContext requestContext, AttributeMetadataProvider attributeMetadataProvider) {
     // Set hasGroupBy=true in the request context since we will group by the timestamp column
     // regardless of the presence of a groupBy or not.
+    ExploreRequest request = requestContext.getRequest();
     requestContext.setHasGroupBy(true);
     QueryRequest.Builder builder = QueryRequest.newBuilder();
 

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

@@ -1,8 +1,11 @@
 package org.hypertrace.gateway.service.explore;
 
 import com.google.common.collect.ImmutableSet;
+import java.util.Map;
 import java.util.Set;
+import org.hypertrace.core.attribute.service.v1.AttributeMetadata;
 import org.hypertrace.gateway.service.common.AttributeMetadataProvider;
+import org.hypertrace.gateway.service.common.ExpressionContext;
 import org.hypertrace.gateway.service.common.util.ExpressionReader;
 import org.hypertrace.gateway.service.common.util.QueryServiceClient;
 import org.hypertrace.gateway.service.v1.common.Expression;
@@ -16,27 +19,40 @@
 
 public class TimeAggregationsWithGroupByRequestHandler implements IRequestHandler {
 
+  private final AttributeMetadataProvider attributeMetadataProvider;
   private final RequestHandler normalRequestHandler;
   private final TimeAggregationsRequestHandler timeAggregationsRequestHandler;
 
   TimeAggregationsWithGroupByRequestHandler(
       QueryServiceClient queryServiceClient, AttributeMetadataProvider attributeMetadataProvider) {
+    this.attributeMetadataProvider = attributeMetadataProvider;
     this.normalRequestHandler = new RequestHandler(queryServiceClient, attributeMetadataProvider);
     this.timeAggregationsRequestHandler =
         new TimeAggregationsRequestHandler(queryServiceClient, attributeMetadataProvider);
   }
 
   @Override
   public ExploreResponse.Builder handleRequest(
-      ExploreRequestContext requestContext, ExploreRequest request) {
+      ExploreRequestContext requestContext, ExpressionContext expressionContext) {
     // This type of handler is always a group by
+    ExploreRequest request = requestContext.getRequest();
     requestContext.setHasGroupBy(true);
+    Map<String, AttributeMetadata> attributeMetadataMap =
+        attributeMetadataProvider.getAttributesMetadata(requestContext, request.getContext());
     // 1. Create a GroupBy request and get the response for the GroupBy
     ExploreRequest groupByRequest = buildGroupByRequest(request);
     ExploreRequestContext groupByRequestContext =
         new ExploreRequestContext(requestContext.getGrpcContext(), groupByRequest);
+    ExpressionContext groupByExpressionContext =
+        new ExpressionContext(
+            attributeMetadataMap,
+            groupByRequest.getFilter(),
+            groupByRequest.getSelectionList(),
+            groupByRequest.getTimeAggregationList(),
+            groupByRequest.getOrderByList(),
+            groupByRequest.getGroupByList());
     ExploreResponse.Builder groupByResponse =
-        normalRequestHandler.handleRequest(groupByRequestContext, groupByRequest);
+        normalRequestHandler.handleRequest(groupByRequestContext, groupByExpressionContext);
 
     // No need for a second query if no results.
     if (groupByResponse.getRowBuilderList().isEmpty()) {
@@ -48,9 +64,17 @@ public ExploreResponse.Builder handleRequest(
     ExploreRequest timeAggregationsRequest = buildTimeAggregationsRequest(request, groupByResponse);
     ExploreRequestContext timeAggregationsRequestContext =
         new ExploreRequestContext(requestContext.getGrpcContext(), timeAggregationsRequest);
+    ExpressionContext timeAggregationsExpressionContext =
+        new ExpressionContext(
+            attributeMetadataMap,
+            timeAggregationsRequest.getFilter(),
+            timeAggregationsRequest.getSelectionList(),
+            timeAggregationsRequest.getTimeAggregationList(),
+            timeAggregationsRequest.getOrderByList(),
+            timeAggregationsRequest.getGroupByList());
     ExploreResponse.Builder timeAggregationsResponse =
         timeAggregationsRequestHandler.handleRequest(
-            timeAggregationsRequestContext, timeAggregationsRequest);
+            timeAggregationsRequestContext, timeAggregationsExpressionContext);
 
     // 3. If includeRestGroup is set, invoke TheRestGroupRequestHandler
     if (request.getIncludeRestGroup()) {