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..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,8 +15,10 @@ 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; 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 +40,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 +70,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 +111,19 @@ 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()) { + responseBuilder.setTotal(fetchTotal(requestContext, request, queryRequest)); + } - return handleQueryServiceResponse( - request, requestContext, resultSetChunkIterator, requestContext, attributeMetadataProvider); + return responseBuilder; } QueryRequest buildQueryRequest( @@ -227,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/EntityRequestHandler.java b/gateway-service-impl/src/main/java/org/hypertrace/gateway/service/explore/entity/EntityRequestHandler.java index ed2bb7b3..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 @@ -78,6 +78,10 @@ public ExploreResponse.Builder handleRequest( .getRowsForTheRestGroup(requestContext, exploreRequest, builder); } + if (exploreRequest.getFetchTotal()) { + builder.setTotal(entityServiceEntityFetcher.getTotal(requestContext, exploreRequest)); + } + return builder; } 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 4a4aad1c..ab15b41a 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; @@ -56,6 +60,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<>(); @@ -154,6 +182,43 @@ 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()); + 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(groupBysExpression) + .setAlias(TOTAL_ALIAS_NAME) + .build()) + .build()); + return builder.build(); + } + private EntityQueryRequest buildRequest( ExploreRequestContext requestContext, ExploreRequest exploreRequest, Set entityIds) { String entityType = exploreRequest.getContext();