diff --git a/gateway-service-impl/build.gradle.kts b/gateway-service-impl/build.gradle.kts index 3a4c5380..dbe39281 100644 --- a/gateway-service-impl/build.gradle.kts +++ b/gateway-service-impl/build.gradle.kts @@ -25,8 +25,8 @@ dependencies { implementation("org.hypertrace.core.query.service:query-service-client:0.5.2") implementation("org.hypertrace.core.attribute.service:attribute-service-client:0.9.3") - implementation("org.hypertrace.entity.service:entity-service-client:0.5.0") - implementation("org.hypertrace.entity.service:entity-service-api:0.5.0") + implementation("org.hypertrace.entity.service:entity-service-client:0.5.6") + implementation("org.hypertrace.entity.service:entity-service-api:0.5.6") implementation("org.hypertrace.core.grpcutils:grpc-context-utils:0.3.2") implementation("org.hypertrace.core.serviceframework:platform-metrics:0.1.19") diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/converters/QueryRequestUtil.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/converters/QueryRequestUtil.java index 80ee3da2..32c2d086 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/converters/QueryRequestUtil.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/converters/QueryRequestUtil.java @@ -18,6 +18,9 @@ public class QueryRequestUtil { public static final String DATE_TIME_CONVERTER = "dateTimeConvert"; + private static final String COUNT_FUNCTION_NAME = "COUNT"; + private static final String DISTINCTCOUNT_FUNCTION_NAME = "DISTINCTCOUNT"; + public static Filter createBetweenTimesFilter(String columnName, long lower, long higher) { return Filter.newBuilder() .setOperator(Operator.AND) @@ -102,7 +105,16 @@ public static Expression createCountByColumnSelection(String columnName) { return Expression.newBuilder() .setFunction( Function.newBuilder() - .setFunctionName("COUNT") + .setFunctionName(COUNT_FUNCTION_NAME) + .addArguments(createColumnExpression(columnName))) + .build(); + } + + public static Expression createDistinctCountByColumnSelection(String columnName) { + return Expression.newBuilder() + .setFunction( + Function.newBuilder() + .setFunctionName(DISTINCTCOUNT_FUNCTION_NAME) .addArguments(createColumnExpression(columnName))) .build(); } 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 841500c1..ddc0e116 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 @@ -11,6 +11,8 @@ import org.hypertrace.entity.query.service.v1.EntityQueryRequest; import org.hypertrace.entity.query.service.v1.ResultSetChunk; import org.hypertrace.entity.query.service.v1.Row; +import org.hypertrace.entity.query.service.v1.TotalEntitiesRequest; +import org.hypertrace.entity.query.service.v1.TotalEntitiesResponse; import org.hypertrace.gateway.service.common.AttributeMetadataProvider; import org.hypertrace.gateway.service.common.converters.EntityServiceAndGatewayServiceConverter; import org.hypertrace.gateway.service.common.util.AttributeMetadataUtil; @@ -50,7 +52,10 @@ public EntityFetcherResponse getEntities( EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest) { List mappedEntityIdAttributeIds = AttributeMetadataUtil.getIdAttributeIds( - attributeMetadataProvider, entityIdColumnsConfigs, requestContext, entitiesRequest.getEntityType()); + attributeMetadataProvider, + entityIdColumnsConfigs, + requestContext, + entitiesRequest.getEntityType()); EntityQueryRequest.Builder builder = EntityQueryRequest.newBuilder() .setEntityType(entitiesRequest.getEntityType()) @@ -174,4 +179,35 @@ public EntityFetcherResponse getTimeAggregatedMetrics( EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest) { throw new UnsupportedOperationException("Fetching time series data not supported by EDS"); } + + @Override + public long getTotal(EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest) { + EntityQueryRequest.Builder builder = + EntityQueryRequest.newBuilder() + .setEntityType(entitiesRequest.getEntityType()) + .setFilter( + EntityServiceAndGatewayServiceConverter.convertToEntityServiceFilter( + entitiesRequest.getFilter())); + + // add time filter for supported scope + EntityServiceAndGatewayServiceConverter.addBetweenTimeFilter( + entitiesRequest.getStartTimeMillis(), + entitiesRequest.getEndTimeMillis(), + attributeMetadataProvider, + entitiesRequest, + builder, + requestContext); + + EntityQueryRequest entityQueryRequest = builder.build(); + + TotalEntitiesRequest totalEntitiesRequest = + TotalEntitiesRequest.newBuilder() + .setEntityType(entitiesRequest.getEntityType()) + .setFilter(entityQueryRequest.getFilter()) + .build(); + + TotalEntitiesResponse response = + entityQueryServiceClient.total(totalEntitiesRequest, requestContext.getHeaders()); + return response.getTotal(); + } } diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/EntityResponse.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/EntityResponse.java index e589adda..a83b1200 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/EntityResponse.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/EntityResponse.java @@ -1,30 +1,24 @@ package org.hypertrace.gateway.service.common.datafetcher; -import org.hypertrace.gateway.service.entity.EntityKey; - -import java.util.HashSet; -import java.util.Set; - public class EntityResponse { private final EntityFetcherResponse entityFetcherResponse; - // set of entity keys, irrespective of offset and limit - private final Set entityKeys; + private final long total; public EntityResponse() { this.entityFetcherResponse = new EntityFetcherResponse(); - this.entityKeys = new HashSet<>(); + this.total = 0; } - public EntityResponse(EntityFetcherResponse entityFetcherResponse, Set entityKeys) { + public EntityResponse(EntityFetcherResponse entityFetcherResponse, long total) { this.entityFetcherResponse = entityFetcherResponse; - this.entityKeys = entityKeys; + this.total = total; } public EntityFetcherResponse getEntityFetcherResponse() { return entityFetcherResponse; } - public Set getEntityKeys() { - return entityKeys; + public long getTotal() { + return total; } } 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 b81f0dac..0bdddeb4 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 @@ -3,10 +3,10 @@ import java.util.ArrayList; import java.util.Comparator; import java.util.List; -import java.util.Set; import org.hypertrace.gateway.service.entity.EntitiesRequestContext; import org.hypertrace.gateway.service.entity.EntityKey; +import org.hypertrace.gateway.service.v1.common.Filter; import org.hypertrace.gateway.service.v1.common.Interval; import org.hypertrace.gateway.service.v1.common.MetricSeries; import org.hypertrace.gateway.service.v1.entity.EntitiesRequest; @@ -46,4 +46,6 @@ default MetricSeries getSortedMetricSeries(MetricSeries.Builder builder) { .addAllValue(sortedIntervals) .build(); } + + long getTotal(EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest); } 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 52fde8bb..da7a5b42 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 @@ -2,6 +2,7 @@ import static org.hypertrace.gateway.service.common.converters.QueryAndGatewayDtoConverter.convertToQueryExpression; import static org.hypertrace.gateway.service.common.converters.QueryRequestUtil.createCountByColumnSelection; +import static org.hypertrace.gateway.service.common.converters.QueryRequestUtil.createDistinctCountByColumnSelection; import static org.hypertrace.gateway.service.common.converters.QueryRequestUtil.createFilter; import static org.hypertrace.gateway.service.common.converters.QueryRequestUtil.createStringNullLiteralExpression; import static org.hypertrace.gateway.service.common.converters.QueryRequestUtil.createTimeColumnGroupByExpression; @@ -488,6 +489,66 @@ public EntityFetcherResponse getTimeAggregatedMetrics( return new EntityFetcherResponse(resultMap); } + @Override + public long getTotal(EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest) { + Map attributeMetadataMap = + attributeMetadataProvider.getAttributesMetadata( + requestContext, entitiesRequest.getEntityType()); + // Validate EntitiesRequest + entitiesRequestValidator.validate(entitiesRequest, attributeMetadataMap); + + List entityIdAttributes = + AttributeMetadataUtil.getIdAttributeIds( + attributeMetadataProvider, + entityIdColumnsConfigs, + requestContext, + entitiesRequest.getEntityType()); + + Filter.Builder filterBuilder = + constructQueryServiceFilter(entitiesRequest, requestContext, entityIdAttributes); + + QueryRequest queryRequest = + QueryRequest.newBuilder() + .addSelection( + createDistinctCountByColumnSelection( + Optional.ofNullable(entityIdAttributes.get(0)).orElseThrow())) + .setFilter(filterBuilder) + .build(); + + if (LOG.isDebugEnabled()) { + LOG.debug("Sending Query to Query Service ======== \n {}", queryRequest); + } + + Iterator resultSetChunkIterator = + queryServiceClient.executeQuery(queryRequest, requestContext.getHeaders(), + requestTimeout); + + while (resultSetChunkIterator.hasNext()) { + ResultSetChunk chunk = resultSetChunkIterator.next(); + if (LOG.isDebugEnabled()) { + LOG.debug("Received chunk: " + chunk.toString()); + } + + if (chunk.getRowCount() < 1) { + break; + } + + for (Row row : chunk.getRowList()) { + // only DISTINCTCOUNT column is requested as a selection + if (row.getColumnList().size() != 1) { + break; + } + + return Long.parseLong(row.getColumn(0).getString()); + } + } + + LOG.error( + "Unable to query total number of entities from query service for query request: {}", + queryRequest); + return 0; + } + private QueryRequest buildTimeSeriesQueryRequest( EntitiesRequest entitiesRequest, EntitiesRequestContext context, diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/EntityService.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/EntityService.java index 4d23f3a5..2a11fe05 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/EntityService.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/EntityService.java @@ -144,7 +144,6 @@ public EntitiesResponse getEntities( executionTree.acceptVisitor(new ExecutionVisitor(executionContext, EntityQueryHandlerRegistry.get())); EntityFetcherResponse entityFetcherResponse = response.getEntityFetcherResponse(); - Set allEntityKeys = response.getEntityKeys(); List results = this.responsePostProcessor.transform( @@ -157,7 +156,7 @@ public EntitiesResponse getEntities( } EntitiesResponse.Builder responseBuilder = - EntitiesResponse.newBuilder().setTotal(allEntityKeys.size()); + EntitiesResponse.newBuilder().setTotal(Long.valueOf(response.getTotal()).intValue()); results.forEach(e -> responseBuilder.addEntity(e.build())); 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 fe21393a..5fc7a0aa 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 @@ -17,7 +17,6 @@ import java.util.stream.Collectors; import java.util.stream.IntStream; -import com.google.common.collect.Sets; import org.hypertrace.gateway.service.common.datafetcher.EntityFetcherResponse; import org.hypertrace.gateway.service.common.datafetcher.EntityResponse; import org.hypertrace.gateway.service.common.datafetcher.IEntityFetcher; @@ -90,13 +89,8 @@ protected static EntityResponse intersect(List entityResponses) entityResponses.parallelStream() .map(EntityResponse::getEntityFetcherResponse) .collect(Collectors.toList())); - Set entityKeys = - entityResponses.parallelStream() - .map(EntityResponse::getEntityKeys) - .reduce(Sets::intersection) - .orElse(Collections.emptySet()); - return new EntityResponse(entityFetcherResponse, entityKeys); + return new EntityResponse(entityFetcherResponse, entityFetcherResponse.size()); } private static EntityFetcherResponse unionEntities(List builders) { @@ -121,13 +115,8 @@ protected static EntityResponse union(List entityResponses) { entityResponses.parallelStream() .map(EntityResponse::getEntityFetcherResponse) .collect(Collectors.toList())); - Set entityKeys = - entityResponses.parallelStream() - .map(EntityResponse::getEntityKeys) - .reduce(Sets::union) - .orElse(Collections.emptySet()); - return new EntityResponse(entityFetcherResponse, entityKeys); + return new EntityResponse(entityFetcherResponse, entityFetcherResponse.size()); } @Override @@ -175,27 +164,16 @@ public EntityResponse visit(DataFetcherNode dataFetcherNode) { // if the data fetcher node is fetching paginated records and the client has requested for // total, the total number of entities has to be fetched separately if (dataFetcherNode.canFetchTotal()) { - EntitiesRequest totalEntitiesRequest = - EntitiesRequest.newBuilder(executionContext.getEntitiesRequest()) - .clearSelection() - .clearTimeAggregation() - .clearOrderBy() - .clearLimit() - .setOffset(0) - .setFilter(dataFetcherNode.getFilter()) - .build(); - + // since, the pagination is pushed down to the data store, total can be requested directly + // from the data store return new EntityResponse( entityFetcher.getEntities(context, request), - entityFetcher - .getEntities(context, totalEntitiesRequest) - .getEntityKeyBuilderMap() - .keySet()); + entityFetcher.getTotal(context, entitiesRequest)); } else { // if the data fetcher node is not paginating, the total number of entities is equal to number // of records fetched EntityFetcherResponse response = entityFetcher.getEntities(context, request); - return new EntityResponse(response, response.getEntityKeyBuilderMap().keySet()); + return new EntityResponse(response, response.getEntityKeyBuilderMap().size()); } } @@ -328,14 +306,14 @@ public EntityResponse visit(SelectionNode selectionNode) { .reduce(childEntityFetcherResponse, (r1, r2) -> unionEntities(Arrays.asList(r1, r2))); if (!childEntityFetcherResponse.isEmpty()) { - // if the child fetcher response is non empty, the total set of entity keys + // if the child fetcher response is non empty, the total // has already been fetched by node below it. // Could be DataFetcherNode or a child SelectionNode - return new EntityResponse(response, childNodeResponse.getEntityKeys()); + return new EntityResponse(response, childNodeResponse.getTotal()); } else { - // if the child fetcher response is empty, the total set of entity keys + // if the child fetcher response is empty, the total // is equal to the response fetched by the current SelectionNode - return new EntityResponse(response, response.getEntityKeyBuilderMap().keySet()); + return new EntityResponse(response, response.size()); } } @@ -413,7 +391,7 @@ public EntityResponse visit(SortAndPaginateNode sortAndPaginateNode) { Map linkedHashMap = new LinkedHashMap<>(); sortedList.forEach(entry -> linkedHashMap.put(entry.getKey(), entry.getValue())); return new EntityResponse( - new EntityFetcherResponse(linkedHashMap), childNodeResponse.getEntityKeys()); + new EntityFetcherResponse(linkedHashMap), childNodeResponse.getTotal()); } @Override @@ -442,6 +420,6 @@ public EntityResponse visit(PaginateOnlyNode paginateOnlyNode) { sortedList.forEach(entry -> linkedHashMap.put(entry.getKey(), entry.getValue())); return new EntityResponse( - new EntityFetcherResponse(linkedHashMap), childNodeResponse.getEntityKeys()); + new EntityFetcherResponse(linkedHashMap), childNodeResponse.getTotal()); } } 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 cb5cf3e0..fafc85c8 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 @@ -157,7 +157,7 @@ public static void compareEntityResponses(EntityResponse expectedEntityResponse, expectedEntityResponse.getEntityFetcherResponse(), actualEntityResponse.getEntityFetcherResponse()); - assertEquals(expectedEntityResponse.getEntityKeys(), actualEntityResponse.getEntityKeys()); + assertEquals(expectedEntityResponse.getTotal(), actualEntityResponse.getTotal()); } public static Filter getTimeRangeFilter(String colName, long startTime, long endTime) { 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 9c5cd354..f0b530e7 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 @@ -9,6 +9,8 @@ import org.hypertrace.entity.query.service.v1.ResultSetChunk; import org.hypertrace.entity.query.service.v1.ResultSetMetadata; import org.hypertrace.entity.query.service.v1.Row; +import org.hypertrace.entity.query.service.v1.TotalEntitiesRequest; +import org.hypertrace.entity.query.service.v1.TotalEntitiesResponse; import org.hypertrace.entity.query.service.v1.Value; import org.hypertrace.entity.query.service.v1.ValueType; import org.hypertrace.gateway.service.common.AttributeMetadataProvider; @@ -22,6 +24,7 @@ import org.hypertrace.gateway.service.v1.common.OrderByExpression; import org.hypertrace.gateway.service.v1.entity.EntitiesRequest; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; import java.util.List; @@ -187,6 +190,50 @@ public void test_getTimeAggregatedMetrics() { }); } + @Nested + class TotalEntities { + @Test + public void shouldReturnTotal() { + 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) + .setFilter( + Filter.newBuilder() + .setOperator(AND) + .addChildFilter(generateEQFilter(API_TYPE_ATTR, "HTTP")) + .addChildFilter(generateEQFilter(API_NAME_ATTR, "DISCOVERED"))) + .addAllOrderBy(orderByExpressions) + .setLimit(limit) + .setOffset(offset) + .build(); + EntitiesRequestContext entitiesRequestContext = + new EntitiesRequestContext( + tenantId, startTime, endTime, entityType.name(), "API.startTime", requestHeaders); + + TotalEntitiesRequest expectedQueryRequest = + TotalEntitiesRequest.newBuilder() + .setEntityType("API") + .setFilter(convertToEntityServiceFilter(entitiesRequest.getFilter())) + .build(); + + when(entityQueryServiceClient.total(eq(expectedQueryRequest), eq(requestHeaders))) + .thenReturn(TotalEntitiesResponse.newBuilder().setTotal(100L).build()); + + assertEquals( + 100, entityDataServiceEntityFetcher.getTotal(entitiesRequestContext, entitiesRequest)); + } + } + private void mockAttributeMetadataProvider(String attributeScope) { AttributeMetadata idAttributeMetadata = AttributeMetadata.newBuilder() 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 5b70071a..cba8b042 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 @@ -47,6 +47,7 @@ import org.hypertrace.gateway.service.v1.entity.Entity; import org.hypertrace.gateway.service.v1.entity.Entity.Builder; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; public class QueryServiceEntityFetcherTests { @@ -358,6 +359,63 @@ public void test_getEntitiesBySpace() { queryServiceEntityFetcher.getEntities(entitiesRequestContext, entitiesRequest)); } + @Nested + class TotalEntities { + @Test + public void shouldReturnTotal() { + 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) + .addTimeAggregation( + buildTimeAggregation( + 30, API_NUM_CALLS_ATTR, FunctionType.SUM, "SUM_API.numCalls", List.of())) + .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(createQsAggregationExpression("DISTINCTCOUNT", API_ID_ATTR)) + .setFilter( + createQsRequestFilter( + API_START_TIME_ATTR, + API_ID_ATTR, + startTime, + endTime, + createStringFilter(API_DISCOVERY_STATE_ATTR, Operator.EQ, "DISCOVERED"))) + .build(); + + List resultSetChunks = + List.of(getResultSetChunk(List.of("DISTINCCOUNT"), new String[][] {{"100"}})); + + when(queryServiceClient.executeQuery(eq(expectedQueryRequest), eq(requestHeaders), eq(500))) + .thenReturn(resultSetChunks.iterator()); + + assertEquals( + 100, queryServiceEntityFetcher.getTotal(entitiesRequestContext, entitiesRequest)); + } + } private void mockAttributeMetadataProvider(String attributeScope) { AttributeMetadata idAttributeMetadata = AttributeMetadata.newBuilder() 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 5498acfa..56657767 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 @@ -8,6 +8,7 @@ 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.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; @@ -134,15 +135,15 @@ public void testIntersect() { EntityResponse finalResult = ExecutionVisitor.intersect( Arrays.asList( - new EntityResponse(result1, result1.getEntityKeyBuilderMap().keySet()), - new EntityResponse(result2, result2.getEntityKeyBuilderMap().keySet()), - new EntityResponse(result3, result3.getEntityKeyBuilderMap().keySet()))); + new EntityResponse(result1, result1.getEntityKeyBuilderMap().size()), + new EntityResponse(result2, result2.getEntityKeyBuilderMap().size()), + new EntityResponse(result3, result3.getEntityKeyBuilderMap().size()))); Map finalEntities = finalResult.getEntityFetcherResponse().getEntityKeyBuilderMap(); - Set finalEntityKeys = finalResult.getEntityKeys(); + long total = finalResult.getTotal(); Assertions.assertEquals(1, finalEntities.size()); - Assertions.assertEquals(1, finalEntityKeys.size()); + Assertions.assertEquals(1, total); Entity.Builder builder = finalEntities.get(EntityKey.of("id1")); Assertions.assertNotNull(builder); @@ -154,14 +155,14 @@ public void testIntersect() { EntityResponse finalResult = ExecutionVisitor.intersect( Arrays.asList( - new EntityResponse(result1, result1.getEntityKeyBuilderMap().keySet()), - new EntityResponse(result2, result2.getEntityKeyBuilderMap().keySet()), - new EntityResponse(result4, result4.getEntityKeyBuilderMap().keySet()))); + new EntityResponse(result1, result1.getEntityKeyBuilderMap().size()), + new EntityResponse(result2, result2.getEntityKeyBuilderMap().size()), + new EntityResponse(result4, result4.getEntityKeyBuilderMap().size()))); Map finalEntities = finalResult.getEntityFetcherResponse().getEntityKeyBuilderMap(); - Set finalEntityKeys = finalResult.getEntityKeys(); + long total = finalResult.getTotal(); assertTrue(finalEntities.isEmpty()); - assertTrue(finalEntityKeys.isEmpty()); + assertEquals(0, total); } } @@ -171,14 +172,14 @@ public void testUnion() { EntityResponse finalResult = ExecutionVisitor.union( Arrays.asList( - new EntityResponse(result1, result1.getEntityKeyBuilderMap().keySet()), - new EntityResponse(result4, result4.getEntityKeyBuilderMap().keySet()))); + new EntityResponse(result1, result1.getEntityKeyBuilderMap().size()), + new EntityResponse(result4, result4.getEntityKeyBuilderMap().size()))); Map finalEntities = finalResult.getEntityFetcherResponse().getEntityKeyBuilderMap(); - Set finalEntityKeys = finalResult.getEntityKeys(); + long total = finalResult.getTotal(); Assertions.assertEquals(4, finalEntities.size()); - Assertions.assertEquals(4, finalEntityKeys.size()); + Assertions.assertEquals(4, total); assertTrue( finalEntities .keySet() @@ -186,12 +187,7 @@ public void testUnion() { Stream.of("id1", "id2", "id3", "id4") .map(EntityKey::of) .collect(Collectors.toList()))); - assertTrue( - finalEntityKeys - .containsAll( - Stream.of("id1", "id2", "id3", "id4") - .map(EntityKey::of) - .collect(Collectors.toList()))); + assertEquals(4, total); Assertions.assertEquals( result1.getEntityKeyBuilderMap().get(EntityKey.of("id1")), finalEntities.get(EntityKey.of("id1"))); @@ -392,14 +388,6 @@ public void test_visitDataFetcherNodeQs() { .setLimit(limit) .setOffset(offset) .build(); - EntitiesRequest totalEntitiesRequest = - EntitiesRequest.newBuilder(entitiesRequest) - .clearSelection() - .clearTimeAggregation() - .clearLimit() - .setOffset(0) - .clearOrderBy() - .build(); EntitiesRequestContext entitiesRequestContext = new EntitiesRequestContext( tenantId, startTime, @@ -413,17 +401,8 @@ public void test_visitDataFetcherNodeQs() { EntityKey.of("entity-id-2"), Entity.newBuilder().putAttribute("API.name", getStringValue("entity-2")) ); - Map totalEntityKeyBuilderResponseMap = Map.of( - EntityKey.of("entity-id-0"), Entity.newBuilder().putAttribute("API.id", getStringValue("entity-0")), - EntityKey.of("entity-id-1"), Entity.newBuilder().putAttribute("API.id", getStringValue("entity-1")), - EntityKey.of("entity-id-2"), Entity.newBuilder().putAttribute("API.id", getStringValue("entity-2")), - EntityKey.of("entity-id-3"), Entity.newBuilder().putAttribute("API.id", getStringValue("entity-3")), - EntityKey.of("entity-id-4"), Entity.newBuilder().putAttribute("API.id", getStringValue("entity-4")) - ); - EntityFetcherResponse entityFetcherResponse = new EntityFetcherResponse(entityKeyBuilderResponseMap); - EntityFetcherResponse totalEntityFetcherResponse = - new EntityFetcherResponse(totalEntityKeyBuilderResponseMap); + when(executionContext.getSourceToSelectionExpressionMap()) .thenReturn(Map.of("QS", List.of(selectionExpression))); when(executionContext.getEntitiesRequest()).thenReturn(entitiesRequest); @@ -432,9 +411,9 @@ public void test_visitDataFetcherNodeQs() { when(executionContext.getTimestampAttributeId()).thenReturn("API.startTime"); when(queryServiceEntityFetcher.getEntities(eq(entitiesRequestContext), eq(entitiesRequest))) .thenReturn(entityFetcherResponse); - when(queryServiceEntityFetcher.getEntities( - eq(entitiesRequestContext), eq(totalEntitiesRequest))) - .thenReturn(totalEntityFetcherResponse); + when(queryServiceEntityFetcher.getTotal( + eq(entitiesRequestContext), eq(entitiesRequest))) + .thenReturn(100L); when(queryServiceEntityFetcher.getTimeAggregatedMetrics(eq(entitiesRequestContext), eq(entitiesRequest))) .thenReturn(new EntityFetcherResponse()); @@ -443,8 +422,7 @@ public void test_visitDataFetcherNodeQs() { "QS", entitiesRequest.getFilter(), limit, offset, orderByExpressions, true); compareEntityResponses( - new EntityResponse( - entityFetcherResponse, totalEntityFetcherResponse.getEntityKeyBuilderMap().keySet()), + new EntityResponse(entityFetcherResponse, 100L), executionVisitor.visit(dataFetcherNode)); } @@ -501,7 +479,7 @@ public void test_visitDataFetcherNode_cannotFetchTotal() { compareEntityResponses( new EntityResponse( - entityFetcherResponse, entityFetcherResponse.getEntityKeyBuilderMap().keySet()), + entityFetcherResponse, entityFetcherResponse.getEntityKeyBuilderMap().size()), executionVisitor.visit(dataFetcherNode)); } @@ -527,14 +505,6 @@ public void test_visitDataFetcherNodeEds() { .setLimit(limit) .setOffset(offset) .build(); - EntitiesRequest totalEntitiesRequest = - EntitiesRequest.newBuilder(entitiesRequest) - .clearSelection() - .clearTimeAggregation() - .clearLimit() - .setOffset(0) - .clearOrderBy() - .build(); EntitiesRequestContext entitiesRequestContext = new EntitiesRequestContext( tenantId, startTime, @@ -547,17 +517,8 @@ public void test_visitDataFetcherNodeEds() { EntityKey.of("entity-id-1"), Entity.newBuilder().putAttribute("API.name", getStringValue("entity-1")), EntityKey.of("entity-id-2"), Entity.newBuilder().putAttribute("API.name", getStringValue("entity-2")) ); - Map totalEntityKeyBuilderResponseMap = Map.of( - EntityKey.of("entity-id-0"), Entity.newBuilder().putAttribute("API.id", getStringValue("entity-0")), - EntityKey.of("entity-id-1"), Entity.newBuilder().putAttribute("API.id", getStringValue("entity-1")), - EntityKey.of("entity-id-2"), Entity.newBuilder().putAttribute("API.id", getStringValue("entity-2")), - EntityKey.of("entity-id-3"), Entity.newBuilder().putAttribute("API.id", getStringValue("entity-3")), - EntityKey.of("entity-id-4"), Entity.newBuilder().putAttribute("API.id", getStringValue("entity-4")) - ); - EntityFetcherResponse entityFetcherResponse = new EntityFetcherResponse(entityKeyBuilderResponseMap); - EntityFetcherResponse totalEntityFetcherResponse = - new EntityFetcherResponse(totalEntityKeyBuilderResponseMap); + when(executionContext.getSourceToSelectionExpressionMap()) .thenReturn(Map.of("EDS", List.of(selectionExpression))); when(executionContext.getEntitiesRequest()).thenReturn(entitiesRequest); @@ -566,15 +527,14 @@ public void test_visitDataFetcherNodeEds() { when(executionContext.getTimestampAttributeId()).thenReturn("API.startTime"); when(entityDataServiceEntityFetcher.getEntities(eq(entitiesRequestContext), eq(entitiesRequest))) .thenReturn(entityFetcherResponse); - when(entityDataServiceEntityFetcher.getEntities( - eq(entitiesRequestContext), eq(totalEntitiesRequest))) - .thenReturn(totalEntityFetcherResponse); + when(entityDataServiceEntityFetcher.getTotal( + eq(entitiesRequestContext), eq(entitiesRequest))) + .thenReturn(100L); DataFetcherNode dataFetcherNode = new DataFetcherNode("EDS", entitiesRequest.getFilter(), limit, offset, orderByExpressions, true); compareEntityResponses( - new EntityResponse( - entityFetcherResponse, totalEntityFetcherResponse.getEntityKeyBuilderMap().keySet()), + new EntityResponse(entityFetcherResponse, 100), executionVisitor.visit(dataFetcherNode)); } @@ -624,7 +584,7 @@ public void test_visitDataFetcherNodeWithoutPagination() { compareEntityResponses( new EntityResponse( - entityFetcherResponse, entityFetcherResponse.getEntityKeyBuilderMap().keySet()), + entityFetcherResponse, entityFetcherResponse.getEntityKeyBuilderMap().size()), executionVisitor.visit(dataFetcherNode)); verify(queryServiceEntityFetcher, times(1)).getEntities(any(), any()); } @@ -659,14 +619,6 @@ public void test_visitPaginateOnlyNode() { .setLimit(limit) .setOffset(offset) .build(); - EntitiesRequest totalEntitiesRequest = - EntitiesRequest.newBuilder(entitiesRequest) - .clearSelection() - .clearTimeAggregation() - .clearLimit() - .setOffset(0) - .clearOrderBy() - .build(); EntitiesRequestContext entitiesRequestContext = new EntitiesRequestContext( tenantId, startTime, @@ -718,17 +670,6 @@ public void test_visitPaginateOnlyNode() { .putMetricSeries("SUM_API.numCalls", getMockMetricSeries(30, "SUM")) ); - Map totalEntityKeyBuilderResponseMap = Map.of( - EntityKey.of("entity-id-0"), Entity.newBuilder().putAttribute("API.id", getStringValue("entity-0")), - EntityKey.of("entity-id-1"), Entity.newBuilder().putAttribute("API.id", getStringValue("entity-1")), - EntityKey.of("entity-id-2"), Entity.newBuilder().putAttribute("API.id", getStringValue("entity-2")), - EntityKey.of("entity-id-3"), Entity.newBuilder().putAttribute("API.id", getStringValue("entity-3")), - EntityKey.of("entity-id-4"), Entity.newBuilder().putAttribute("API.id", getStringValue("entity-4")) - ); - - EntityFetcherResponse totalEntityFetcherResponse = - new EntityFetcherResponse(totalEntityKeyBuilderResponseMap); - when(executionContext.getEntityIdExpressions()).thenReturn(List.of(buildExpression(API_ID_ATTR))); when(executionContext.getSourceToSelectionExpressionMap()) .thenReturn(Map.of("QS", List.of(selectionExpression))); @@ -769,9 +710,9 @@ public void test_visitPaginateOnlyNode() { when(queryServiceEntityFetcher.getEntities( entitiesRequestContext, entitiesRequestForAttributes)) .thenReturn(attributesResponse); - when(queryServiceEntityFetcher.getEntities( - eq(entitiesRequestContext), eq(totalEntitiesRequest))) - .thenReturn(totalEntityFetcherResponse); + when(queryServiceEntityFetcher.getTotal( + eq(entitiesRequestContext), eq(entitiesRequest))) + .thenReturn(100L); when(queryServiceEntityFetcher.getEntities( entitiesRequestContext, entitiesRequestForMetricAggregation)) .thenReturn(new EntityFetcherResponse(entityKeyBuilderResponseMap2)); @@ -793,9 +734,7 @@ public void test_visitPaginateOnlyNode() { .build(); compareEntityResponses( - new EntityResponse( - new EntityFetcherResponse(expectedEntityKeyBuilderResponseMap), - totalEntityFetcherResponse.getEntityKeyBuilderMap().keySet()), + new EntityResponse(new EntityFetcherResponse(expectedEntityKeyBuilderResponseMap), 100), executionVisitor.visit(selectionNode)); } @@ -816,7 +755,7 @@ public void test_visitSelectionNode_differentSource_callSeparatedCalls() { when(entityDataServiceEntityFetcher.getEntities(any(), any())).thenReturn(result4); when(queryServiceEntityFetcher.getEntities(any(), any())).thenReturn(result4); when(executionVisitor.visit(any(NoOpNode.class))) - .thenReturn(new EntityResponse(result4, result4.getEntityKeyBuilderMap().keySet())); + .thenReturn(new EntityResponse(result4, result4.getEntityKeyBuilderMap().size())); executionVisitor.visit(selectionNode); verify(entityDataServiceEntityFetcher).getEntities(any(), any()); verify(queryServiceEntityFetcher).getEntities(any(), any()); @@ -953,7 +892,7 @@ public void test_visitOnlySelectionsNode_shouldSetTotalEntityKeys() { { compareEntityResponses( new EntityResponse( - new EntityFetcherResponse(entityKeyBuilderResponseMap1), totalEntityKeys), + new EntityFetcherResponse(entityKeyBuilderResponseMap1), 4), executionVisitor.visit(childSelectionNode)); } @@ -962,7 +901,7 @@ public void test_visitOnlySelectionsNode_shouldSetTotalEntityKeys() { { compareEntityResponses( new EntityResponse( - new EntityFetcherResponse(expectedEntityKeyBuilderResponseMap), totalEntityKeys), + new EntityFetcherResponse(expectedEntityKeyBuilderResponseMap), 4), executionVisitor.visit(selectionNode)); } } @@ -990,7 +929,7 @@ public void test_visitSelectionNode_nonEmptyFilter_emptyResult() { EntityResponse response = executionVisitor.visit(selectionNode); Assertions.assertTrue(response.getEntityFetcherResponse().isEmpty()); - Assertions.assertTrue(response.getEntityKeys().isEmpty()); + Assertions.assertEquals(0, response.getTotal()); verify(queryServiceEntityFetcher, never()).getEntities(any(), any()); } diff --git a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/entity/query/visitor/FilterOptimizingVisitorTest.java b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/entity/query/visitor/FilterOptimizingVisitorTest.java index c30a7b53..a02a9a2f 100644 --- a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/entity/query/visitor/FilterOptimizingVisitorTest.java +++ b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/entity/query/visitor/FilterOptimizingVisitorTest.java @@ -9,6 +9,7 @@ import org.hypertrace.gateway.service.entity.query.AndNode; import org.hypertrace.gateway.service.entity.query.DataFetcherNode; +import org.hypertrace.gateway.service.entity.query.OrNode; import org.hypertrace.gateway.service.entity.query.PaginateOnlyNode; import org.hypertrace.gateway.service.entity.query.QueryNode; import org.hypertrace.gateway.service.entity.query.SelectionNode; @@ -105,6 +106,69 @@ public void testAndNodes_dataFetcherPaginationNodes() { assertEquals(orderByExpressions, mergedDataFetcherNode.getOrderByExpressionList()); } + + @Test + public void testAndNodes_dataFetcherNodes_sameSource() { + Filter filter1 = generateEQFilter("API.name", "apiName1"); + Filter filter2 = generateEQFilter("API.id", "apiId1"); + Filter filter3 = generateEQFilter("API.type", "HTTP"); + DataFetcherNode dataFetcherNode1 = + new DataFetcherNode("QS", filter1); + DataFetcherNode dataFetcherNode2 = + new DataFetcherNode("QS", filter2); + DataFetcherNode dataFetcherNode3 = + new DataFetcherNode("QS", filter3); + AndNode andNode = new AndNode(List.of(dataFetcherNode1, dataFetcherNode2)); + AndNode parentAndNode = new AndNode(List.of(andNode, dataFetcherNode3)); + + QueryNode queryNode = parentAndNode.acceptVisitor(new FilterOptimizingVisitor()); + + DataFetcherNode mergedDataFetcherNode = (DataFetcherNode) queryNode; + assertEquals("QS", mergedDataFetcherNode.getSource()); + assertEquals( + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter( + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter(filter1) + .addChildFilter(filter2)) + .addChildFilter(filter3) + .build(), + mergedDataFetcherNode.getFilter()); + } + + @Test + public void testOrNodes_dataFetcherNodes_sameSource() { + Filter filter1 = generateEQFilter("API.name", "apiName1"); + Filter filter2 = generateEQFilter("API.id", "apiId1"); + Filter filter3 = generateEQFilter("API.type", "HTTP"); + DataFetcherNode dataFetcherNode1 = + new DataFetcherNode("QS", filter1); + DataFetcherNode dataFetcherNode2 = + new DataFetcherNode("QS", filter2); + DataFetcherNode dataFetcherNode3 = + new DataFetcherNode("QS", filter3); + OrNode orNode = new OrNode(List.of(dataFetcherNode1, dataFetcherNode2)); + OrNode parentOrNode = new OrNode(List.of(orNode, dataFetcherNode3)); + + QueryNode queryNode = parentOrNode.acceptVisitor(new FilterOptimizingVisitor()); + + DataFetcherNode mergedDataFetcherNode = (DataFetcherNode) queryNode; + assertEquals("QS", mergedDataFetcherNode.getSource()); + assertEquals( + Filter.newBuilder() + .setOperator(Operator.OR) + .addChildFilter( + Filter.newBuilder() + .setOperator(Operator.OR) + .addChildFilter(filter1) + .addChildFilter(filter2)) + .addChildFilter(filter3) + .build(), + mergedDataFetcherNode.getFilter()); + } + @Test public void testAndNodes_dataFetcherNonPaginationNodes() { Filter filter1 = generateEQFilter("API.name", "apiName1");