diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/QueryRequestContext.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/QueryRequestContext.java index 145949e7..7e30ecd8 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/QueryRequestContext.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/QueryRequestContext.java @@ -24,6 +24,10 @@ public void mapAliasToFunctionExpression(String alias, FunctionExpression functi aliasToFunctionExpressionMap.put(alias, functionExpression); } + public boolean containsFunctionExpression(String alias) { + return aliasToFunctionExpressionMap.containsKey(alias); + } + public FunctionExpression getFunctionExpressionByAlias(String alias) { return aliasToFunctionExpressionMap.get(alias); } diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/EntityDataServiceEntityFetcher.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/EntityDataServiceEntityFetcher.java index ee1259d6..841500c1 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/EntityDataServiceEntityFetcher.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/EntityDataServiceEntityFetcher.java @@ -169,12 +169,6 @@ public EntityFetcherResponse getEntities( return new EntityFetcherResponse(entityBuilders); } - @Override - public EntityFetcherResponse getAggregatedMetrics( - EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest) { - throw new UnsupportedOperationException("Fetching aggregated metrics not supported by EDS"); - } - @Override public EntityFetcherResponse getTimeAggregatedMetrics( EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest) { diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/IEntityFetcher.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/IEntityFetcher.java index b00e6d73..b81f0dac 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/IEntityFetcher.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/IEntityFetcher.java @@ -27,16 +27,6 @@ public interface IEntityFetcher { EntityFetcherResponse getEntities( EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest); - /** - * Get aggregated metrics - * - * @param requestContext Additional context for the incoming request - * @param entitiesRequest encapsulates the aggregated metrics query (selection, filter, order) - * @return Map of the Entity Builders keyed by the EntityId - */ - EntityFetcherResponse getAggregatedMetrics( - EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest); - /** * Get time series data * diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/QueryServiceEntityFetcher.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/QueryServiceEntityFetcher.java index 7ebe972d..52fde8bb 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/QueryServiceEntityFetcher.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/QueryServiceEntityFetcher.java @@ -157,7 +157,10 @@ public EntityFetcherResponse getEntities( i++) { ColumnMetadata metadata = chunk.getResultSetMetadata().getColumnMetadata(i); org.hypertrace.core.query.service.api.Value columnValue = row.getColumn(i); - addEntityAttribute(entityBuilder, + buildEntity( + entityBuilder, + requestContext, + entitiesRequest, metadata, columnValue, attributeMetadataMap, @@ -169,97 +172,6 @@ public EntityFetcherResponse getEntities( return new EntityFetcherResponse(entityBuilders); } - @Override - public EntityFetcherResponse getAggregatedMetrics( - EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest) { - // Only supported filter is entityIds IN ["id1", "id2", "id3"] - Map attributeMetadataMap = - attributeMetadataProvider.getAttributesMetadata( - requestContext, entitiesRequest.getEntityType()); - entitiesRequestValidator.validate(entitiesRequest, attributeMetadataMap); - - List aggregates = - ExpressionReader.getFunctionExpressions(entitiesRequest.getSelectionList().stream()); - if (aggregates.isEmpty()) { - return new EntityFetcherResponse(); - } - - List entityIdAttributes = - AttributeMetadataUtil.getIdAttributeIds( - attributeMetadataProvider, entityIdColumnsConfigs, requestContext, entitiesRequest.getEntityType()); - - QueryRequest.Builder builder = - constructSelectionQuery(requestContext, entitiesRequest, entityIdAttributes, aggregates); - adjustLimitAndOffset(builder, entitiesRequest.getLimit(), entitiesRequest.getOffset()); - - QueryRequest request = builder.build(); - - if (LOG.isDebugEnabled()) { - LOG.debug("Sending Aggregated Metrics Request to Query Service ======== \n {}", request); - } - - Iterator resultSetChunkIterator = - queryServiceClient.executeQuery(request, requestContext.getHeaders(), - requestTimeout); - - // We want to retain the order as returned from the respective source. Hence using a - // LinkedHashMap - Map entityMap = new LinkedHashMap<>(); - - while (resultSetChunkIterator.hasNext()) { - ResultSetChunk chunk = resultSetChunkIterator.next(); - if (LOG.isDebugEnabled()) { - LOG.debug("Received chunk: " + chunk.toString()); - } - - if (chunk.getRowCount() < 1) { - break; - } - - if (!chunk.hasResultSetMetadata()) { - LOG.warn("Chunk doesn't have result metadata so couldn't process the response."); - break; - } - - for (Row row : chunk.getRowList()) { - // Construct the EntityKey from the EntityId attributes columns - EntityKey entityKey = - EntityKey.of( - IntStream.range(0, entityIdAttributes.size()) - .mapToObj(value -> row.getColumn(value).getString()) - .toArray(String[]::new)); - Builder entityBuilder = entityMap.computeIfAbsent(entityKey, k -> Entity.newBuilder()); - entityBuilder.setEntityType(entitiesRequest.getEntityType()); - entityBuilder.setId(entityKey.toString()); - // Always include the id in entity since that's needed to make follow up queries in - // optimal fashion. If this wasn't really requested by the client, it should be removed - // as post processing. - for (int i = 0; i < entityIdAttributes.size(); i++) { - entityBuilder.putAttribute( - entityIdAttributes.get(i), - Value.newBuilder() - .setString(entityKey.getAttributes().get(i)) - .setValueType(ValueType.STRING) - .build()); - } - - for (int i = entityIdAttributes.size(); - i < chunk.getResultSetMetadata().getColumnMetadataCount(); - i++) { - ColumnMetadata metadata = chunk.getResultSetMetadata().getColumnMetadata(i); - org.hypertrace.core.query.service.api.Value columnValue = row.getColumn(i); - addAggregateMetric(entityBuilder, - requestContext, - entitiesRequest, - metadata, - columnValue, - attributeMetadataMap); - } - } - } - return new EntityFetcherResponse(entityMap); - } - private void adjustLimitAndOffset(QueryRequest.Builder builder, int limit, int offset) { // If there is more than one groupBy column, we cannot set the same limit that came // in the request since that might return less entities than needed when the same @@ -336,28 +248,51 @@ private QueryRequest.Builder constructSelectionQuery(EntitiesRequestContext requ return builder; } - private void addEntityAttribute(Entity.Builder entityBuilder, + private void buildEntity( + Entity.Builder entityBuilder, + QueryRequestContext requestContext, + EntitiesRequest entitiesRequest, ColumnMetadata metadata, org.hypertrace.core.query.service.api.Value columnValue, Map attributeMetadataMap, boolean isSkipCountColumn) { // Ignore the count column since we introduced that ourselves into the query - if (isSkipCountColumn && - StringUtils.equalsIgnoreCase(COUNT_COLUMN_NAME, metadata.getColumnName())) { + if (isSkipCountColumn + && StringUtils.equalsIgnoreCase(COUNT_COLUMN_NAME, metadata.getColumnName())) { return; } + // aggregate + if (requestContext.containsFunctionExpression(metadata.getColumnName())) { + addAggregateMetric( + entityBuilder, + requestContext, + entitiesRequest, + metadata, + columnValue, + attributeMetadataMap); + } else { + // attribute + addEntityAttribute(entityBuilder, metadata, columnValue, attributeMetadataMap); + } + } + + private void addEntityAttribute( + Entity.Builder entityBuilder, + ColumnMetadata metadata, + org.hypertrace.core.query.service.api.Value columnValue, + Map attributeMetadataMap) { + String attributeName = metadata.getColumnName(); entityBuilder.putAttribute( attributeName, QueryAndGatewayDtoConverter.convertToGatewayValue( - attributeName, - columnValue, - attributeMetadataMap)); + attributeName, columnValue, attributeMetadataMap)); } - private void addAggregateMetric(Entity.Builder entityBuilder, + private void addAggregateMetric( + Entity.Builder entityBuilder, QueryRequestContext requestContext, EntitiesRequest entitiesRequest, ColumnMetadata metadata, diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/ExecutionContext.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/ExecutionContext.java index 8cb4fdf7..57a6efe0 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/ExecutionContext.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/ExecutionContext.java @@ -206,7 +206,7 @@ public void removePendingSelectionSource(String source) { pendingSelectionSources.remove(source); } - public void removePendingMetricAggregationSources(String source) { + public void removePendingMetricAggregationSource(String source) { pendingMetricAggregationSources.remove(source); } @@ -214,6 +214,10 @@ public void removePendingSelectionSourceForOrderBy(String source) { pendingSelectionSourcesForOrderBy.remove(source); } + public void removePendingMetricAggregationSourceForOrderBy(String source) { + pendingMetricAggregationSourcesForOrderBy.remove(source); + } + public Map> getSourceToFilterExpressionMap() { return sourceToFilterExpressionMap; } diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/ExecutionTreeBuilder.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/ExecutionTreeBuilder.java index 75eb1821..af8a1820 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/ExecutionTreeBuilder.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/ExecutionTreeBuilder.java @@ -197,7 +197,7 @@ QueryNode buildExecutionTree(ExecutionContext executionContext, QueryNode filter new SelectionNode.Builder(rootNode) .setAggMetricSelectionSources(metricSourcesForOrderBy) .build(); - metricSourcesForOrderBy.forEach(executionContext::removePendingMetricAggregationSources); + metricSourcesForOrderBy.forEach(executionContext::removePendingMetricAggregationSource); } // Try adding SortAndPaginateNode diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionContextBuilderVisitor.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionContextBuilderVisitor.java index 4472a18e..f6307b1b 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionContextBuilderVisitor.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionContextBuilderVisitor.java @@ -58,6 +58,13 @@ public Void visit(DataFetcherNode dataFetcherNode) { executionContext.removePendingSelectionSource(source); // TODO: Currently, assumes that the order by attribute is also present in the selection set executionContext.removePendingSelectionSourceForOrderBy(source); + // TODO: Remove redundant attributes for metric aggregation source for order by + // The current metric aggregation source is only QS + + // The order by on metric aggregations will also be added in the selections list of + // DataFetcherNode, so that the order by metric aggregations can be fetched before + // and only the required data set is paginated + executionContext.removePendingMetricAggregationSourceForOrderBy(source); // set of attributes which were fetched from the source Map> sourceToSelectionAttributeMap = diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionVisitor.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionVisitor.java index 310eddfb..fe167699 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionVisitor.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionVisitor.java @@ -6,18 +6,18 @@ import java.util.ArrayList; import java.util.Arrays; +import java.util.Collection; import java.util.Collections; import java.util.LinkedHashMap; import java.util.LinkedList; import java.util.List; import java.util.Map; import java.util.Set; -import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.Future; import java.util.stream.Collectors; import java.util.stream.IntStream; +import java.util.stream.Stream; import com.google.common.collect.Sets; import org.hypertrace.gateway.service.common.datafetcher.EntityFetcherResponse; @@ -41,6 +41,7 @@ import org.hypertrace.gateway.service.v1.common.Filter; import org.hypertrace.gateway.service.v1.common.LiteralConstant; import org.hypertrace.gateway.service.v1.common.Operator; +import org.hypertrace.gateway.service.v1.common.OrderByExpression; import org.hypertrace.gateway.service.v1.common.Value; import org.hypertrace.gateway.service.v1.common.ValueType; import org.hypertrace.gateway.service.v1.entity.EntitiesRequest; @@ -145,17 +146,32 @@ public EntityResponse visit(DataFetcherNode dataFetcherNode) { executionContext.getTimestampAttributeId(), executionContext.getRequestHeaders()); + // fetching both attribute selections and metric order by selections for + // optimized pagination + List attributeSelections = + executionContext + .getSourceToSelectionExpressionMap() + .getOrDefault(source, executionContext.getEntityIdExpressions()); + List metricOrderBySelections = + executionContext + .getSourceToMetricOrderByExpressionMap() + .getOrDefault(source, Collections.emptyList()) + .stream() + .map(OrderByExpression::getExpression) + .collect(Collectors.toList()); + List selections = Stream.of(attributeSelections, metricOrderBySelections) + .flatMap(Collection::stream) + .collect(Collectors.toList()); + EntitiesRequest.Builder requestBuilder = EntitiesRequest.newBuilder(entitiesRequest) .clearSelection() + .clearTimeAggregation() .clearFilter() .clearOrderBy() .clearLimit() .clearOffset() - .addAllSelection( - executionContext - .getSourceToSelectionExpressionMap() - .getOrDefault(source, executionContext.getEntityIdExpressions())) + .addAllSelection(selections) .setFilter(dataFetcherNode.getFilter()); if (dataFetcherNode.getLimit() != null) { @@ -243,6 +259,7 @@ public EntityResponse visit(SelectionNode selectionNode) { EntitiesRequest request = EntitiesRequest.newBuilder(executionContext.getEntitiesRequest()) .clearSelection() + .clearTimeAggregation() .clearFilter() // TODO: Should we push order by, limit and offet down to the data source? // If we want to push the order by down, we would also have to divide order by into @@ -273,6 +290,7 @@ public EntityResponse visit(SelectionNode selectionNode) { EntitiesRequest request = EntitiesRequest.newBuilder(executionContext.getEntitiesRequest()) .clearSelection() + .clearTimeAggregation() .clearFilter() .clearOrderBy() .clearOffset() @@ -290,7 +308,7 @@ public EntityResponse visit(SelectionNode selectionNode) { request.getEntityType(), executionContext.getTimestampAttributeId(), executionContext.getRequestHeaders()); - return entityFetcher.getAggregatedMetrics(context, request); + return entityFetcher.getEntities(context, request); }) .collect(Collectors.toList())); resultMapList.addAll( @@ -299,6 +317,7 @@ public EntityResponse visit(SelectionNode selectionNode) { source -> { EntitiesRequest request = EntitiesRequest.newBuilder(executionContext.getEntitiesRequest()) + .clearSelection() .clearTimeAggregation() .clearFilter() .clearOrderBy() diff --git a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/EntitiesRequestAndResponseUtils.java b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/EntitiesRequestAndResponseUtils.java index 10854738..cb5cf3e0 100644 --- a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/EntitiesRequestAndResponseUtils.java +++ b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/EntitiesRequestAndResponseUtils.java @@ -115,6 +115,13 @@ public static AggregatedMetricValue getAggregatedMetricValue(FunctionType functi .build(); } + public static AggregatedMetricValue getAggregatedMetricValue(FunctionType functionType, long value) { + return AggregatedMetricValue.newBuilder() + .setFunction(functionType) + .setValue(Value.newBuilder().setLong(value).setValueType(ValueType.LONG)) + .build(); + } + public static Expression getLiteralExpression(long value) { return Expression.newBuilder() .setLiteral( diff --git a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/datafetcher/EntityDataServiceEntityFetcherTests.java b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/datafetcher/EntityDataServiceEntityFetcherTests.java index 0c8d0fc6..9c5cd354 100644 --- a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/datafetcher/EntityDataServiceEntityFetcherTests.java +++ b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/datafetcher/EntityDataServiceEntityFetcherTests.java @@ -176,17 +176,6 @@ public void test_getEntities_WithoutPagination() { entityDataServiceEntityFetcher.getEntities(entitiesRequestContext, entitiesRequest).size()); } - @Test - public void test_getAggregatedMetrics() { - assertThrows( - UnsupportedOperationException.class, - () -> { - entityDataServiceEntityFetcher.getAggregatedMetrics( - new EntitiesRequestContext(TENANT_ID, 0, 1, "API", "API.startTime", Map.of()), - EntitiesRequest.newBuilder().build()); - }); - } - @Test public void test_getTimeAggregatedMetrics() { assertThrows( diff --git a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/datafetcher/QueryServiceEntityFetcherTests.java b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/datafetcher/QueryServiceEntityFetcherTests.java index d786f328..5b70071a 100644 --- a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/datafetcher/QueryServiceEntityFetcherTests.java +++ b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/datafetcher/QueryServiceEntityFetcherTests.java @@ -6,12 +6,12 @@ import static org.hypertrace.gateway.service.common.EntitiesRequestAndResponseUtils.buildTimeAggregation; import static org.hypertrace.gateway.service.common.EntitiesRequestAndResponseUtils.compareEntityFetcherResponses; import static org.hypertrace.gateway.service.common.EntitiesRequestAndResponseUtils.generateEQFilter; +import static org.hypertrace.gateway.service.common.EntitiesRequestAndResponseUtils.getAggregatedMetricValue; import static org.hypertrace.gateway.service.common.EntitiesRequestAndResponseUtils.getStringValue; import static org.hypertrace.gateway.service.common.QueryServiceRequestAndResponseUtils.createQsAggregationExpression; import static org.hypertrace.gateway.service.common.QueryServiceRequestAndResponseUtils.createQsRequestFilter; import static org.hypertrace.gateway.service.common.QueryServiceRequestAndResponseUtils.getResultSetChunk; import static org.hypertrace.gateway.service.common.converters.QueryRequestUtil.createColumnExpression; -import static org.hypertrace.gateway.service.common.converters.QueryRequestUtil.createCountByColumnSelection; import static org.hypertrace.gateway.service.common.converters.QueryRequestUtil.createStringFilter; import static org.hypertrace.gateway.service.v1.common.Operator.AND; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -20,6 +20,8 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import java.util.Collections; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Optional; @@ -77,6 +79,92 @@ public void setup () { attributeMetadataProvider, entityIdColumnsConfigs); } + @Test + public void testGetEntities() { + List orderByExpressions = List.of(buildOrderByExpression(API_ID_ATTR)); + long startTime = 1L; + long endTime = 10L; + int limit = 10; + int offset = 5; + String tenantId = "TENANT_ID"; + Map requestHeaders = Map.of("x-tenant-id", tenantId); + AttributeScope entityType = AttributeScope.API; + EntitiesRequest entitiesRequest = + EntitiesRequest.newBuilder() + .setEntityType(entityType.name()) + .setStartTimeMillis(startTime) + .setEndTimeMillis(endTime) + .addSelection(buildExpression(API_NAME_ATTR)) + .addSelection( + buildAggregateExpression( + API_NUM_CALLS_ATTR, FunctionType.SUM, "Sum_numCalls", Collections.emptyList())) + .setFilter( + Filter.newBuilder() + .setOperator(AND) + .addChildFilter( + EntitiesRequestAndResponseUtils.getTimeRangeFilter( + "API.startTime", startTime, endTime)) + .addChildFilter(generateEQFilter(API_DISCOVERY_STATE_ATTR, "DISCOVERED"))) + .addAllOrderBy(orderByExpressions) + .setLimit(limit) + .setOffset(offset) + .build(); + EntitiesRequestContext entitiesRequestContext = + new EntitiesRequestContext( + tenantId, startTime, endTime, entityType.name(), "API.startTime", requestHeaders); + + QueryRequest expectedQueryRequest = + QueryRequest.newBuilder() + .addSelection(createColumnExpression(API_ID_ATTR)) + .addSelection(createQsAggregationExpression("SUM", API_NUM_CALLS_ATTR, "Sum_numCalls")) + .addSelection(createColumnExpression(API_NAME_ATTR)) + .setFilter( + createQsRequestFilter( + API_START_TIME_ATTR, + API_ID_ATTR, + startTime, + endTime, + createStringFilter(API_DISCOVERY_STATE_ATTR, Operator.EQ, "DISCOVERED"))) + .addGroupBy(createColumnExpression(API_ID_ATTR)) + .addGroupBy(createColumnExpression(API_NAME_ATTR)) + .setOffset(offset) + .setLimit(QueryServiceClient.DEFAULT_QUERY_SERVICE_GROUP_BY_LIMIT) + .addAllOrderBy( + QueryAndGatewayDtoConverter.convertToQueryOrderByExpressions(orderByExpressions)) + .build(); + + List resultSetChunks = + List.of( + getResultSetChunk( + List.of("API.id", "API.name", "Sum_numCalls"), + new String[][] {{"apiId1", "api 1", "3"}, {"apiId2", "api 2", "5"}})); + + when(queryServiceClient.executeQuery(eq(expectedQueryRequest), eq(requestHeaders), eq(500))) + .thenReturn(resultSetChunks.iterator()); + + EntityFetcherResponse response = + queryServiceEntityFetcher.getEntities(entitiesRequestContext, entitiesRequest); + assertEquals(2, response.size()); + + Map expectedEntityKeyBuilderResponseMap = new LinkedHashMap<>(); + expectedEntityKeyBuilderResponseMap.put(EntityKey.of("apiId1"), Entity.newBuilder() + .setId("apiId1") + .setEntityType("API") + .putAttribute("API.id", getStringValue("apiId1")) + .putAttribute("API.name", getStringValue("api 1")) + .putMetric("Sum_numCalls", getAggregatedMetricValue(FunctionType.SUM, 3)) + ); + expectedEntityKeyBuilderResponseMap.put(EntityKey.of("apiId2"), Entity.newBuilder() + .setId("apiId2") + .setEntityType("API") + .putAttribute("API.id", getStringValue("apiId2")) + .putAttribute("API.name", getStringValue("api 2")) + .putMetric("Sum_numCalls", getAggregatedMetricValue(FunctionType.SUM, 5)) + ); + compareEntityFetcherResponses( + new EntityFetcherResponse(expectedEntityKeyBuilderResponseMap), response); + } + @Test public void test_getEntitiesWithPagination() { List orderByExpressions = List.of(buildOrderByExpression(API_ID_ATTR)); diff --git a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionVisitorTest.java b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionVisitorTest.java index bd75f289..28e70f9f 100644 --- a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionVisitorTest.java +++ b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/entity/query/visitor/ExecutionVisitorTest.java @@ -684,6 +684,7 @@ public void test_visitPaginateOnlyNode() { when(executionContext.getTimestampAttributeId()).thenReturn("API.startTime"); EntitiesRequest entitiesRequestForAttributes = EntitiesRequest.newBuilder(entitiesRequest) .clearSelection() + .clearTimeAggregation() .addSelection(selectionExpression) .setLimit(limit + offset) .setOffset(0) @@ -694,11 +695,13 @@ public void test_visitPaginateOnlyNode() { .clearOffset() .clearOrderBy() .clearSelection() + .clearTimeAggregation() .addSelection(metricExpression) .clearFilter() .setFilter(generateInFilter(API_ID_ATTR, List.of("entity-id-3", "entity-id-2"))) .build(); EntitiesRequest entitiesRequestForTimeAggregation = EntitiesRequest.newBuilder(entitiesRequest) + .clearSelection() .clearLimit() .clearOffset() .clearOrderBy() @@ -711,7 +714,7 @@ public void test_visitPaginateOnlyNode() { when(queryServiceEntityFetcher.getEntities( eq(entitiesRequestContext), eq(totalEntitiesRequest))) .thenReturn(totalEntityFetcherResponse); - when(queryServiceEntityFetcher.getAggregatedMetrics( + when(queryServiceEntityFetcher.getEntities( entitiesRequestContext, entitiesRequestForMetricAggregation)) .thenReturn(new EntityFetcherResponse(entityKeyBuilderResponseMap2)); when(queryServiceEntityFetcher.getTimeAggregatedMetrics( @@ -753,12 +756,12 @@ public void test_visitSelectionNode_differentSource_callSeparatedCalls() { Map.of(EDS_SOURCE, Collections.emptyList()), Map.of(QS_SOURCE, Collections.emptyList())); when(entityDataServiceEntityFetcher.getEntities(any(), any())).thenReturn(result4); - when(queryServiceEntityFetcher.getAggregatedMetrics(any(), any())).thenReturn(result4); + when(queryServiceEntityFetcher.getEntities(any(), any())).thenReturn(result4); when(executionVisitor.visit(any(NoOpNode.class))) .thenReturn(new EntityResponse(result4, result4.getEntityKeyBuilderMap().keySet())); executionVisitor.visit(selectionNode); verify(entityDataServiceEntityFetcher).getEntities(any(), any()); - verify(queryServiceEntityFetcher).getAggregatedMetrics(any(), any()); + verify(queryServiceEntityFetcher).getEntities(any(), any()); } @Test @@ -865,7 +868,7 @@ public void test_visitOnlySelectionsNode_shouldSetTotalEntityKeys() { when(queryServiceEntityFetcher.getEntities( entitiesRequestContext, entitiesRequestForAttributes)) .thenReturn(attributesResponse); - when(queryServiceEntityFetcher.getAggregatedMetrics( + when(queryServiceEntityFetcher.getEntities( entitiesRequestContext, entitiesRequestForMetricAggregation)) .thenReturn(new EntityFetcherResponse(entityKeyBuilderResponseMap2)); @@ -931,7 +934,6 @@ public void test_visitSelectionNode_nonEmptyFilter_emptyResult() { Assertions.assertTrue(response.getEntityFetcherResponse().isEmpty()); Assertions.assertTrue(response.getEntityKeys().isEmpty()); verify(queryServiceEntityFetcher, never()).getEntities(any(), any()); - verify(queryServiceEntityFetcher, never()).getAggregatedMetrics(any(), any()); } private MetricSeries getMockMetricSeries(int period, String aggregation) {