diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/converters/QueryAndGatewayDtoConverter.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/converters/QueryAndGatewayDtoConverter.java index 9b6f461b..dd4535ff 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/converters/QueryAndGatewayDtoConverter.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/converters/QueryAndGatewayDtoConverter.java @@ -1,10 +1,10 @@ package org.hypertrace.gateway.service.common.converters; import static java.lang.String.format; -import static org.hypertrace.gateway.service.v1.common.FunctionType.DISTINCT_ARRAY; import com.google.common.base.Strings; import java.time.Duration; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Optional; @@ -472,6 +472,27 @@ public static Filter addTimeAndSpaceFiltersAndConvertToQueryFilter( String spacesAttributeId, org.hypertrace.gateway.service.v1.common.Filter providedFilter) { + return addTimeSpaceAndIdFiltersAndConvertToQueryFilter( + startTimeMillis, + endTimeMillis, + spaceId, + Collections.emptyList(), + timestampAttributeId, + spacesAttributeId, + Collections.emptyList(), + providedFilter); + } + + public static Filter addTimeSpaceAndIdFiltersAndConvertToQueryFilter( + long startTimeMillis, + long endTimeMillis, + String spaceId, + List entityIds, + String timestampAttributeId, + String spacesAttributeId, + List entityIdAttributes, + org.hypertrace.gateway.service.v1.common.Filter providedFilter) { + Filter.Builder compositeFilter = Filter.newBuilder().setOperator(Operator.AND); Filter convertedProvidedFilter = isNonDefaultFilter(providedFilter) @@ -493,6 +514,14 @@ public static Filter addTimeAndSpaceFiltersAndConvertToQueryFilter( QueryRequestUtil.createStringFilter(spacesAttributeId, Operator.EQ, spaceId)); } + if (!entityIds.isEmpty()) { + compositeFilter.addChildFilter( + createEntityIdAttributeFilter( + entityIdAttributes, + Operator.IN, + QueryRequestUtil.createStringArrayLiteralExpression(entityIds))); + } + // If only one filter was added, unwrap the one child filter and use that return compositeFilter.getChildFilterCount() == 1 ? compositeFilter.getChildFilter(0) @@ -531,6 +560,14 @@ private static boolean isNonDefaultFilter(Filter filter) { return filter != null && !Filter.getDefaultInstance().equals(filter); } + private static Filter createEntityIdAttributeFilter( + List entityIdAttributes, Operator operator, Expression expression) { + if (entityIdAttributes.size() != 1) { + throw new RuntimeException("Cannot have more than one id attribute for an entity"); + } + return QueryRequestUtil.createFilter(entityIdAttributes.get(0), operator, expression); + } + public static List convertToQueryOrderByExpressions( List gatewayOrderBy) { return gatewayOrderBy.stream() diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/ExploreService.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/ExploreService.java index 1a27d90e..f2bd3c5f 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/ExploreService.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/ExploreService.java @@ -19,9 +19,11 @@ import org.hypertrace.gateway.service.common.ExpressionContext; import org.hypertrace.gateway.service.common.RequestContext; import org.hypertrace.gateway.service.common.config.ScopeFilterConfigs; +import org.hypertrace.gateway.service.common.datafetcher.QueryServiceEntityFetcher; import org.hypertrace.gateway.service.common.util.QueryServiceClient; import org.hypertrace.gateway.service.entity.config.EntityIdColumnsConfigs; import org.hypertrace.gateway.service.explore.entity.EntityRequestHandler; +import org.hypertrace.gateway.service.explore.entity.EntityServiceEntityFetcher; import org.hypertrace.gateway.service.v1.explore.ExploreRequest; import org.hypertrace.gateway.service.v1.explore.ExploreResponse; @@ -46,10 +48,27 @@ public ExploreService( ScopeFilterConfigs scopeFiltersConfig, EntityIdColumnsConfigs entityIdColumnsConfigs, EntityTypesProvider entityTypesProvider) { + QueryServiceEntityFetcher queryServiceEntityFetcher = + new QueryServiceEntityFetcher( + queryServiceClient, attributeMetadataProvider, entityIdColumnsConfigs); + EntityServiceEntityFetcher entityServiceEntityFetcher = + new EntityServiceEntityFetcher( + attributeMetadataProvider, entityIdColumnsConfigs, entityQueryServiceClient); this.attributeMetadataProvider = attributeMetadataProvider; - this.normalRequestHandler = new RequestHandler(queryServiceClient, attributeMetadataProvider); + this.normalRequestHandler = + new RequestHandler( + queryServiceClient, + attributeMetadataProvider, + entityIdColumnsConfigs, + queryServiceEntityFetcher, + entityServiceEntityFetcher); this.timeAggregationsRequestHandler = - new TimeAggregationsRequestHandler(queryServiceClient, attributeMetadataProvider); + new TimeAggregationsRequestHandler( + queryServiceClient, + attributeMetadataProvider, + entityIdColumnsConfigs, + queryServiceEntityFetcher, + entityServiceEntityFetcher); this.timeAggregationsWithGroupByRequestHandler = new TimeAggregationsWithGroupByRequestHandler( attributeMetadataProvider, normalRequestHandler, timeAggregationsRequestHandler); @@ -58,7 +77,8 @@ public ExploreService( attributeMetadataProvider, entityIdColumnsConfigs, queryServiceClient, - entityQueryServiceClient); + queryServiceEntityFetcher, + entityServiceEntityFetcher); this.scopeFilterConfigs = scopeFiltersConfig; this.entityTypesProvider = entityTypesProvider; initMetrics(); diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/RequestHandler.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/RequestHandler.java index 3fc02442..7509aa32 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/RequestHandler.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/RequestHandler.java @@ -5,11 +5,16 @@ import com.google.common.collect.Streams; import com.google.protobuf.InvalidProtocolBufferException; import com.google.protobuf.util.JsonFormat; +import java.util.ArrayList; +import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.Optional; +import java.util.Set; import java.util.stream.Collectors; import org.hypertrace.core.attribute.service.v1.AttributeMetadata; +import org.hypertrace.core.attribute.service.v1.AttributeSource; import org.hypertrace.core.query.service.api.ColumnMetadata; import org.hypertrace.core.query.service.api.Filter; import org.hypertrace.core.query.service.api.QueryRequest; @@ -19,32 +24,65 @@ import org.hypertrace.core.query.service.api.Value; import org.hypertrace.gateway.service.common.AttributeMetadataProvider; import org.hypertrace.gateway.service.common.converters.QueryAndGatewayDtoConverter; +import org.hypertrace.gateway.service.common.datafetcher.EntityFetcherResponse; +import org.hypertrace.gateway.service.common.datafetcher.QueryServiceEntityFetcher; import org.hypertrace.gateway.service.common.util.AttributeMetadataUtil; import org.hypertrace.gateway.service.common.util.ExpressionReader; import org.hypertrace.gateway.service.common.util.MetricAggregationFunctionUtil; import org.hypertrace.gateway.service.common.util.OrderByUtil; +import org.hypertrace.gateway.service.common.util.QueryExpressionUtil; import org.hypertrace.gateway.service.common.util.QueryServiceClient; +import org.hypertrace.gateway.service.entity.EntitiesRequestContext; +import org.hypertrace.gateway.service.entity.config.EntityIdColumnsConfigs; +import org.hypertrace.gateway.service.explore.entity.EntityServiceEntityFetcher; import org.hypertrace.gateway.service.v1.common.Expression; import org.hypertrace.gateway.service.v1.common.FunctionExpression; import org.hypertrace.gateway.service.v1.common.OrderByExpression; import org.hypertrace.gateway.service.v1.common.TimeAggregation; +import org.hypertrace.gateway.service.v1.entity.EntitiesRequest; +import org.hypertrace.gateway.service.v1.entity.Entity.Builder; import org.hypertrace.gateway.service.v1.explore.ExploreRequest; import org.hypertrace.gateway.service.v1.explore.ExploreResponse; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +/** + * {@link RequestHandler} is currently used only when the selections, group bys and filters are on + * QS. Multiple sources are supported but when filters are on EDS only source. + * + *

If there are filters on attributes from EDS source only, a. handler will first query QS source + * with time range filter to get all the possible Ids b. Then it will filter the ids from step (a), + * with EDS filters c. And finally, with entity ids queries from step (b) we will query the + * selections from QS. In this approach we are making three network calls but handling less data in + * memory. + * + *

