From 7d11119ed72e106f4464f475bd1e399cee5e2ed2 Mon Sep 17 00:00:00 2001 From: aman-bansal Date: Thu, 6 Jun 2024 11:05:54 +0530 Subject: [PATCH 1/2] chore | adding the support for total and pagination in explore --- .../gateway/service/v1/explore.proto | 2 + .../service/explore/RequestHandler.java | 71 +++++++++++++++++- .../explore/entity/EntityRequestHandler.java | 26 ++----- .../entity/EntityServiceEntityFetcher.java | 72 ++++++++++++++++++- 4 files changed, 149 insertions(+), 22 deletions(-) diff --git a/gateway-service-api/src/main/proto/org/hypertrace/gateway/service/v1/explore.proto b/gateway-service-api/src/main/proto/org/hypertrace/gateway/service/v1/explore.proto index 2173c754..cd20f233 100644 --- a/gateway-service-api/src/main/proto/org/hypertrace/gateway/service/v1/explore.proto +++ b/gateway-service-api/src/main/proto/org/hypertrace/gateway/service/v1/explore.proto @@ -33,6 +33,7 @@ message ExploreRequest { string space_id = 22; optional ContextOption context_option = 23; + bool fetch_total = 24; } message ContextOption { @@ -47,6 +48,7 @@ message EntityOption { message ExploreResponse { repeated org.hypertrace.gateway.service.v1.common.Row row = 2; + int32 total = 3; } // Used to set column names that are not there in the ExploreRequest selections eg. interval start time. To maintain 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 b74faed1..2ce531df 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 @@ -17,6 +17,7 @@ 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.Function; import org.hypertrace.core.query.service.api.QueryRequest; import org.hypertrace.core.query.service.api.ResultSetChunk; import org.hypertrace.core.query.service.api.ResultSetMetadata; @@ -38,6 +39,7 @@ import org.hypertrace.gateway.service.v1.common.AttributeExpression; import org.hypertrace.gateway.service.v1.common.Expression; 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; @@ -67,6 +69,7 @@ */ public class RequestHandler implements RequestHandlerWithSorting { private static final Logger LOG = LoggerFactory.getLogger(RequestHandler.class); + private static final String TOTAL_ALIAS_NAME = "total"; private static final Expression NULL_VALUE_EXPRESSION = Expression.newBuilder() .setLiteral( @@ -107,9 +110,73 @@ public ExploreResponse.Builder handleRequest( buildQueryRequest(requestContext, request, attributeMetadataProvider); Iterator resultSetChunkIterator = executeQuery(requestContext, queryRequest); + ExploreResponse.Builder responseBuilder = + handleQueryServiceResponse( + request, + requestContext, + resultSetChunkIterator, + requestContext, + attributeMetadataProvider); + + if (request.getFetchTotal()) { + QueryRequest totalQueryRequest = buildTotalQueryRequest(request, queryRequest); + Iterator totalIterator = executeQuery(requestContext, totalQueryRequest); + while (totalIterator.hasNext()) { + ResultSetChunk chunk = totalIterator.next(); + LOG.info("Total received chunk is: {}", chunk); + if (chunk.getRowCount() < 1) { + break; + } + + if (!chunk.hasResultSetMetadata() + || chunk.getResultSetMetadata().getColumnMetadataList().size() != 1 + || !TOTAL_ALIAS_NAME.equals( + chunk.getResultSetMetadata().getColumnMetadata(0).getColumnName())) { + LOG.warn("Chunk doesn't have result metadata so couldn't process the response."); + break; + } + responseBuilder.setTotal(chunk.getRow(0).getColumn(0).getInt()); + } + } - return handleQueryServiceResponse( - request, requestContext, resultSetChunkIterator, requestContext, attributeMetadataProvider); + return responseBuilder; + } + + QueryRequest buildTotalQueryRequest(ExploreRequest request, QueryRequest queryRequest) { + QueryRequest.Builder builder = QueryRequest.newBuilder(); + List groupBys = + ExpressionReader.getAttributeExpressions(request.getGroupByList()); + // add all group-by columns as distinct count to get total + builder.addSelection( + org.hypertrace.core.query.service.api.Expression.newBuilder() + .setFunction( + Function.newBuilder() + .addAllArguments( + groupBys.stream() + .map( + attribute -> + org.hypertrace.core.query.service.api.Expression.newBuilder() + .setColumnIdentifier( + org.hypertrace.core.query.service.api.ColumnIdentifier + .newBuilder() + .setColumnName( + attribute + .getAttributeExpression() + .getAttributeId()) + .setAlias( + attribute.getAttributeExpression().getAlias()) + .build()) + .build()) + .collect(Collectors.toUnmodifiableList())) + .setFunctionName(FunctionType.DISTINCTCOUNT.name()) + .setAlias(TOTAL_ALIAS_NAME) + .build()) + .build()); + + // Add filter + builder.setFilter(queryRequest.getFilter()); + builder.setLimit(1); + return builder.build(); } QueryRequest buildQueryRequest( 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 d76df61c..c5212c04 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 @@ -35,8 +35,6 @@ */ public class EntityRequestHandler extends RequestHandler { private final EntityServiceEntityFetcher entityServiceEntityFetcher; - private final AttributeMetadataProvider attributeMetadataProvider; - private final QueryServiceEntityFetcher queryServiceEntityFetcher; public EntityRequestHandler( AttributeMetadataProvider attributeMetadataProvider, @@ -50,9 +48,7 @@ public EntityRequestHandler( entityIdColumnsConfig, queryServiceEntityFetcher, entityServiceEntityFetcher); - this.attributeMetadataProvider = attributeMetadataProvider; this.entityServiceEntityFetcher = entityServiceEntityFetcher; - this.queryServiceEntityFetcher = queryServiceEntityFetcher; } @Override @@ -67,7 +63,8 @@ public ExploreResponse.Builder handleRequest( ExploreResponse.Builder builder = ExploreResponse.newBuilder(); Set entityIds = new HashSet<>(); Optional maybeEntityOption = getEntityOption(exploreRequest); - if (requestOnLiveEntities(maybeEntityOption)) { + boolean requestOnLiveEntities = requestOnLiveEntities(maybeEntityOption); + if (requestOnLiveEntities) { entityIds.addAll(getEntityIdsInTimeRangeFromQueryService(requestContext, exploreRequest)); if (entityIds.isEmpty()) { return builder; @@ -76,21 +73,15 @@ public ExploreResponse.Builder handleRequest( 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()) { - sortAndPaginatePostProcess( - builder, - requestContext.getOrderByExpressions(), - requestContext.getRowLimitBeforeRest(), - requestContext.getOffset()); - } - if (requestContext.hasGroupBy() && requestContext.getIncludeRestGroup()) { getTheRestGroupRequestHandler() .getRowsForTheRestGroup(requestContext, exploreRequest, builder); } + if (exploreRequest.getFetchTotal()) { + builder.setTotal(entityServiceEntityFetcher.getTotal(requestContext, exploreRequest)); + } + return builder; } @@ -122,10 +113,7 @@ protected List sortAndPagi } private boolean requestOnLiveEntities(Optional entityOption) { - if (entityOption.isEmpty()) { - return true; - } - return !entityOption.get().getIncludeNonLiveEntities(); + return entityOption.map(option -> !option.getIncludeNonLiveEntities()).orElse(true); } private Optional getEntityOption(ExploreRequest exploreRequest) { 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 0c55ad90..aae9b45a 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 @@ -2,6 +2,7 @@ import com.google.common.collect.Streams; import java.util.ArrayList; +import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.Map; @@ -9,11 +10,13 @@ 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.ColumnIdentifier; 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; import org.hypertrace.entity.query.service.v1.Filter.Builder; +import org.hypertrace.entity.query.service.v1.Function; import org.hypertrace.entity.query.service.v1.LiteralConstant; import org.hypertrace.entity.query.service.v1.Operator; import org.hypertrace.entity.query.service.v1.ResultSetChunk; @@ -28,6 +31,7 @@ import org.hypertrace.gateway.service.entity.config.EntityIdColumnsConfig; import org.hypertrace.gateway.service.explore.ExploreRequestContext; import org.hypertrace.gateway.service.v1.common.FunctionExpression; +import org.hypertrace.gateway.service.v1.common.FunctionType; import org.hypertrace.gateway.service.v1.common.Row; import org.hypertrace.gateway.service.v1.explore.ExploreRequest; import org.slf4j.Logger; @@ -35,6 +39,7 @@ public class EntityServiceEntityFetcher { private static final Logger LOGGER = LoggerFactory.getLogger(EntityServiceEntityFetcher.class); + private static final String TOTAL_ALIAS_NAME = "total"; private static final int DEFAULT_ENTITY_REQUEST_LIMIT = 10_000; private final AttributeMetadataProvider attributeMetadataProvider; private final EntityIdColumnsConfig entityIdColumnsConfig; @@ -57,6 +62,30 @@ public List getResults( return readChunkResults(requestContext, result); } + public int getTotal(ExploreRequestContext requestContext, ExploreRequest exploreRequest) { + EntityQueryRequest request = buildTotalEntitiesRequest(requestContext, exploreRequest); + Iterator resultSetChunkIterator = + entityQueryServiceClient.execute(request, requestContext.getHeaders()); + while (resultSetChunkIterator.hasNext()) { + ResultSetChunk chunk = resultSetChunkIterator.next(); + LOGGER.info("Total received chunk is: {}", chunk); + if (chunk.getRowCount() < 1) { + break; + } + + if (!chunk.hasResultSetMetadata() + || chunk.getResultSetMetadata().getColumnMetadataList().size() != 1 + || !TOTAL_ALIAS_NAME.equals( + chunk.getResultSetMetadata().getColumnMetadata(0).getColumnName())) { + LOGGER.warn("Chunk doesn't have result metadata so couldn't process the response."); + break; + } + return chunk.getRow(0).getColumn(0).getInt(); + } + + return 0; + } + protected List readChunkResults( ExploreRequestContext requestContext, Iterator resultSetChunkIterator) { List resultRows = new ArrayList<>(); @@ -155,6 +184,46 @@ private Map remapAttributeMetadataByResultName( attributeMetadataByIdMap); } + private EntityQueryRequest buildTotalEntitiesRequest( + ExploreRequestContext requestContext, ExploreRequest exploreRequest) { + String entityType = exploreRequest.getContext(); + + List entityIdAttributeIds = + AttributeMetadataUtil.getIdAttributeIds( + attributeMetadataProvider, entityIdColumnsConfig, requestContext, entityType); + EntityQueryRequest.Builder builder = + EntityQueryRequest.newBuilder() + .setEntityType(entityType) + .setFilter(buildFilter(exploreRequest, entityIdAttributeIds, Collections.emptySet())); + List groupBys = + ExpressionReader.getAttributeExpressions(exploreRequest.getGroupByList()); + builder.addSelection( + Expression.newBuilder() + .setFunction( + Function.newBuilder() + .setFunctionName(FunctionType.DISTINCTCOUNT.name()) + .addAllArguments( + groupBys.stream() + .map( + attribute -> + Expression.newBuilder() + .setColumnIdentifier( + ColumnIdentifier.newBuilder() + .setColumnName( + attribute + .getAttributeExpression() + .getAttributeId()) + .setAlias( + attribute.getAttributeExpression().getAlias()) + .build()) + .build()) + .collect(Collectors.toUnmodifiableList())) + .setAlias(TOTAL_ALIAS_NAME) + .build()) + .build()); + return builder.build(); + } + private EntityQueryRequest buildRequest( ExploreRequestContext requestContext, ExploreRequest exploreRequest, Set entityIds) { String entityType = exploreRequest.getContext(); @@ -169,7 +238,8 @@ private EntityQueryRequest buildRequest( addGroupBys(exploreRequest, builder); addSelections(requestContext, exploreRequest, builder); - builder.setLimit(DEFAULT_ENTITY_REQUEST_LIMIT); + builder.setLimit(exploreRequest.getLimit()); + builder.setOffset(exploreRequest.getOffset()); return builder.build(); } From eb2d1526629891e1b3da8db85ac6100f6098830d Mon Sep 17 00:00:00 2001 From: aman-bansal Date: Thu, 6 Jun 2024 11:18:23 +0530 Subject: [PATCH 2/2] refactor --- .../service/explore/RequestHandler.java | 113 +++++++++--------- .../entity/EntityServiceEntityFetcher.java | 29 ++--- 2 files changed, 71 insertions(+), 71 deletions(-) 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 2ce531df..5ae22add 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 @@ -15,6 +15,7 @@ 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.ColumnIdentifier; import org.hypertrace.core.query.service.api.ColumnMetadata; import org.hypertrace.core.query.service.api.Filter; import org.hypertrace.core.query.service.api.Function; @@ -119,66 +120,12 @@ public ExploreResponse.Builder handleRequest( attributeMetadataProvider); if (request.getFetchTotal()) { - QueryRequest totalQueryRequest = buildTotalQueryRequest(request, queryRequest); - Iterator totalIterator = executeQuery(requestContext, totalQueryRequest); - while (totalIterator.hasNext()) { - ResultSetChunk chunk = totalIterator.next(); - LOG.info("Total received chunk is: {}", chunk); - if (chunk.getRowCount() < 1) { - break; - } - - if (!chunk.hasResultSetMetadata() - || chunk.getResultSetMetadata().getColumnMetadataList().size() != 1 - || !TOTAL_ALIAS_NAME.equals( - chunk.getResultSetMetadata().getColumnMetadata(0).getColumnName())) { - LOG.warn("Chunk doesn't have result metadata so couldn't process the response."); - break; - } - responseBuilder.setTotal(chunk.getRow(0).getColumn(0).getInt()); - } + responseBuilder.setTotal(fetchTotal(requestContext, request, queryRequest)); } return responseBuilder; } - QueryRequest buildTotalQueryRequest(ExploreRequest request, QueryRequest queryRequest) { - QueryRequest.Builder builder = QueryRequest.newBuilder(); - List groupBys = - ExpressionReader.getAttributeExpressions(request.getGroupByList()); - // add all group-by columns as distinct count to get total - builder.addSelection( - org.hypertrace.core.query.service.api.Expression.newBuilder() - .setFunction( - Function.newBuilder() - .addAllArguments( - groupBys.stream() - .map( - attribute -> - org.hypertrace.core.query.service.api.Expression.newBuilder() - .setColumnIdentifier( - org.hypertrace.core.query.service.api.ColumnIdentifier - .newBuilder() - .setColumnName( - attribute - .getAttributeExpression() - .getAttributeId()) - .setAlias( - attribute.getAttributeExpression().getAlias()) - .build()) - .build()) - .collect(Collectors.toUnmodifiableList())) - .setFunctionName(FunctionType.DISTINCTCOUNT.name()) - .setAlias(TOTAL_ALIAS_NAME) - .build()) - .build()); - - // Add filter - builder.setFilter(queryRequest.getFilter()); - builder.setLimit(1); - return builder.build(); - } - QueryRequest buildQueryRequest( ExploreRequestContext requestContext, ExploreRequest request, @@ -294,6 +241,62 @@ private List getEntityIdsToFilterFromSourceEDS( .collect(Collectors.toUnmodifiableList()); } + private int fetchTotal( + ExploreRequestContext requestContext, ExploreRequest request, QueryRequest queryRequest) { + QueryRequest totalQueryRequest = buildTotalQueryRequest(request, queryRequest); + Iterator totalIterator = executeQuery(requestContext, totalQueryRequest); + while (totalIterator.hasNext()) { + ResultSetChunk chunk = totalIterator.next(); + LOG.info("Total received chunk is: {}", chunk); + if (chunk.getRowCount() < 1) { + break; + } + + if (!chunk.hasResultSetMetadata() + || chunk.getResultSetMetadata().getColumnMetadataList().size() != 1 + || !TOTAL_ALIAS_NAME.equals( + chunk.getResultSetMetadata().getColumnMetadata(0).getColumnName())) { + LOG.warn("Chunk doesn't have result metadata so couldn't process the response."); + break; + } + return chunk.getRow(0).getColumn(0).getInt(); + } + return 0; + } + + private QueryRequest buildTotalQueryRequest(ExploreRequest request, QueryRequest queryRequest) { + QueryRequest.Builder builder = QueryRequest.newBuilder(); + List groupBys = + ExpressionReader.getAttributeExpressions(request.getGroupByList()); + // add all group-by columns as distinct count to get total + List groupBysExpression = + groupBys.stream() + .map( + attribute -> + org.hypertrace.core.query.service.api.Expression.newBuilder() + .setColumnIdentifier( + ColumnIdentifier.newBuilder() + .setColumnName(attribute.getAttributeExpression().getAttributeId()) + .setAlias(attribute.getAttributeExpression().getAlias()) + .build()) + .build()) + .collect(Collectors.toUnmodifiableList()); + builder.addSelection( + org.hypertrace.core.query.service.api.Expression.newBuilder() + .setFunction( + Function.newBuilder() + .addAllArguments(groupBysExpression) + .setFunctionName(FunctionType.DISTINCTCOUNT.name()) + .setAlias(TOTAL_ALIAS_NAME) + .build()) + .build()); + + // Add filter + builder.setFilter(queryRequest.getFilter()); + builder.setLimit(1); + return builder.build(); + } + private ExploreRequest buildExploreRequest( ExploreRequestContext exploreRequestContext, String context, 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 aae9b45a..2d82a3e5 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 @@ -197,27 +197,24 @@ private EntityQueryRequest buildTotalEntitiesRequest( .setFilter(buildFilter(exploreRequest, entityIdAttributeIds, Collections.emptySet())); List groupBys = ExpressionReader.getAttributeExpressions(exploreRequest.getGroupByList()); + List groupBysExpression = + groupBys.stream() + .map( + attribute -> + Expression.newBuilder() + .setColumnIdentifier( + ColumnIdentifier.newBuilder() + .setColumnName(attribute.getAttributeExpression().getAttributeId()) + .setAlias(attribute.getAttributeExpression().getAlias()) + .build()) + .build()) + .collect(Collectors.toUnmodifiableList()); builder.addSelection( Expression.newBuilder() .setFunction( Function.newBuilder() .setFunctionName(FunctionType.DISTINCTCOUNT.name()) - .addAllArguments( - groupBys.stream() - .map( - attribute -> - Expression.newBuilder() - .setColumnIdentifier( - ColumnIdentifier.newBuilder() - .setColumnName( - attribute - .getAttributeExpression() - .getAttributeId()) - .setAlias( - attribute.getAttributeExpression().getAlias()) - .build()) - .build()) - .collect(Collectors.toUnmodifiableList())) + .addAllArguments(groupBysExpression) .setAlias(TOTAL_ALIAS_NAME) .build()) .build());