refactor(expression-context): expression location based single source

gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/ExpressionLocation.java

@@ -0,0 +1,11 @@
+package org.hypertrace.gateway.service.common;
+
+public enum ExpressionLocation {
+  COLUMN_SELECTION,
+  METRIC_AGGREGATION,
+  TIME_AGGREGATION,
+  COLUMN_ORDER_BY,
+  METRIC_ORDER_BY,
+  COLUMN_FILTER,
+  COLUMN_GROUP_BY,
+}

gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/transformer/ResponsePostProcessor.java

@@ -30,7 +30,7 @@ public List<Entity.Builder> transform(
             .flatMap(Optional::stream)
             .collect(Collectors.toSet());
     Set<String> aggregations =
-        executionContext.getExpressionContext().getSourceToMetricExpressionMap().values().stream()
+        executionContext.getExpressionContext().getSourceToMetricAggregationExpressionMap().values().stream()
             .flatMap(Collection::stream)
             .map(expression -> expression.getFunction().getAlias())
             .collect(Collectors.toSet());

gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/EntityExecutionContext.java

@@ -186,7 +186,7 @@ public void removeSelectionAttributes(String source, Set<String> attributes) {
   private void buildSelectionPendingSources() {
     pendingSelectionSources.addAll(expressionContext.getSourceToSelectionExpressionMap().keySet());
     pendingMetricAggregationSources.addAll(
-        expressionContext.getSourceToMetricExpressionMap().keySet());
+        expressionContext.getSourceToMetricAggregationExpressionMap().keySet());
     pendingTimeAggregationSources.addAll(
         expressionContext.getSourceToTimeAggregationMap().keySet());
   }

gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionVisitor.java

@@ -261,7 +261,7 @@ public EntityResponse visit(SelectionNode selectionNode) {
                           .addAllSelection(
                               executionContext
                                   .getExpressionContext()
-                                  .getSourceToMetricExpressionMap()
+                                  .getSourceToMetricAggregationExpressionMap()
                                   .get(source))
                           .setFilter(filter)
                           .build();

gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/ExpressionContextTest.java

@@ -31,12 +31,7 @@ void testGetSingleSourceForAllAttributes_allSourceExpressionMapKeySetsHaveOneSou
             createSourceToExpressionsMap(List.of("QS")),
             createSourceToExpressionsMap(List.of("QS")),
             createSourceToExpressionsMap(List.of("QS")),
-            createSourceToExpressionsMap(List.of("QS")),
-            Map.of(
-                "API.id", Set.of(QS.name()),
-                "API.name", Set.of(QS.name()),
-                "API.duration", Set.of(QS.name()),
-                "API.startTime", Set.of(QS.name())));
+            createSourceToExpressionsMap(List.of("QS")));
 
     assertEquals(
         Optional.of("QS"), ExpressionContext.getSingleSourceForAllAttributes(expressionContext));
@@ -51,12 +46,7 @@ void testGetSingleSourceForAllAttributes_someSourceExpressionMapKeySetsHaveMulti
             createSourceToExpressionsMap(List.of("QS", "EDS")),
             createSourceToExpressionsMap(List.of("QS")),
             createSourceToExpressionsMap(List.of("QS")),
-            createSourceToExpressionsMap(List.of("QS")),
-            Map.of(
-                "API.id", Set.of(QS.name()),
-                "API.name", Set.of(QS.name()),
-                "API.duration", Set.of(QS.name(), EDS.name()),
-                "API.startTime", Set.of(QS.name())));
+            createSourceToExpressionsMap(List.of("QS")));
 
     assertEquals(
         Optional.of("QS"), ExpressionContext.getSingleSourceForAllAttributes(expressionContext));
@@ -71,12 +61,7 @@ void testGetSingleSourceForAllAttributes_someSourceExpressionMapKeySetsHaveDiffe
             createSourceToExpressionsMap(List.of("EDS")),
             createSourceToExpressionsMap(List.of("QS")),
             createSourceToExpressionsMap(List.of("QS")),
-            createSourceToExpressionsMap(List.of("QS")),
-            Map.of(
-                "API.id", Set.of(QS.name()),
-                "API.name", Set.of(QS.name()),
-                "API.duration", Set.of(EDS.name()),
-                "API.startTime", Set.of(QS.name())));
+            createSourceToExpressionsMap(List.of("QS")));
 
     assertEquals(
         Optional.empty(), ExpressionContext.getSingleSourceForAllAttributes(expressionContext));