Other approach could be to first query all the data from QS source (limit of 10000) along with + * selections and then filter out data from EDS filters. In this way we will only send two queries + * but need to handle lots of data in memory. + * + */ public class RequestHandler implements RequestHandlerWithSorting { private static final Logger LOG = LoggerFactory.getLogger(RequestHandler.class); private final QueryServiceClient queryServiceClient; private final AttributeMetadataProvider attributeMetadataProvider; private final TheRestGroupRequestHandler theRestGroupRequestHandler; + private final EntityIdColumnsConfigs entityIdColumnsConfigs; + private final QueryServiceEntityFetcher queryServiceEntityFetcher; + private final EntityServiceEntityFetcher entityServiceEntityFetcher; public RequestHandler( - QueryServiceClient queryServiceClient, AttributeMetadataProvider attributeMetadataProvider) { + QueryServiceClient queryServiceClient, + AttributeMetadataProvider attributeMetadataProvider, + EntityIdColumnsConfigs entityIdColumnsConfigs, + QueryServiceEntityFetcher queryServiceEntityFetcher, + EntityServiceEntityFetcher entityServiceEntityFetcher) { this.queryServiceClient = queryServiceClient; this.attributeMetadataProvider = attributeMetadataProvider; this.theRestGroupRequestHandler = new TheRestGroupRequestHandler(this); + this.entityIdColumnsConfigs = entityIdColumnsConfigs; + this.queryServiceEntityFetcher = queryServiceEntityFetcher; + this.entityServiceEntityFetcher = entityServiceEntityFetcher; } @Override @@ -69,8 +107,18 @@ QueryRequest buildQueryRequest( requestContext.setHasGroupBy(true); } - QueryRequest.Builder builder = QueryRequest.newBuilder(); + List entityIds = new ArrayList<>(); + org.hypertrace.gateway.service.v1.common.Filter qsSourceFilter = request.getFilter(); + Map attributeMetadataMap = + attributeMetadataProvider.getAttributesMetadata(requestContext, request.getContext()); + if (hasOnlyAttributeSource(request.getFilter(), AttributeSource.EDS, attributeMetadataMap)) { + entityIds = getEntityIdsToFilter(requestContext, request, attributeMetadataMap); + qsSourceFilter = + buildFilter(request.getFilter(), AttributeSource.QS, attributeMetadataMap) + .orElse(request.getFilter()); + } + QueryRequest.Builder builder = QueryRequest.newBuilder(); // 1. Add selections. All selections should either be only column or only function, never both. // The validator should catch this. List aggregatedSelections = @@ -92,7 +140,8 @@ QueryRequest buildQueryRequest( // 2. Add filter builder.setFilter( - constructQueryServiceFilter(request, requestContext, attributeMetadataProvider)); + constructQueryServiceFilter( + request, qsSourceFilter, requestContext, attributeMetadataProvider, entityIds)); if (requestContext.hasGroupBy() && request.getIncludeRestGroup() && request.getOffset() > 0) { // including rest group with offset is an invalid combination @@ -114,6 +163,96 @@ QueryRequest buildQueryRequest( return builder.build(); } + // This is to get all the entity Ids for the EDS source filter. + // 1. First filter out entity ids based on the time range from QS filter + // 2. Then filter out entity ids return in 1 based on EDS filter. + private List getEntityIdsToFilter( + ExploreRequestContext context, + ExploreRequest exploreRequest, + Map attributeMetadataMap) { + LOG.debug("Querying entity ids from EDS source {}", exploreRequest); + // Check if there is any filter present with EDS only source. If not then return, + // else query the respective entityIds from the EDS source. + Optional maybeEdsFilter = + buildFilter(exploreRequest.getFilter(), AttributeSource.EDS, attributeMetadataMap); + if (maybeEdsFilter.isEmpty()) { + return Collections.emptyList(); + } + + Set allEntityIds = + this.getEntityIdsInTimeRangeFromQueryService(context, exploreRequest); + ExploreRequest edsExploreRequest = + buildExploreRequest(context, exploreRequest.getContext(), maybeEdsFilter.orElseThrow()); + List resultRows = + this.entityServiceEntityFetcher.getResults(context, edsExploreRequest, allEntityIds); + + return resultRows.stream() + .map(row -> row.getColumnsMap().values().stream().findFirst()) + .filter(Optional::isPresent) + .map(Optional::get) + .map(org.hypertrace.gateway.service.v1.common.Value::getString) + .collect(Collectors.toUnmodifiableList()); + } + + private ExploreRequest buildExploreRequest( + ExploreRequestContext exploreRequestContext, + String context, + org.hypertrace.gateway.service.v1.common.Filter edsFilter) { + List entityIdAttributeIds = + AttributeMetadataUtil.getIdAttributeIds( + attributeMetadataProvider, entityIdColumnsConfigs, exploreRequestContext, context); + List groupBySelections = + entityIdAttributeIds.stream() + .map(attributeId -> QueryExpressionUtil.buildAttributeExpression(attributeId).build()) + .collect(Collectors.toUnmodifiableList()); + + return ExploreRequest.newBuilder() + .setContext(context) + .setFilter(edsFilter) + .addAllGroupBy(groupBySelections) + .build(); + } + + protected Set getEntityIdsInTimeRangeFromQueryService( + ExploreRequestContext requestContext, ExploreRequest exploreRequest) { + EntitiesRequestContext entitiesRequestContext = + convert(attributeMetadataProvider, requestContext); + Map attributeMetadataMap = + attributeMetadataProvider.getAttributesMetadata( + requestContext, exploreRequest.getContext()); + + EntitiesRequest.Builder entitiesRequest = + EntitiesRequest.newBuilder() + .setEntityType(exploreRequest.getContext()) + .setStartTimeMillis(exploreRequest.getStartTimeMillis()) + .setEndTimeMillis(exploreRequest.getEndTimeMillis()); + + Optional maybeQsFilters = + buildFilter(exploreRequest.getFilter(), AttributeSource.QS, attributeMetadataMap); + maybeQsFilters.ifPresent(entitiesRequest::setFilter); + + EntityFetcherResponse response = + queryServiceEntityFetcher.getEntities(entitiesRequestContext, entitiesRequest.build()); + return response.getEntityKeyBuilderMap().values().stream() + .map(Builder::getId) + .collect(Collectors.toUnmodifiableSet()); + } + + private EntitiesRequestContext convert( + AttributeMetadataProvider attributeMetadataProvider, ExploreRequestContext requestContext) { + String entityType = requestContext.getContext(); + + String timestampAttributeId = + AttributeMetadataUtil.getTimestampAttributeId( + attributeMetadataProvider, requestContext, entityType); + return new EntitiesRequestContext( + requestContext.getGrpcContext(), + requestContext.getStartTimeMillis(), + requestContext.getEndTimeMillis(), + entityType, + timestampAttributeId); + } + private Iterator executeQuery( ExploreRequestContext context, QueryRequest queryRequest) { if (getLogger().isDebugEnabled()) { @@ -137,15 +276,35 @@ Filter constructQueryServiceFilter( ExploreRequest request, ExploreRequestContext exploreRequestContext, AttributeMetadataProvider attributeMetadataProvider) { - return QueryAndGatewayDtoConverter.addTimeAndSpaceFiltersAndConvertToQueryFilter( + return this.constructQueryServiceFilter( + request, + request.getFilter(), + exploreRequestContext, + attributeMetadataProvider, + Collections.emptyList()); + } + + Filter constructQueryServiceFilter( + ExploreRequest request, + org.hypertrace.gateway.service.v1.common.Filter requestFilter, + ExploreRequestContext exploreRequestContext, + AttributeMetadataProvider attributeMetadataProvider, + List entityIds) { + return QueryAndGatewayDtoConverter.addTimeSpaceAndIdFiltersAndConvertToQueryFilter( request.getStartTimeMillis(), request.getEndTimeMillis(), request.getSpaceId(), + entityIds, AttributeMetadataUtil.getTimestampAttributeId( attributeMetadataProvider, exploreRequestContext, request.getContext()), AttributeMetadataUtil.getSpaceAttributeId( attributeMetadataProvider, exploreRequestContext, request.getContext()), - request.getFilter()); + AttributeMetadataUtil.getIdAttributeIds( + attributeMetadataProvider, + entityIdColumnsConfigs, + exploreRequestContext, + request.getContext()), + requestFilter); } void addGroupByExpressions(QueryRequest.Builder builder, ExploreRequest request) { @@ -369,4 +528,85 @@ private Map remapAttributeMetadataByResultName( .collect(Collectors.toUnmodifiableList()), attributeMetadataByIdMap); } + + private boolean hasOnlyAttributeSource( + org.hypertrace.gateway.service.v1.common.Filter filter, + AttributeSource source, + Map attributeMetadataMap) { + if (filter.equals(org.hypertrace.gateway.service.v1.common.Filter.getDefaultInstance())) { + return false; + } + + org.hypertrace.gateway.service.v1.common.Operator operator = filter.getOperator(); + switch (operator) { + case UNDEFINED: + return false; + case AND: + case OR: + for (org.hypertrace.gateway.service.v1.common.Filter childFilter : + filter.getChildFilterList()) { + if (hasOnlyAttributeSource(childFilter, source, attributeMetadataMap)) { + return true; + } + } + return false; + default: + List availableSources = + attributeMetadataMap + .get( + ExpressionReader.getAttributeIdFromAttributeSelection(filter.getLhs()) + .orElseThrow()) + .getSourcesList(); + return availableSources.size() == 1 && availableSources.contains(source); + } + } + + private Optional buildFilter( + org.hypertrace.gateway.service.v1.common.Filter filter, + AttributeSource source, + Map attributeMetadataMap) { + if (filter.equals(org.hypertrace.gateway.service.v1.common.Filter.getDefaultInstance())) { + return Optional.empty(); + } + + org.hypertrace.gateway.service.v1.common.Operator operator = filter.getOperator(); + switch (operator) { + case UNDEFINED: + return Optional.empty(); + case AND: + case OR: + return buildCompositeFilter(filter, source, operator, attributeMetadataMap); + default: + List availableSources = + attributeMetadataMap + .get( + ExpressionReader.getAttributeIdFromAttributeSelection(filter.getLhs()) + .orElseThrow()) + .getSourcesList(); + return availableSources.contains(source) + ? Optional.of( + org.hypertrace.gateway.service.v1.common.Filter.newBuilder(filter).build()) + : Optional.empty(); + } + } + + private Optional buildCompositeFilter( + org.hypertrace.gateway.service.v1.common.Filter filter, + AttributeSource source, + org.hypertrace.gateway.service.v1.common.Operator operator, + Map attributeMetadataMap) { + org.hypertrace.gateway.service.v1.common.Filter.Builder filterBuilder = + org.hypertrace.gateway.service.v1.common.Filter.newBuilder(); + for (org.hypertrace.gateway.service.v1.common.Filter childFilter : + filter.getChildFilterList()) { + buildFilter(childFilter, source, attributeMetadataMap) + .ifPresent(filterBuilder::addChildFilter); + } + if (filterBuilder.getChildFilterCount() > 0) { + filterBuilder.setOperator(operator); + return Optional.of(filterBuilder.build()); + } else { + return Optional.empty(); + } + } } diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/TimeAggregationsRequestHandler.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/TimeAggregationsRequestHandler.java index 022c98e5..4dfe6484 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/TimeAggregationsRequestHandler.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/TimeAggregationsRequestHandler.java @@ -15,10 +15,13 @@ import org.hypertrace.core.query.service.api.Value; import org.hypertrace.gateway.service.common.AttributeMetadataProvider; import org.hypertrace.gateway.service.common.converters.QueryAndGatewayDtoConverter; +import org.hypertrace.gateway.service.common.datafetcher.QueryServiceEntityFetcher; import org.hypertrace.gateway.service.common.util.AttributeMetadataUtil; import org.hypertrace.gateway.service.common.util.ExpressionReader; import org.hypertrace.gateway.service.common.util.QueryExpressionUtil; import org.hypertrace.gateway.service.common.util.QueryServiceClient; +import org.hypertrace.gateway.service.entity.config.EntityIdColumnsConfigs; +import org.hypertrace.gateway.service.explore.entity.EntityServiceEntityFetcher; import org.hypertrace.gateway.service.v1.common.OrderByExpression; import org.hypertrace.gateway.service.v1.common.Period; import org.hypertrace.gateway.service.v1.common.SortOrder; @@ -34,8 +37,17 @@ public class TimeAggregationsRequestHandler extends RequestHandler { private static final Logger LOG = LoggerFactory.getLogger(TimeAggregationsRequestHandler.class); TimeAggregationsRequestHandler( - QueryServiceClient queryServiceClient, AttributeMetadataProvider attributeMetadataProvider) { - super(queryServiceClient, attributeMetadataProvider); + QueryServiceClient queryServiceClient, + AttributeMetadataProvider attributeMetadataProvider, + EntityIdColumnsConfigs entityIdColumnsConfigs, + QueryServiceEntityFetcher queryServiceEntityFetcher, + EntityServiceEntityFetcher entityServiceEntityFetcher) { + super( + queryServiceClient, + attributeMetadataProvider, + entityIdColumnsConfigs, + queryServiceEntityFetcher, + entityServiceEntityFetcher); } @Override diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/TimeAggregationsWithGroupByRequestHandler.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/TimeAggregationsWithGroupByRequestHandler.java index 2c988008..3c11f2ce 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/TimeAggregationsWithGroupByRequestHandler.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/TimeAggregationsWithGroupByRequestHandler.java @@ -11,7 +11,6 @@ import org.hypertrace.gateway.service.common.AttributeMetadataProvider; import org.hypertrace.gateway.service.common.util.AttributeMetadataUtil; import org.hypertrace.gateway.service.common.util.ExpressionReader; -import org.hypertrace.gateway.service.common.util.QueryServiceClient; import org.hypertrace.gateway.service.v1.common.Expression; import org.hypertrace.gateway.service.v1.common.Filter; import org.hypertrace.gateway.service.v1.common.LiteralConstant; @@ -36,14 +35,6 @@ public class TimeAggregationsWithGroupByRequestHandler implements IRequestHandle 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); - } - TimeAggregationsWithGroupByRequestHandler( AttributeMetadataProvider attributeMetadataProvider, RequestHandler normalRequestHandler, diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/entity/EntityRequestHandler.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/entity/EntityRequestHandler.java index 07806382..5c0f8594 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/entity/EntityRequestHandler.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/entity/EntityRequestHandler.java @@ -1,41 +1,18 @@ package org.hypertrace.gateway.service.explore.entity; -import static java.util.Collections.unmodifiableSet; -import static java.util.stream.Collectors.toUnmodifiableSet; - -import com.google.common.annotations.VisibleForTesting; -import com.google.common.collect.Streams; import java.util.HashSet; -import java.util.Iterator; import java.util.List; -import java.util.Map; import java.util.Optional; import java.util.Set; -import java.util.stream.Collectors; -import org.hypertrace.core.attribute.service.v1.AttributeMetadata; -import org.hypertrace.entity.query.service.client.EntityQueryServiceClient; -import org.hypertrace.entity.query.service.v1.ColumnMetadata; -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.Value; import org.hypertrace.gateway.service.common.AttributeMetadataProvider; -import org.hypertrace.gateway.service.common.converters.EntityServiceAndGatewayServiceConverter; -import org.hypertrace.gateway.service.common.datafetcher.EntityFetcherResponse; import org.hypertrace.gateway.service.common.datafetcher.QueryServiceEntityFetcher; -import org.hypertrace.gateway.service.common.util.AttributeMetadataUtil; import org.hypertrace.gateway.service.common.util.DataCollectionUtil; -import org.hypertrace.gateway.service.common.util.MetricAggregationFunctionUtil; import org.hypertrace.gateway.service.common.util.QueryServiceClient; -import org.hypertrace.gateway.service.entity.EntitiesRequestContext; import org.hypertrace.gateway.service.entity.config.EntityIdColumnsConfigs; import org.hypertrace.gateway.service.explore.ExploreRequestContext; import org.hypertrace.gateway.service.explore.RequestHandler; import org.hypertrace.gateway.service.explore.RowComparator; -import org.hypertrace.gateway.service.v1.common.FunctionExpression; import org.hypertrace.gateway.service.v1.common.OrderByExpression; -import org.hypertrace.gateway.service.v1.entity.EntitiesRequest; -import org.hypertrace.gateway.service.v1.entity.Entity.Builder; import org.hypertrace.gateway.service.v1.explore.EntityOption; import org.hypertrace.gateway.service.v1.explore.ExploreRequest; import org.hypertrace.gateway.service.v1.explore.ExploreResponse; @@ -57,38 +34,25 @@ * */ public class EntityRequestHandler extends RequestHandler { + private final EntityServiceEntityFetcher entityServiceEntityFetcher; private final AttributeMetadataProvider attributeMetadataProvider; - private final QueryServiceEntityFetcher queryServiceEntityFetcher; - private final EntityServiceEntityFetcher entityServiceEntityFetcher; public EntityRequestHandler( AttributeMetadataProvider attributeMetadataProvider, EntityIdColumnsConfigs entityIdColumnsConfigs, QueryServiceClient queryServiceClient, - EntityQueryServiceClient entityQueryServiceClient) { - super(queryServiceClient, attributeMetadataProvider); - - this.attributeMetadataProvider = attributeMetadataProvider; - this.queryServiceEntityFetcher = - new QueryServiceEntityFetcher( - queryServiceClient, attributeMetadataProvider, entityIdColumnsConfigs); - this.entityServiceEntityFetcher = - new EntityServiceEntityFetcher( - attributeMetadataProvider, entityIdColumnsConfigs, entityQueryServiceClient); - } - - @VisibleForTesting - public EntityRequestHandler( - AttributeMetadataProvider attributeMetadataProvider, - QueryServiceClient queryServiceClient, QueryServiceEntityFetcher queryServiceEntityFetcher, EntityServiceEntityFetcher entityServiceEntityFetcher) { - super(queryServiceClient, attributeMetadataProvider); - + super( + queryServiceClient, + attributeMetadataProvider, + entityIdColumnsConfigs, + queryServiceEntityFetcher, + entityServiceEntityFetcher); this.attributeMetadataProvider = attributeMetadataProvider; - this.queryServiceEntityFetcher = queryServiceEntityFetcher; this.entityServiceEntityFetcher = entityServiceEntityFetcher; + this.queryServiceEntityFetcher = queryServiceEntityFetcher; } @Override @@ -104,40 +68,14 @@ public ExploreResponse.Builder handleRequest( Set entityIds = new HashSet<>(); Optional maybeEntityOption = getEntityOption(exploreRequest); if (requestOnLiveEntities(maybeEntityOption)) { - entityIds.addAll(getEntityIdsFromQueryService(requestContext, exploreRequest)); + entityIds.addAll(getEntityIdsInTimeRangeFromQueryService(requestContext, exploreRequest)); if (entityIds.isEmpty()) { return builder; } } - Iterator resultSetChunkIterator = - entityServiceEntityFetcher.getResults( - requestContext, exploreRequest, unmodifiableSet(entityIds)); - - while (resultSetChunkIterator.hasNext()) { - org.hypertrace.entity.query.service.v1.ResultSetChunk chunk = resultSetChunkIterator.next(); - getLogger().debug("Received chunk: {}", chunk); - - if (chunk.getRowCount() < 1) { - break; - } - - if (!chunk.hasResultSetMetadata()) { - getLogger().warn("Chunk doesn't have result metadata so couldn't process the response."); - break; - } - - chunk - .getRowList() - .forEach( - row -> - handleRow( - row, - chunk.getResultSetMetadata(), - builder, - requestContext, - attributeMetadataProvider)); - } + builder.addAllRow( + entityServiceEntityFetcher.getResults(requestContext, exploreRequest, entityIds)); // If there's a Group By in the request, we need to do the sorting and pagination ourselves. if (requestContext.hasGroupBy()) { @@ -183,113 +121,6 @@ protected List sortAndPagi rowBuilders.stream(), limit, offset, orderByExpressions.size(), rowComparator); } - private Set getEntityIdsFromQueryService( - ExploreRequestContext requestContext, ExploreRequest exploreRequest) { - EntitiesRequestContext entitiesRequestContext = - convert(attributeMetadataProvider, requestContext); - EntitiesRequest entitiesRequest = - EntitiesRequest.newBuilder() - .setEntityType(exploreRequest.getContext()) - .setStartTimeMillis(exploreRequest.getStartTimeMillis()) - .setEndTimeMillis(exploreRequest.getEndTimeMillis()) - .build(); - EntityFetcherResponse response = - queryServiceEntityFetcher.getEntities(entitiesRequestContext, entitiesRequest); - return response.getEntityKeyBuilderMap().values().stream() - .map(Builder::getId) - .collect(toUnmodifiableSet()); - } - - private EntitiesRequestContext convert( - AttributeMetadataProvider attributeMetadataProvider, ExploreRequestContext requestContext) { - String entityType = requestContext.getContext(); - - String timestampAttributeId = - AttributeMetadataUtil.getTimestampAttributeId( - attributeMetadataProvider, requestContext, entityType); - return new EntitiesRequestContext( - requestContext.getGrpcContext(), - requestContext.getStartTimeMillis(), - requestContext.getEndTimeMillis(), - entityType, - timestampAttributeId); - } - - private void handleRow( - Row row, - ResultSetMetadata resultSetMetadata, - ExploreResponse.Builder builder, - ExploreRequestContext requestContext, - AttributeMetadataProvider attributeMetadataProvider) { - var rowBuilder = org.hypertrace.gateway.service.v1.common.Row.newBuilder(); - for (int i = 0; i < resultSetMetadata.getColumnMetadataCount(); i++) { - handleColumn( - row.getColumn(i), - resultSetMetadata.getColumnMetadata(i), - rowBuilder, - requestContext, - attributeMetadataProvider); - } - builder.addRow(rowBuilder); - } - - private void handleColumn( - Value value, - ColumnMetadata metadata, - org.hypertrace.gateway.service.v1.common.Row.Builder rowBuilder, - ExploreRequestContext requestContext, - AttributeMetadataProvider attributeMetadataProvider) { - FunctionExpression function = - requestContext.getFunctionExpressionByAlias(metadata.getColumnName()); - handleColumn(value, metadata, rowBuilder, requestContext, attributeMetadataProvider, function); - } - - void handleColumn( - Value value, - ColumnMetadata metadata, - org.hypertrace.gateway.service.v1.common.Row.Builder rowBuilder, - ExploreRequestContext requestContext, - AttributeMetadataProvider attributeMetadataProvider, - FunctionExpression function) { - Map attributeMetadataMap = - attributeMetadataProvider.getAttributesMetadata( - requestContext, requestContext.getContext()); - Map resultKeyToAttributeMetadataMap = - this.remapAttributeMetadataByResultName( - requestContext.getExploreRequest(), attributeMetadataMap); - org.hypertrace.gateway.service.v1.common.Value gwValue; - if (function != null) { - // Function expression value - gwValue = - EntityServiceAndGatewayServiceConverter.convertToGatewayValueForMetricValue( - MetricAggregationFunctionUtil.getValueTypeForFunctionType( - function, attributeMetadataMap), - resultKeyToAttributeMetadataMap, - metadata, - value); - } else { - // Simple columnId expression value eg. groupBy columns or column selections - gwValue = - EntityServiceAndGatewayServiceConverter.convertToGatewayValue( - metadata.getColumnName(), value, resultKeyToAttributeMetadataMap); - } - - rowBuilder.putColumns(metadata.getColumnName(), gwValue); - } - - private Map remapAttributeMetadataByResultName( - ExploreRequest request, Map attributeMetadataByIdMap) { - return AttributeMetadataUtil.remapAttributeMetadataByResultKey( - Streams.concat( - request.getSelectionList().stream(), - // Add groupBy to Selection list. - // The expectation from the Gateway service client is that they do not add the group - // by expressions to the selection expressions in the request - request.getGroupByList().stream()) - .collect(Collectors.toUnmodifiableList()), - attributeMetadataByIdMap); - } - private boolean requestOnLiveEntities(Optional entityOption) { if (entityOption.isEmpty()) { return true; diff --git a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/entity/EntityServiceEntityFetcher.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/entity/EntityServiceEntityFetcher.java index 605291c9..1bbeb0c8 100644 --- a/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/entity/EntityServiceEntityFetcher.java +++ b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/entity/EntityServiceEntityFetcher.java @@ -1,10 +1,15 @@ package org.hypertrace.gateway.service.explore.entity; +import com.google.common.collect.Streams; +import java.util.ArrayList; import java.util.Iterator; import java.util.List; +import java.util.Map; import java.util.Set; import java.util.stream.Collectors; +import org.hypertrace.core.attribute.service.v1.AttributeMetadata; import org.hypertrace.entity.query.service.client.EntityQueryServiceClient; +import org.hypertrace.entity.query.service.v1.ColumnMetadata; import org.hypertrace.entity.query.service.v1.EntityQueryRequest; import org.hypertrace.entity.query.service.v1.Expression; import org.hypertrace.entity.query.service.v1.Filter; @@ -12,17 +17,24 @@ import org.hypertrace.entity.query.service.v1.LiteralConstant; import org.hypertrace.entity.query.service.v1.Operator; import org.hypertrace.entity.query.service.v1.ResultSetChunk; +import org.hypertrace.entity.query.service.v1.ResultSetMetadata; import org.hypertrace.entity.query.service.v1.Value; import org.hypertrace.entity.query.service.v1.ValueType; import org.hypertrace.gateway.service.common.AttributeMetadataProvider; import org.hypertrace.gateway.service.common.converters.EntityServiceAndGatewayServiceConverter; import org.hypertrace.gateway.service.common.util.AttributeMetadataUtil; import org.hypertrace.gateway.service.common.util.ExpressionReader; +import org.hypertrace.gateway.service.common.util.MetricAggregationFunctionUtil; import org.hypertrace.gateway.service.entity.config.EntityIdColumnsConfigs; import org.hypertrace.gateway.service.explore.ExploreRequestContext; +import org.hypertrace.gateway.service.v1.common.FunctionExpression; +import org.hypertrace.gateway.service.v1.common.Row; import org.hypertrace.gateway.service.v1.explore.ExploreRequest; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class EntityServiceEntityFetcher { + private static final Logger LOGGER = LoggerFactory.getLogger(EntityServiceEntityFetcher.class); private static final int DEFAULT_ENTITY_REQUEST_LIMIT = 10_000; private final AttributeMetadataProvider attributeMetadataProvider; private final EntityIdColumnsConfigs entityIdColumnsConfigs; @@ -37,10 +49,110 @@ public EntityServiceEntityFetcher( this.entityQueryServiceClient = entityQueryServiceClient; } - public Iterator getResults( + public List getResults( ExploreRequestContext requestContext, ExploreRequest exploreRequest, Set entityIds) { EntityQueryRequest request = buildRequest(requestContext, exploreRequest, entityIds); - return entityQueryServiceClient.execute(request, requestContext.getHeaders()); + Iterator result = + entityQueryServiceClient.execute(request, requestContext.getHeaders()); + return readChunkResults(requestContext, result); + } + + protected List readChunkResults( + ExploreRequestContext requestContext, Iterator resultSetChunkIterator) { + List resultRows = new ArrayList<>(); + while (resultSetChunkIterator.hasNext()) { + ResultSetChunk chunk = resultSetChunkIterator.next(); + LOGGER.debug("Received chunk: {}", chunk); + + if (chunk.getRowCount() < 1) { + break; + } + + if (!chunk.hasResultSetMetadata()) { + LOGGER.warn("Chunk doesn't have result metadata so couldn't process the response."); + break; + } + + resultRows.addAll( + chunk.getRowList().stream() + .map( + row -> + handleRow( + row, + chunk.getResultSetMetadata(), + requestContext, + attributeMetadataProvider)) + .collect(Collectors.toUnmodifiableList())); + } + return resultRows; + } + + private Row handleRow( + org.hypertrace.entity.query.service.v1.Row row, + ResultSetMetadata resultSetMetadata, + ExploreRequestContext requestContext, + AttributeMetadataProvider attributeMetadataProvider) { + var rowBuilder = org.hypertrace.gateway.service.v1.common.Row.newBuilder(); + for (int i = 0; i < resultSetMetadata.getColumnMetadataCount(); i++) { + org.hypertrace.entity.query.service.v1.ColumnMetadata metadata = + resultSetMetadata.getColumnMetadata(i); + FunctionExpression function = + requestContext.getFunctionExpressionByAlias(metadata.getColumnName()); + handleColumn( + row.getColumn(i), + metadata, + rowBuilder, + requestContext, + attributeMetadataProvider, + function); + } + return rowBuilder.build(); + } + + private void handleColumn( + Value value, + ColumnMetadata metadata, + Row.Builder rowBuilder, + ExploreRequestContext requestContext, + AttributeMetadataProvider attributeMetadataProvider, + FunctionExpression function) { + Map attributeMetadataMap = + attributeMetadataProvider.getAttributesMetadata( + requestContext, requestContext.getContext()); + Map resultKeyToAttributeMetadataMap = + this.remapAttributeMetadataByResultName( + requestContext.getExploreRequest(), attributeMetadataMap); + org.hypertrace.gateway.service.v1.common.Value gwValue; + if (function != null) { + // Function expression value + gwValue = + EntityServiceAndGatewayServiceConverter.convertToGatewayValueForMetricValue( + MetricAggregationFunctionUtil.getValueTypeForFunctionType( + function, attributeMetadataMap), + resultKeyToAttributeMetadataMap, + metadata, + value); + } else { + // Simple columnId expression value eg. groupBy columns or column selections + gwValue = + EntityServiceAndGatewayServiceConverter.convertToGatewayValue( + metadata.getColumnName(), value, resultKeyToAttributeMetadataMap); + } + + rowBuilder.putColumns(metadata.getColumnName(), gwValue); + } + + private Map remapAttributeMetadataByResultName( + ExploreRequest request, Map attributeMetadataByIdMap) { + return AttributeMetadataUtil.remapAttributeMetadataByResultKey( + Streams.concat( + request.getSelectionList().stream(), + // Add groupBy to Selection list. + // The expectation from the Gateway service client is that they do not add the group + // by expressions to the selection expressions in the request + request.getGroupByList().stream()) + .collect(Collectors.toUnmodifiableList()), + attributeMetadataByIdMap); } private EntityQueryRequest buildRequest( diff --git a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/AbstractServiceTest.java b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/AbstractServiceTest.java index 076719a7..2b510586 100644 --- a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/AbstractServiceTest.java +++ b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/AbstractServiceTest.java @@ -19,6 +19,7 @@ import java.io.Reader; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; @@ -32,9 +33,11 @@ import org.hypertrace.core.attribute.service.v1.AttributeMetadataFilter; import org.hypertrace.core.query.service.api.QueryRequest; import org.hypertrace.core.query.service.api.ResultSetChunk; +import org.hypertrace.entity.query.service.client.EntityQueryServiceClient; import org.hypertrace.gateway.service.EntityTypesProvider; import org.hypertrace.gateway.service.common.config.ScopeFilterConfigs; import org.hypertrace.gateway.service.common.util.QueryServiceClient; +import org.hypertrace.gateway.service.entity.config.EntityIdColumnsConfigs; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.params.ParameterizedTest; @@ -58,6 +61,7 @@ public abstract class AbstractServiceTest< private static List attributeMetadataList; private static ScopeFilterConfigs scopeFilterConfigs; private static EntityTypesProvider entityTypesProvider; + private static EntityIdColumnsConfigs entityIdColumnsConfigs; @BeforeAll public static void setUp() throws IOException { @@ -78,6 +82,7 @@ public static void setUp() throws IOException { + "]"; Config config = ConfigFactory.parseString(scopeFiltersConfig); scopeFilterConfigs = new ScopeFilterConfigs(config); + entityIdColumnsConfigs = new EntityIdColumnsConfigs(Collections.emptyMap()); entityTypesProvider = mock(EntityTypesProvider.class); } @@ -152,13 +157,16 @@ protected static Stream getTestFileNames(String suiteName) { @MethodSource("data") public void runTest(String fileName) throws IOException { QueryServiceClient queryServiceClient = createMockQueryServiceClient(fileName); + EntityQueryServiceClient entityQueryServiceClient = mock(EntityQueryServiceClient.class); TGatewayServiceRequestType testRequest = readGatewayServiceRequest(fileName); TGatewayServiceResponseType actualResponse = executeApi( testRequest, queryServiceClient, + entityQueryServiceClient, attributeMetadataProvider, scopeFilterConfigs, + entityIdColumnsConfigs, entityTypesProvider); TGatewayServiceResponseType expectedResponse = readGatewayServiceResponse(fileName); @@ -278,7 +286,9 @@ private String createResourceFileName(String filesBaseDir, String fileName) { protected abstract TGatewayServiceResponseType executeApi( TGatewayServiceRequestType request, QueryServiceClient queryServiceClient, + EntityQueryServiceClient entityQueryServiceClient, AttributeMetadataProvider attributeMetadataProvider, ScopeFilterConfigs scopeFilterConfigs, + EntityIdColumnsConfigs entityIdColumnsConfigs, EntityTypesProvider entityTypesProvider); } diff --git a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/ExploreServiceTest.java b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/ExploreServiceTest.java index 42e29bec..70ff7734 100644 --- a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/ExploreServiceTest.java +++ b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/ExploreServiceTest.java @@ -2,12 +2,14 @@ import com.google.protobuf.GeneratedMessageV3; import java.util.stream.Stream; +import org.hypertrace.entity.query.service.client.EntityQueryServiceClient; import org.hypertrace.gateway.service.EntityTypesProvider; import org.hypertrace.gateway.service.common.AbstractServiceTest; import org.hypertrace.gateway.service.common.AttributeMetadataProvider; import org.hypertrace.gateway.service.common.RequestContext; import org.hypertrace.gateway.service.common.config.ScopeFilterConfigs; import org.hypertrace.gateway.service.common.util.QueryServiceClient; +import org.hypertrace.gateway.service.entity.config.EntityIdColumnsConfigs; import org.hypertrace.gateway.service.v1.explore.ExploreRequest; import org.hypertrace.gateway.service.v1.explore.ExploreResponse; @@ -37,16 +39,18 @@ protected GeneratedMessageV3.Builder getGatewayServiceResponseBuilder() { protected ExploreResponse executeApi( ExploreRequest request, QueryServiceClient queryServiceClient, + EntityQueryServiceClient entityQueryServiceClient, AttributeMetadataProvider attributeMetadataProvider, ScopeFilterConfigs scopeFilterConfigs, + EntityIdColumnsConfigs entityIdColumnsConfigs, EntityTypesProvider entityTypesProvider) { ExploreService exploreService = new ExploreService( queryServiceClient, - null, + entityQueryServiceClient, attributeMetadataProvider, scopeFilterConfigs, - null, + entityIdColumnsConfigs, entityTypesProvider); return exploreService.explore( new RequestContext( diff --git a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/RequestHandlerTest.java b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/RequestHandlerTest.java index 08bd314e..eea6ce7d 100644 --- a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/RequestHandlerTest.java +++ b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/RequestHandlerTest.java @@ -1,22 +1,49 @@ package org.hypertrace.gateway.service.explore; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import com.google.protobuf.InvalidProtocolBufferException; +import com.google.protobuf.util.JsonFormat; import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import org.hypertrace.core.attribute.service.v1.AttributeMetadata; +import org.hypertrace.core.attribute.service.v1.AttributeSource; +import org.hypertrace.core.grpcutils.context.RequestContext; +import org.hypertrace.core.query.service.api.ColumnMetadata; +import org.hypertrace.core.query.service.api.QueryRequest; +import org.hypertrace.core.query.service.api.ResultSetChunk; +import org.hypertrace.core.query.service.api.ResultSetMetadata; import org.hypertrace.gateway.service.common.AttributeMetadataProvider; +import org.hypertrace.gateway.service.common.datafetcher.EntityFetcherResponse; +import org.hypertrace.gateway.service.common.datafetcher.QueryServiceEntityFetcher; +import org.hypertrace.gateway.service.common.util.ExpressionReader; import org.hypertrace.gateway.service.common.util.QueryServiceClient; +import org.hypertrace.gateway.service.entity.EntityKey; +import org.hypertrace.gateway.service.entity.config.EntityIdColumnsConfigs; +import org.hypertrace.gateway.service.explore.entity.EntityServiceEntityFetcher; +import org.hypertrace.gateway.service.v1.common.AttributeExpression; import org.hypertrace.gateway.service.v1.common.ColumnIdentifier; import org.hypertrace.gateway.service.v1.common.Expression; +import org.hypertrace.gateway.service.v1.common.Filter; import org.hypertrace.gateway.service.v1.common.FunctionExpression; import org.hypertrace.gateway.service.v1.common.FunctionType; 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.Period; +import org.hypertrace.gateway.service.v1.common.Row; import org.hypertrace.gateway.service.v1.common.SortOrder; import org.hypertrace.gateway.service.v1.common.TimeAggregation; import org.hypertrace.gateway.service.v1.common.Value; import org.hypertrace.gateway.service.v1.common.ValueType; +import org.hypertrace.gateway.service.v1.entity.Entity; import org.hypertrace.gateway.service.v1.explore.ExploreRequest; +import org.hypertrace.gateway.service.v1.explore.ExploreResponse.Builder; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -74,7 +101,12 @@ public void orderByExpressionsWithFunction_shouldMatchCorrespondingSelections() .build(); RequestHandler requestHandler = - new RequestHandler(mock(QueryServiceClient.class), mock(AttributeMetadataProvider.class)); + new RequestHandler( + mock(QueryServiceClient.class), + mock(AttributeMetadataProvider.class), + mock(EntityIdColumnsConfigs.class), + mock(QueryServiceEntityFetcher.class), + mock(EntityServiceEntityFetcher.class)); List orderByExpressions = requestHandler.getRequestOrderByExpressions(exploreRequest); @@ -150,7 +182,12 @@ public void noChangeIfOrderByExpressionHasSameAliasAsSelection() { .build(); RequestHandler requestHandler = - new RequestHandler(mock(QueryServiceClient.class), mock(AttributeMetadataProvider.class)); + new RequestHandler( + mock(QueryServiceClient.class), + mock(AttributeMetadataProvider.class), + mock(EntityIdColumnsConfigs.class), + mock(QueryServiceEntityFetcher.class), + mock(EntityServiceEntityFetcher.class)); List orderByExpressions = requestHandler.getRequestOrderByExpressions(exploreRequest); @@ -226,7 +263,12 @@ public void orderByExpressionsWithFunction_shouldMatchCorrespondingTimeAggregati .build(); RequestHandler requestHandler = - new RequestHandler(mock(QueryServiceClient.class), mock(AttributeMetadataProvider.class)); + new RequestHandler( + mock(QueryServiceClient.class), + mock(AttributeMetadataProvider.class), + mock(EntityIdColumnsConfigs.class), + mock(QueryServiceEntityFetcher.class), + mock(EntityServiceEntityFetcher.class)); List orderByExpressions = requestHandler.getRequestOrderByExpressions(exploreRequest); @@ -297,7 +339,12 @@ public void noChangeIfOrderByExpressionIsAColumnAndOrderExpressionOrderIsMaintai .build(); RequestHandler requestHandler = - new RequestHandler(mock(QueryServiceClient.class), mock(AttributeMetadataProvider.class)); + new RequestHandler( + mock(QueryServiceClient.class), + mock(AttributeMetadataProvider.class), + mock(EntityIdColumnsConfigs.class), + mock(QueryServiceEntityFetcher.class), + mock(EntityServiceEntityFetcher.class)); List orderByExpressions = requestHandler.getRequestOrderByExpressions(exploreRequest); @@ -329,4 +376,441 @@ public void noChangeIfOrderByExpressionIsAColumnAndOrderExpressionOrderIsMaintai .build(), orderByExpressions.get(1)); } + + @Test + public void testExploreQueryWithEDSFilter() throws InvalidProtocolBufferException { + ExploreRequest exploreRequest = + ExploreRequest.newBuilder() + .setContext("API") + .setFilter( + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter( + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter( + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter( + Filter.newBuilder() + .setLhs( + Expression.newBuilder() + .setAttributeExpression( + AttributeExpression.newBuilder() + .setAttributeId("API.attributeId3") + .build()) + .build()) + .setOperator(Operator.EQ) + .setRhs( + Expression.newBuilder() + .setLiteral( + LiteralConstant.newBuilder() + .setValue( + Value.newBuilder() + .setValueType(ValueType.STRING) + .setString("value") + .build()) + .build()) + .build()) + .build()) + .addChildFilter( + Filter.newBuilder() + .setLhs( + Expression.newBuilder() + .setAttributeExpression( + AttributeExpression.newBuilder() + .setAttributeId("API.attributeId4") + .build()) + .build()) + .setOperator(Operator.EQ) + .setRhs( + Expression.newBuilder() + .setLiteral( + LiteralConstant.newBuilder() + .setValue( + Value.newBuilder() + .setValueType(ValueType.STRING) + .setString("value") + .build()) + .build()) + .build()) + .build()) + .build()) + .build()) + .addChildFilter( + Filter.newBuilder() + .setLhs( + Expression.newBuilder() + .setAttributeExpression( + AttributeExpression.newBuilder() + .setAttributeId("API.attributeId1") + .build()) + .build()) + .setOperator(Operator.EQ) + .setRhs( + Expression.newBuilder() + .setLiteral( + LiteralConstant.newBuilder() + .setValue( + Value.newBuilder() + .setValueType(ValueType.STRING) + .setString("value") + .build()) + .build()) + .build()) + .build()) + .addChildFilter( + Filter.newBuilder() + .setLhs( + Expression.newBuilder() + .setAttributeExpression( + AttributeExpression.newBuilder() + .setAttributeId("API.attributeId2") + .build()) + .build()) + .setOperator(Operator.EQ) + .setRhs( + Expression.newBuilder() + .setLiteral( + LiteralConstant.newBuilder() + .setValue( + Value.newBuilder() + .setValueType(ValueType.STRING) + .setString("value") + .build()) + .build()) + .build()) + .build()) + .build()) + .addSelection( + Expression.newBuilder() + .setFunction( + FunctionExpression.newBuilder() + .setFunction(FunctionType.AVG) + .setAlias("AVG_Duration") + .addArguments( + Expression.newBuilder() + .setColumnIdentifier( + ColumnIdentifier.newBuilder() + .setColumnName("Api.Trace.metrics.duration_millis"))))) + .addSelection( + Expression.newBuilder() + .setFunction( + FunctionExpression.newBuilder() + .setFunction(FunctionType.AVGRATE) + .setAlias("RATE_Duration") + .addArguments( + Expression.newBuilder() + .setColumnIdentifier( + ColumnIdentifier.newBuilder() + .setColumnName("Api.Trace.metrics.duration_millis"))) + .addArguments( + Expression.newBuilder() + .setLiteral( + LiteralConstant.newBuilder() + .setValue( + Value.newBuilder() + .setValueType(ValueType.LONG) + .setLong(30)))))) + .build(); + + Map attributeMap = + Map.of( + "API.attributeId1", + AttributeMetadata.newBuilder() + .addAllSources(List.of(AttributeSource.EDS, AttributeSource.QS)) + .build(), + "API.attributeId2", + AttributeMetadata.newBuilder().addSources(AttributeSource.QS).build(), + "API.attributeId3", + AttributeMetadata.newBuilder() + .addAllSources(List.of(AttributeSource.EDS, AttributeSource.QS)) + .build(), + "API.attributeId4", + AttributeMetadata.newBuilder().addSources(AttributeSource.EDS).build()); + Filter edsFilter = + buildFilter(exploreRequest.getFilter(), AttributeSource.EDS, attributeMap).get(); + + AttributeMetadataProvider attributeMetadataProvider = mock(AttributeMetadataProvider.class); + QueryServiceEntityFetcher queryServiceEntityFetcher = mock(QueryServiceEntityFetcher.class); + EntityServiceEntityFetcher entityServiceEntityFetcher = mock(EntityServiceEntityFetcher.class); + EntityIdColumnsConfigs entityIdColumnsConfigs = mock(EntityIdColumnsConfigs.class); + QueryServiceClient queryServiceClient = mock(QueryServiceClient.class); + RequestHandler requestHandler = + new RequestHandler( + queryServiceClient, + attributeMetadataProvider, + entityIdColumnsConfigs, + queryServiceEntityFetcher, + entityServiceEntityFetcher); + ExploreRequestContext newExploreRequestContext = + new ExploreRequestContext(RequestContext.forTenantId("tenantId"), exploreRequest); + + when(entityIdColumnsConfigs.getIdKey("API")).thenReturn(Optional.of("entityId")); + when(attributeMetadataProvider.getAttributesMetadata(any(), any())).thenReturn(attributeMap); + when(attributeMetadataProvider.getAttributeMetadata(any(), any(), eq("entityId"))) + .thenReturn(Optional.of(AttributeMetadata.newBuilder().setId("API.id").build())); + when(attributeMetadataProvider.getAttributeMetadata(any(), any(), any())) + .thenReturn(Optional.of(AttributeMetadata.newBuilder().setId("API.timestampId").build())); + when(queryServiceEntityFetcher.getEntities(any(), any())) + .thenReturn( + new EntityFetcherResponse( + Map.of( + EntityKey.from("entityId1"), + Entity.newBuilder().setEntityType("API").setId("entityId1"), + EntityKey.from("entityId2"), + Entity.newBuilder().setEntityType("API").setId("entityId2"), + EntityKey.from("entityId3"), + Entity.newBuilder().setEntityType("API").setId("entityId3")))); + when(entityServiceEntityFetcher.getResults( + any(), + eq( + ExploreRequest.newBuilder() + .setContext("API") + .setFilter(edsFilter) + .addGroupBy( + Expression.newBuilder() + .setAttributeExpression( + AttributeExpression.newBuilder() + .setAttributeId("API.timestampId") + .build()) + .build()) + .build()), + eq(Set.of("entityId1", "entityId2", "entityId3")))) + .thenReturn( + List.of( + Row.newBuilder() + .putColumns("entityId", Value.newBuilder().setString("entityId1").build()) + .build(), + Row.newBuilder() + .putColumns("entityId", Value.newBuilder().setString("entityId2").build()) + .build())); + + when(queryServiceClient.executeQuery(any(), eq(getQueryRequest()))) + .thenReturn( + List.of( + ResultSetChunk.newBuilder() + .setResultSetMetadata( + ResultSetMetadata.newBuilder() + .addColumnMetadata( + ColumnMetadata.newBuilder() + .setColumnName("columnName1") + .build()) + .addColumnMetadata( + ColumnMetadata.newBuilder() + .setColumnName("columnName2") + .build()) + .build()) + .addRow( + org.hypertrace.core.query.service.api.Row.newBuilder() + .addColumn( + org.hypertrace.core.query.service.api.Value.newBuilder() + .setString("value1") + .build()) + .addColumn( + org.hypertrace.core.query.service.api.Value.newBuilder() + .setString("value2") + .build()) + .build()) + .addRow( + org.hypertrace.core.query.service.api.Row.newBuilder() + .addColumn( + org.hypertrace.core.query.service.api.Value.newBuilder() + .setString("value3") + .build()) + .addColumn( + org.hypertrace.core.query.service.api.Value.newBuilder() + .setString("value4") + .build()) + .build()) + .build()) + .iterator()); + Builder responseBuilder = + requestHandler.handleRequest(newExploreRequestContext, exploreRequest); + Assertions.assertEquals(2, responseBuilder.getRowList().size()); + } + + private QueryRequest getQueryRequest() throws InvalidProtocolBufferException { + QueryRequest.Builder requestBuilder = QueryRequest.newBuilder(); + JsonFormat.parser() + .ignoringUnknownFields() + .merge( + "{\n" + + " \"filter\": {\n" + + " \"childFilter\": [{\n" + + " \"childFilter\": [{\n" + + " \"lhs\": {\n" + + " \"attributeExpression\": {\n" + + " \"attributeId\": \"API.timestampId\"\n" + + " }\n" + + " },\n" + + " \"operator\": \"GE\",\n" + + " \"rhs\": {\n" + + " \"literal\": {\n" + + " \"value\": {\n" + + " \"valueType\": \"LONG\"\n" + + " }\n" + + " }\n" + + " }\n" + + " }, {\n" + + " \"lhs\": {\n" + + " \"attributeExpression\": {\n" + + " \"attributeId\": \"API.timestampId\"\n" + + " }\n" + + " },\n" + + " \"operator\": \"LT\",\n" + + " \"rhs\": {\n" + + " \"literal\": {\n" + + " \"value\": {\n" + + " \"valueType\": \"LONG\"\n" + + " }\n" + + " }\n" + + " }\n" + + " }]\n" + + " }, {\n" + + " \"childFilter\": [{\n" + + " \"childFilter\": [{\n" + + " \"childFilter\": [{\n" + + " \"lhs\": {\n" + + " \"attributeExpression\": {\n" + + " \"attributeId\": \"API.attributeId3\"\n" + + " }\n" + + " },\n" + + " \"operator\": \"EQ\",\n" + + " \"rhs\": {\n" + + " \"literal\": {\n" + + " \"value\": {\n" + + " \"string\": \"value\"\n" + + " }\n" + + " }\n" + + " }\n" + + " }]\n" + + " }]\n" + + " }, {\n" + + " \"lhs\": {\n" + + " \"attributeExpression\": {\n" + + " \"attributeId\": \"API.attributeId1\"\n" + + " }\n" + + " },\n" + + " \"operator\": \"EQ\",\n" + + " \"rhs\": {\n" + + " \"literal\": {\n" + + " \"value\": {\n" + + " \"string\": \"value\"\n" + + " }\n" + + " }\n" + + " }\n" + + " }, {\n" + + " \"lhs\": {\n" + + " \"attributeExpression\": {\n" + + " \"attributeId\": \"API.attributeId2\"\n" + + " }\n" + + " },\n" + + " \"operator\": \"EQ\",\n" + + " \"rhs\": {\n" + + " \"literal\": {\n" + + " \"value\": {\n" + + " \"string\": \"value\"\n" + + " }\n" + + " }\n" + + " }\n" + + " }]\n" + + " }, {\n" + + " \"lhs\": {\n" + + " \"attributeExpression\": {\n" + + " \"attributeId\": \"API.timestampId\"\n" + + " }\n" + + " },\n" + + " \"operator\": \"IN\",\n" + + " \"rhs\": {\n" + + " \"literal\": {\n" + + " \"value\": {\n" + + " \"valueType\": \"STRING_ARRAY\",\n" + + " \"stringArray\": [\"entityId1\", \"entityId2\"]\n" + + " }\n" + + " }\n" + + " }\n" + + " }]\n" + + " },\n" + + " \"selection\": [{\n" + + " \"function\": {\n" + + " \"functionName\": \"AVG\",\n" + + " \"arguments\": [{\n" + + " \"attributeExpression\": {\n" + + " \"attributeId\": \"Api.Trace.metrics.duration_millis\"\n" + + " }\n" + + " }],\n" + + " \"alias\": \"AVG_Duration\"\n" + + " }\n" + + " }, {\n" + + " \"function\": {\n" + + " \"functionName\": \"AVGRATE\",\n" + + " \"arguments\": [{\n" + + " \"attributeExpression\": {\n" + + " \"attributeId\": \"Api.Trace.metrics.duration_millis\"\n" + + " }\n" + + " }, {\n" + + " \"literal\": {\n" + + " \"value\": {\n" + + " \"string\": \"PT30S\"\n" + + " }\n" + + " }\n" + + " }],\n" + + " \"alias\": \"RATE_Duration\"\n" + + " }\n" + + " }]\n" + + "}", + requestBuilder); + + return requestBuilder.build(); + } + + private Optional buildFilter( + org.hypertrace.gateway.service.v1.common.Filter filter, + AttributeSource source, + Map attributeMetadataMap) { + if (filter.equals(org.hypertrace.gateway.service.v1.common.Filter.getDefaultInstance())) { + return Optional.empty(); + } + + org.hypertrace.gateway.service.v1.common.Operator operator = filter.getOperator(); + switch (operator) { + case UNDEFINED: + return Optional.empty(); + case AND: + case OR: + return buildCompositeFilter(filter, source, operator, attributeMetadataMap); + default: + List availableSources = + attributeMetadataMap + .get( + ExpressionReader.getAttributeIdFromAttributeSelection(filter.getLhs()) + .orElseThrow()) + .getSourcesList(); + return availableSources.contains(source) + ? Optional.of( + org.hypertrace.gateway.service.v1.common.Filter.newBuilder(filter).build()) + : Optional.empty(); + } + } + + private Optional buildCompositeFilter( + org.hypertrace.gateway.service.v1.common.Filter filter, + AttributeSource source, + org.hypertrace.gateway.service.v1.common.Operator operator, + Map attributeMetadataMap) { + org.hypertrace.gateway.service.v1.common.Filter.Builder filterBuilder = + org.hypertrace.gateway.service.v1.common.Filter.newBuilder(); + for (org.hypertrace.gateway.service.v1.common.Filter childFilter : + filter.getChildFilterList()) { + buildFilter(childFilter, source, attributeMetadataMap) + .ifPresent(filterBuilder::addChildFilter); + } + if (filterBuilder.getChildFilterCount() > 0) { + filterBuilder.setOperator(operator); + return Optional.of(filterBuilder.build()); + } else { + return Optional.empty(); + } + } } diff --git a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/TimeAggregationsRequestHandlerTest.java b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/TimeAggregationsRequestHandlerTest.java index e7b6c547..78e58b79 100644 --- a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/TimeAggregationsRequestHandlerTest.java +++ b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/TimeAggregationsRequestHandlerTest.java @@ -4,8 +4,11 @@ import java.util.List; import org.hypertrace.gateway.service.common.AttributeMetadataProvider; +import org.hypertrace.gateway.service.common.datafetcher.QueryServiceEntityFetcher; import org.hypertrace.gateway.service.common.util.QueryExpressionUtil; import org.hypertrace.gateway.service.common.util.QueryServiceClient; +import org.hypertrace.gateway.service.entity.config.EntityIdColumnsConfigs; +import org.hypertrace.gateway.service.explore.entity.EntityServiceEntityFetcher; import org.hypertrace.gateway.service.v1.common.Expression; import org.hypertrace.gateway.service.v1.common.FunctionExpression; import org.hypertrace.gateway.service.v1.common.FunctionType; @@ -64,7 +67,11 @@ public void intervalStartTimeOrderByShouldBeAddedToOrderByListAndAliasShouldMatc TimeAggregationsRequestHandler requestHandler = new TimeAggregationsRequestHandler( - mock(QueryServiceClient.class), mock(AttributeMetadataProvider.class)); + mock(QueryServiceClient.class), + mock(AttributeMetadataProvider.class), + mock(EntityIdColumnsConfigs.class), + mock(QueryServiceEntityFetcher.class), + mock(EntityServiceEntityFetcher.class)); List orderByExpressions = requestHandler.getRequestOrderByExpressions(exploreRequest); @@ -112,7 +119,11 @@ public void intervalStartTimeOrderingNotAddedIfAlreadyRequested() { TimeAggregationsRequestHandler requestHandler = new TimeAggregationsRequestHandler( - mock(QueryServiceClient.class), mock(AttributeMetadataProvider.class)); + mock(QueryServiceClient.class), + mock(AttributeMetadataProvider.class), + mock(EntityIdColumnsConfigs.class), + mock(QueryServiceEntityFetcher.class), + mock(EntityServiceEntityFetcher.class)); List orderByExpressions = requestHandler.getRequestOrderByExpressions(exploreRequest); diff --git a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/entity/EntityRequestHandlerTest.java b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/entity/EntityRequestHandlerTest.java index 7b5aaa31..88b94515 100644 --- a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/entity/EntityRequestHandlerTest.java +++ b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/entity/EntityRequestHandlerTest.java @@ -7,23 +7,20 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; import org.hypertrace.core.attribute.service.v1.AttributeKind; import org.hypertrace.core.attribute.service.v1.AttributeMetadata; -import org.hypertrace.entity.query.service.v1.ColumnMetadata; -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.core.attribute.service.v1.AttributeSource; import org.hypertrace.gateway.service.common.AttributeMetadataProvider; import org.hypertrace.gateway.service.common.datafetcher.EntityFetcherResponse; import org.hypertrace.gateway.service.common.datafetcher.QueryServiceEntityFetcher; import org.hypertrace.gateway.service.common.util.QueryServiceClient; import org.hypertrace.gateway.service.entity.EntitiesRequestContext; import org.hypertrace.gateway.service.entity.EntityKey; +import org.hypertrace.gateway.service.entity.config.EntityIdColumnsConfigs; import org.hypertrace.gateway.service.explore.ExploreRequestContext; import org.hypertrace.gateway.service.v1.common.ColumnIdentifier; import org.hypertrace.gateway.service.v1.common.Expression; @@ -58,6 +55,7 @@ void setup() { this.entityRequestHandler = new EntityRequestHandler( attributeMetadataProvider, + mock(EntityIdColumnsConfigs.class), mock(QueryServiceClient.class), queryServiceEntityFetcher, entityServiceEntityFetcher); @@ -89,11 +87,19 @@ void shouldBuildEntityResponse_multipleDataSources() { AttributeMetadata.newBuilder() .setKey("API.type") .setValueKind(AttributeKind.TYPE_STRING) + .addSources(AttributeSource.EDS) .build(), "API.external", AttributeMetadata.newBuilder() .setKey("API.external") .setValueKind(AttributeKind.TYPE_STRING) + .addSources(AttributeSource.EDS) + .build(), + "API.name", + AttributeMetadata.newBuilder() + .setKey("API.name") + .setValueKind(AttributeKind.TYPE_STRING) + .addSources(AttributeSource.EDS) .build())); EntitiesRequest entitiesRequest = @@ -154,11 +160,19 @@ void testHandleRequest_emptyEntityIds() { AttributeMetadata.newBuilder() .setKey("API.type") .setValueKind(AttributeKind.TYPE_STRING) + .addSources(AttributeSource.EDS) .build(), "API.external", AttributeMetadata.newBuilder() .setKey("API.external") .setValueKind(AttributeKind.TYPE_STRING) + .addSources(AttributeSource.EDS) + .build(), + "API.name", + AttributeMetadata.newBuilder() + .setKey("API.name") + .setValueKind(AttributeKind.TYPE_STRING) + .addSources(AttributeSource.EDS) .build())); EntitiesRequest entitiesRequest = @@ -185,43 +199,36 @@ private EntityFetcherResponse mockEntityFetcherResponse() { Entity.newBuilder().setEntityType("API").setId("api2"))); } - private Iterator mockResults() { - ResultSetChunk chunk = - ResultSetChunk.newBuilder() - .setResultSetMetadata( - ResultSetMetadata.newBuilder() - .addColumnMetadata( - ColumnMetadata.newBuilder().setColumnName("API.type").build()) - .addColumnMetadata( - ColumnMetadata.newBuilder().setColumnName("COUNT_API.external_[]").build()) + private List mockResults() { + return List.of( + org.hypertrace.gateway.service.v1.common.Row.newBuilder() + .putColumns( + "API.type", + org.hypertrace.gateway.service.v1.common.Value.newBuilder() + .setValueType(org.hypertrace.gateway.service.v1.common.ValueType.STRING) + .setString("HTTP") .build()) - .addRow( - Row.newBuilder() - .addColumn( - org.hypertrace.entity.query.service.v1.Value.newBuilder() - .setValueType(org.hypertrace.entity.query.service.v1.ValueType.STRING) - .setString("HTTP") - .build()) - .addColumn( - org.hypertrace.entity.query.service.v1.Value.newBuilder() - .setValueType(org.hypertrace.entity.query.service.v1.ValueType.LONG) - .setLong(12)) + .putColumns( + "COUNT_API.external_[]", + org.hypertrace.gateway.service.v1.common.Value.newBuilder() + .setValueType(org.hypertrace.gateway.service.v1.common.ValueType.LONG) + .setLong(12) .build()) - .addRow( - Row.newBuilder() - .addColumn( - org.hypertrace.entity.query.service.v1.Value.newBuilder() - .setValueType(org.hypertrace.entity.query.service.v1.ValueType.STRING) - .setString("GRPC") - .build()) - .addColumn( - org.hypertrace.entity.query.service.v1.Value.newBuilder() - .setValueType(org.hypertrace.entity.query.service.v1.ValueType.LONG) - .setLong(24)) + .build(), + org.hypertrace.gateway.service.v1.common.Row.newBuilder() + .putColumns( + "API.type", + org.hypertrace.gateway.service.v1.common.Value.newBuilder() + .setValueType(org.hypertrace.gateway.service.v1.common.ValueType.STRING) + .setString("GRPC") .build()) - .build(); - - return List.of(chunk).iterator(); + .putColumns( + "COUNT_API.external_[]", + org.hypertrace.gateway.service.v1.common.Value.newBuilder() + .setValueType(org.hypertrace.gateway.service.v1.common.ValueType.LONG) + .setLong(24) + .build()) + .build()); } private Filter createEqFilter(String column, String value) { diff --git a/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/entity/EntityServiceEntityFetcherTest.java b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/entity/EntityServiceEntityFetcherTest.java new file mode 100644 index 00000000..734a243a --- /dev/null +++ b/gateway-service-impl/src/test/java/org/hypertrace/gateway/service/explore/entity/EntityServiceEntityFetcherTest.java @@ -0,0 +1,168 @@ +package org.hypertrace.gateway.service.explore.entity; + +import static org.hypertrace.core.grpcutils.context.RequestContext.forTenantId; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.hypertrace.core.attribute.service.v1.AttributeKind; +import org.hypertrace.core.attribute.service.v1.AttributeMetadata; +import org.hypertrace.core.attribute.service.v1.AttributeSource; +import org.hypertrace.entity.query.service.client.EntityQueryServiceClient; +import org.hypertrace.entity.query.service.v1.ColumnMetadata; +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.gateway.service.common.AttributeMetadataProvider; +import org.hypertrace.gateway.service.entity.config.EntityIdColumnsConfigs; +import org.hypertrace.gateway.service.explore.ExploreRequestContext; +import org.hypertrace.gateway.service.v1.common.ColumnIdentifier; +import org.hypertrace.gateway.service.v1.common.Expression; +import org.hypertrace.gateway.service.v1.common.Filter; +import org.hypertrace.gateway.service.v1.common.FunctionExpression; +import org.hypertrace.gateway.service.v1.common.FunctionType; +import org.hypertrace.gateway.service.v1.common.LiteralConstant; +import org.hypertrace.gateway.service.v1.common.Operator; +import org.hypertrace.gateway.service.v1.common.Value; +import org.hypertrace.gateway.service.v1.common.ValueType; +import org.hypertrace.gateway.service.v1.explore.ExploreRequest; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class EntityServiceEntityFetcherTest { + private AttributeMetadataProvider attributeMetadataProvider; + private EntityIdColumnsConfigs entityIdColumnsConfigs; + private EntityQueryServiceClient entityQueryServiceClient; + + @Test + void testGetEntities() { + AttributeMetadataProvider attributeMetadataProvider = mock(AttributeMetadataProvider.class); + EntityIdColumnsConfigs entityIdColumnsConfigs = mock(EntityIdColumnsConfigs.class); + EntityQueryServiceClient entityQueryServiceClient = mock(EntityQueryServiceClient.class); + EntityServiceEntityFetcher entityServiceEntityFetcher = + new EntityServiceEntityFetcher( + attributeMetadataProvider, entityIdColumnsConfigs, entityQueryServiceClient); + + when(entityQueryServiceClient.execute(any(), any())).thenReturn(mockResults()); + + Expression aggregation = createFunctionExpression("API.external"); + ExploreRequest exploreRequest = + ExploreRequest.newBuilder() + .setContext("API") + .setStartTimeMillis(123L) + .setEndTimeMillis(234L) + .setFilter(createEqFilter("API.name", "api1")) + .addSelection(aggregation) + .addGroupBy(createColumnExpression("API.type")) + .build(); + ExploreRequestContext exploreRequestContext = + new ExploreRequestContext(forTenantId("customer1"), exploreRequest); + exploreRequestContext.mapAliasToFunctionExpression( + "COUNT_API.external_[]", aggregation.getFunction()); + + when(attributeMetadataProvider.getAttributesMetadata(exploreRequestContext, "API")) + .thenReturn( + Map.of( + "API.type", + AttributeMetadata.newBuilder() + .setKey("API.type") + .setValueKind(AttributeKind.TYPE_STRING) + .addSources(AttributeSource.EDS) + .build(), + "API.external", + AttributeMetadata.newBuilder() + .setKey("API.external") + .setValueKind(AttributeKind.TYPE_STRING) + .addSources(AttributeSource.EDS) + .build(), + "API.name", + AttributeMetadata.newBuilder() + .setKey("API.name") + .setValueKind(AttributeKind.TYPE_STRING) + .addSources(AttributeSource.EDS) + .build())); + + List rows = + entityServiceEntityFetcher.getResults( + exploreRequestContext, exploreRequest, Set.of("api1", "api2")); + Assertions.assertEquals(rows.size(), 2); + } + + private Iterator mockResults() { + ResultSetChunk chunk = + ResultSetChunk.newBuilder() + .setResultSetMetadata( + ResultSetMetadata.newBuilder() + .addColumnMetadata( + ColumnMetadata.newBuilder().setColumnName("API.type").build()) + .addColumnMetadata( + ColumnMetadata.newBuilder().setColumnName("COUNT_API.external_[]").build()) + .build()) + .addRow( + Row.newBuilder() + .addColumn( + org.hypertrace.entity.query.service.v1.Value.newBuilder() + .setValueType(org.hypertrace.entity.query.service.v1.ValueType.STRING) + .setString("HTTP") + .build()) + .addColumn( + org.hypertrace.entity.query.service.v1.Value.newBuilder() + .setValueType(org.hypertrace.entity.query.service.v1.ValueType.LONG) + .setLong(12)) + .build()) + .addRow( + Row.newBuilder() + .addColumn( + org.hypertrace.entity.query.service.v1.Value.newBuilder() + .setValueType(org.hypertrace.entity.query.service.v1.ValueType.STRING) + .setString("GRPC") + .build()) + .addColumn( + org.hypertrace.entity.query.service.v1.Value.newBuilder() + .setValueType(org.hypertrace.entity.query.service.v1.ValueType.LONG) + .setLong(24)) + .build()) + .build(); + + return List.of(chunk).iterator(); + } + + private Expression createColumnExpression(String column) { + return Expression.newBuilder() + .setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(column).build()) + .build(); + } + + private Expression createFunctionExpression(String column) { + return Expression.newBuilder() + .setFunction( + FunctionExpression.newBuilder() + .setFunction(FunctionType.COUNT) + .setAlias("COUNT_" + column + "_[]") + .addArguments(createColumnExpression(column)) + .build()) + .build(); + } + + private Filter createEqFilter(String column, String value) { + return Filter.newBuilder() + .setLhs(createColumnExpression(column)) + .setOperator(Operator.EQ) + .setRhs( + Expression.newBuilder() + .setLiteral( + LiteralConstant.newBuilder() + .setValue( + Value.newBuilder() + .setString(value) + .setValueType(ValueType.STRING) + .build()) + .build()) + .build()) + .build(); + } +}