Skip to content
This repository has been archived by the owner on Jun 26, 2024. It is now read-only.

Commit

Permalink
refactor
Browse files Browse the repository at this point in the history
  • Loading branch information
@aman-bansal
aman-bansal committed Jun 6, 2024
1 parent 7d11119 commit eb2d152
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 71 deletions.
113 changes: 58 additions & 55 deletions ...way-service-impl/src/main/java/org/hypertrace/gateway/service/explore/RequestHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -119,66 +120,12 @@ public ExploreResponse.Builder handleRequest(
attributeMetadataProvider);

if (request.getFetchTotal()) {
QueryRequest totalQueryRequest = buildTotalQueryRequest(request, queryRequest);
Iterator<ResultSetChunk> 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<org.hypertrace.gateway.service.v1.common.Expression> 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,
Expand Down Expand Up @@ -294,6 +241,62 @@ private List<String> getEntityIdsToFilterFromSourceEDS(
.collect(Collectors.toUnmodifiableList());
}

private int fetchTotal(
ExploreRequestContext requestContext, ExploreRequest request, QueryRequest queryRequest) {
QueryRequest totalQueryRequest = buildTotalQueryRequest(request, queryRequest);
Iterator<ResultSetChunk> 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<org.hypertrace.gateway.service.v1.common.Expression> groupBys =
ExpressionReader.getAttributeExpressions(request.getGroupByList());
// add all group-by columns as distinct count to get total
List<org.hypertrace.core.query.service.api.Expression> 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,
Expand Down
29 changes: 13 additions & 16 deletions ...c/main/java/org/hypertrace/gateway/service/explore/entity/EntityServiceEntityFetcher.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -197,27 +197,24 @@ private EntityQueryRequest buildTotalEntitiesRequest(
.setFilter(buildFilter(exploreRequest, entityIdAttributeIds, Collections.emptySet()));
List<org.hypertrace.gateway.service.v1.common.Expression> groupBys =
ExpressionReader.getAttributeExpressions(exploreRequest.getGroupByList());
List<Expression> 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());
Expand Down

0 comments on commit eb2d152

Please sign in to comment.