@@ -92,12 +77,7 @@ void testGetSingleSourceForAllAttributes_someSourceExpressionMapKeySetsAreEmpty(
             createSourceToExpressionsMap(List.of()),
             createSourceToExpressionsMap(List.of("QS")),
             createSourceToExpressionsMap(List.of("QS")),
-            createSourceToExpressionsMap(List.of()),
-            Map.of(
-                "API.id", Set.of(QS.name()),
-                "API.name", Set.of(QS.name()),
-                "API.duration", Set.of(),
-                "API.startTime", Set.of(QS.name())));
+            createSourceToExpressionsMap(List.of()));
 
     assertEquals(
         Optional.of("QS"), ExpressionContext.getSingleSourceForAllAttributes(expressionContext));
@@ -113,12 +93,7 @@ void testGetSingleSourceForAllAttributes_allEmptySourceExpressionMapKeySetsAndAt
             createSourceToExpressionsMap(List.of()),
             createSourceToExpressionsMap(List.of()),
             createSourceToExpressionsMap(List.of()),
-            createSourceToExpressionsMap(List.of()),
-            Map.of(
-                "API.id", Set.of(),
-                "API.name", Set.of(),
-                "API.duration", Set.of(),
-                "API.startTime", Set.of()));
+            createSourceToExpressionsMap(List.of()));
 
     assertEquals(
         Optional.empty(), ExpressionContext.getSingleSourceForAllAttributes(expressionContext));
@@ -309,12 +284,11 @@ private ExpressionContext getMockExpressionContext(
       Map<String, List<TimeAggregation>> sourceToTimeAggregationMap,
       Map<String, List<OrderByExpression>> sourceToSelectionOrderByExpressionMap,
       Map<String, List<OrderByExpression>> sourceToMetricOrderByExpressionMap,
-      Map<String, List<Expression>> sourceToFilterExpressionMap,
-      Map<String, Set<String>> attributeToSourcesMap) {
+      Map<String, List<Expression>> sourceToFilterExpressionMap) {
     ExpressionContext expressionContext = mock(ExpressionContext.class);
     when(expressionContext.getSourceToSelectionExpressionMap())
         .thenReturn(sourceToSelectionExpressionMap);
-    when(expressionContext.getSourceToMetricExpressionMap())
+    when(expressionContext.getSourceToMetricAggregationExpressionMap())
         .thenReturn(sourceToMetricExpressionMap);
     when(expressionContext.getSourceToTimeAggregationMap()).thenReturn(sourceToTimeAggregationMap);
     when(expressionContext.getSourceToSelectionOrderByExpressionMap())
@@ -323,7 +297,6 @@ private ExpressionContext getMockExpressionContext(
         .thenReturn(sourceToMetricOrderByExpressionMap);
     when(expressionContext.getSourceToFilterExpressionMap())
         .thenReturn(sourceToFilterExpressionMap);
-    when(expressionContext.getAllAttributesToSourcesMap()).thenReturn(attributeToSourcesMap);
 
     return expressionContext;
   }

gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/transformer/ResponsePostProcessorTest.java

@@ -39,7 +39,7 @@ public void shouldAddMissingSelections() {
                 List.of(createAttributeExpression("API.id"), createAttributeExpression("API.name")),
                 "QS",
                 List.of(createAttributeExpression("API.serviceId"))));
-    when(expressionContext.getSourceToMetricExpressionMap()).thenReturn(Collections.emptyMap());
+    when(expressionContext.getSourceToMetricAggregationExpressionMap()).thenReturn(Collections.emptyMap());
     when(expressionContext.getSourceToTimeAggregationMap()).thenReturn(Collections.emptyMap());
 
     EntityExecutionContext executionContext = mock(EntityExecutionContext.class);
@@ -81,7 +81,7 @@ public void shouldAddMissingSelections() {
   public void shouldAddMissingAggregations() {
     ExpressionContext expressionContext = mock(ExpressionContext.class);
     when(expressionContext.getSourceToSelectionExpressionMap()).thenReturn(Collections.emptyMap());
-    when(expressionContext.getSourceToMetricExpressionMap())
+    when(expressionContext.getSourceToMetricAggregationExpressionMap())
         .thenReturn(Map.of("QS", List.of(createAggregateExpression("API.duration"))));
     when(expressionContext.getSourceToTimeAggregationMap()).thenReturn(Collections.emptyMap());
 
@@ -119,7 +119,7 @@ public void shouldAddMissingAggregations() {
   public void shouldAddMissingTimeAggregations() {
     ExpressionContext expressionContext = mock(ExpressionContext.class);
     when(expressionContext.getSourceToSelectionExpressionMap()).thenReturn(Collections.emptyMap());
-    when(expressionContext.getSourceToMetricExpressionMap()).thenReturn(Collections.emptyMap());
+    when(expressionContext.getSourceToMetricAggregationExpressionMap()).thenReturn(Collections.emptyMap());
     when(expressionContext.getSourceToTimeAggregationMap())
         .thenReturn(Map.of("QS", List.of(createTimeAggregation("API.duration"))));
 

gateway-service-impl/src/test/java/org/hypertrace/gateway/service/entity/query/ExecutionTreeUtilsTest.java

@@ -8,14 +8,9 @@
 import static org.mockito.Mockito.when;
 
 import java.util.Collections;
-import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.stream.Collectors;
 import org.hypertrace.gateway.service.common.ExpressionContext;
-import org.hypertrace.gateway.service.v1.common.Expression;
-import org.hypertrace.gateway.service.v1.common.OrderByExpression;
-import org.hypertrace.gateway.service.v1.common.TimeAggregation;
 import org.junit.jupiter.api.Test;
 
 class ExecutionTreeUtilsTest {
@@ -54,34 +49,4 @@ void removeDuplicateSelectionAttributes() {
     verify(executionContext).removeSelectionAttributes("EDS", Set.of("API.name"));
     verify(executionContext).removeSelectionAttributes("AS", Set.of("API.latency"));
   }
-
-  private ExpressionContext getMockExpressionContext(
-      Map<String, List<Expression>> sourceToSelectionExpressionMap,
-      Map<String, List<Expression>> sourceToMetricExpressionMap,
-      Map<String, List<TimeAggregation>> sourceToTimeAggregationMap,
-      Map<String, List<OrderByExpression>> sourceToSelectionOrderByExpressionMap,
-      Map<String, List<OrderByExpression>> sourceToMetricOrderByExpressionMap,
-      Map<String, List<Expression>> sourceToFilterExpressionMap,
-      Map<String, Set<String>> attributeToSourcesMap) {
-    ExpressionContext expressionContext = mock(ExpressionContext.class);
-    when(expressionContext.getSourceToSelectionExpressionMap())
-        .thenReturn(sourceToSelectionExpressionMap);
-    when(expressionContext.getSourceToMetricExpressionMap())
-        .thenReturn(sourceToMetricExpressionMap);
-    when(expressionContext.getSourceToTimeAggregationMap()).thenReturn(sourceToTimeAggregationMap);
-    when(expressionContext.getSourceToSelectionOrderByExpressionMap())
-        .thenReturn(sourceToSelectionOrderByExpressionMap);
-    when(expressionContext.getSourceToMetricOrderByExpressionMap())
-        .thenReturn(sourceToMetricOrderByExpressionMap);
-    when(expressionContext.getSourceToFilterExpressionMap())
-        .thenReturn(sourceToFilterExpressionMap);
-    when(expressionContext.getAllAttributesToSourcesMap()).thenReturn(attributeToSourcesMap);
-
-    return expressionContext;
-  }
-
-  // We don't care about the values. Just they keySets.
-  private <T> Map<String, List<T>> createSourceToExpressionsMap(List<String> sourceKeys) {
-    return sourceKeys.stream().collect(Collectors.toUnmodifiableMap(s -> s, s -> List.of()));
-  }
 }

gateway-service-impl/src/test/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionVisitorTest.java

@@ -684,7 +684,7 @@ public void test_visitPaginateOnlyNode() {
         .thenReturn(List.of(buildExpression(API_ID_ATTR)));
     when(expressionContext.getSourceToSelectionExpressionMap())
         .thenReturn(Map.of("QS", List.of(selectionExpression)));
-    when(expressionContext.getSourceToMetricExpressionMap())
+    when(expressionContext.getSourceToMetricAggregationExpressionMap())
         .thenReturn(Map.of("QS", List.of(metricExpression)));
     when(expressionContext.getSourceToTimeAggregationMap())
         .thenReturn(Map.of("QS", List.of(timeAggregation)));
@@ -854,7 +854,7 @@ public void test_visitOnlySelectionsNode_shouldSetTotalEntityKeys() {
         .thenReturn(List.of(buildExpression(API_ID_ATTR)));
     when(expressionContext.getSourceToSelectionExpressionMap())
         .thenReturn(Map.of("QS", List.of(selectionExpression)));
-    when(expressionContext.getSourceToMetricExpressionMap())
+    when(expressionContext.getSourceToMetricAggregationExpressionMap())
         .thenReturn(Map.of("QS", List.of(metricExpression)));
     when(expressionContext.getSourceToTimeAggregationMap())
         .thenReturn(Map.of("QS", List.of(timeAggregation)));
@@ -986,7 +986,7 @@ private EntityExecutionContext mockExecutionContext(
     when(executionContext.getPendingSelectionSources()).thenReturn(selectionSource);
     when(executionContext.getPendingMetricAggregationSources()).thenReturn(aggregateSource);
     when(expressionContext.getSourceToSelectionExpressionMap()).thenReturn(sourceToSelectionMap);
-    when(expressionContext.getSourceToMetricExpressionMap()).thenReturn(sourceToAggregateMap);
+    when(expressionContext.getSourceToMetricAggregationExpressionMap()).thenReturn(sourceToAggregateMap);
     return executionContext;
   }