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 committed Jun 25, 2024
1 parent 9753432 commit b16a950
Show file tree
Hide file tree
Showing 5 changed files with 123 additions and 19 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,9 @@ public QueryNode build() {
executionContext.setSortAndPaginationNodeAdded(true);
rootNode.acceptVisitor(new ExecutionContextBuilderVisitor(executionContext));
QueryNode executionTree = buildExecutionTree(executionContext, rootNode);
LOG.info("Execution Tree:{}", executionTree.acceptVisitor(new PrintVisitor()));
if (LOG.isDebugEnabled()) {
LOG.debug("Execution Tree:{}", executionTree.acceptVisitor(new PrintVisitor()));
}
return executionTree;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
Expand Down Expand Up @@ -94,15 +95,23 @@ protected static EntityResponse intersect(List<EntityResponse> entityResponses)
return new EntityResponse(entityFetcherResponse, entityFetcherResponse.size());
}

private static void mergeEntities(
private static EntityFetcherResponse mergeEntities(
EntityFetcherResponse rootEntityResponse, EntityFetcherResponse otherResponse) {
otherResponse
.getEntityKeyBuilderMap()
.forEach(
(key, value) ->
rootEntityResponse
.getEntityKeyBuilderMap()
.computeIfPresent(key, (k, v) -> v.mergeFrom(value.build())));
return new EntityFetcherResponse(
rootEntityResponse.getEntityKeyBuilderMap().entrySet().stream()
.collect(
Collectors.toUnmodifiableMap(
Entry::getKey,
entry -> {
Map<EntityKey, Builder> entityKeyBuilderMap =
otherResponse.getEntityKeyBuilderMap();
if (entityKeyBuilderMap.containsKey(entry.getKey())) {
return entry
.getValue()
.mergeFrom(entityKeyBuilderMap.get(entry.getKey()).build());
}
return entry.getValue();
})));
}

private static EntityFetcherResponse unionEntities(List<EntityFetcherResponse> builders) {
Expand Down Expand Up @@ -302,8 +311,9 @@ public EntityResponse visit(JoinNode joinNode) {
EntityFetcherResponse response =
resultMapList.stream()
.reduce(new EntityFetcherResponse(), (r1, r2) -> unionEntities(Arrays.asList(r1, r2)));
mergeEntities(childEntityFetcherResponse, response);
return new EntityResponse(childEntityFetcherResponse, childNodeResponse.getTotal());

return new EntityResponse(
mergeEntities(childEntityFetcherResponse, response), childNodeResponse.getTotal());
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,10 @@ public static Value getStringValue(String value) {
return Value.newBuilder().setString(value).setValueType(ValueType.STRING).build();
}

public static Value getLongValue(long value) {
return Value.newBuilder().setLong(value).setValueType(ValueType.LONG).build();
}

public static AggregatedMetricValue getAggregatedMetricValue(
FunctionType functionType, double value) {
return AggregatedMetricValue.newBuilder()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import static org.hypertrace.gateway.service.common.EntitiesRequestAndResponseUtils.compareEntityResponses;
import static org.hypertrace.gateway.service.common.EntitiesRequestAndResponseUtils.generateEQFilter;
import static org.hypertrace.gateway.service.common.EntitiesRequestAndResponseUtils.getAggregatedMetricValue;
import static org.hypertrace.gateway.service.common.EntitiesRequestAndResponseUtils.getLongValue;
import static org.hypertrace.gateway.service.common.EntitiesRequestAndResponseUtils.getStringValue;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
Expand Down Expand Up @@ -43,6 +44,7 @@
import org.hypertrace.gateway.service.entity.EntityQueryHandlerRegistry;
import org.hypertrace.gateway.service.entity.query.DataFetcherNode;
import org.hypertrace.gateway.service.entity.query.EntityExecutionContext;
import org.hypertrace.gateway.service.entity.query.JoinNode;
import org.hypertrace.gateway.service.entity.query.NoOpNode;
import org.hypertrace.gateway.service.entity.query.PaginateOnlyNode;
import org.hypertrace.gateway.service.entity.query.SelectionNode;
Expand All @@ -58,6 +60,7 @@
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.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;
Expand Down Expand Up @@ -380,6 +383,99 @@ public void testConstructFilterFromChildNodesNonEmptyResultsMultipleEntityIdExpr
new HashSet<>(filter.getChildFilterList()));
}

@Test
public void test_visit_JoinNode() {
List<OrderByExpression> orderByExpressions =
List.of(
buildOrderByExpression(
SortOrder.DESC,
buildAggregateExpression(
API_NUM_CALLS_ATTR, FunctionType.COUNT, "numCalls", List.of())));
int limit = 10;
int offset = 0;
long startTime = 0;
long endTime = 10;
String tenantId = "TENANT_ID";
AttributeScope entityType = AttributeScope.API;
Expression selectionExpression = buildExpression(API_NAME_ATTR);
EntitiesRequest entitiesRequest =
EntitiesRequest.newBuilder()
.setEntityType(entityType.name())
.setStartTimeMillis(startTime)
.setEndTimeMillis(endTime)
.addSelection(orderByExpressions.get(0).getExpression())
.addSelection(selectionExpression)
.setFilter(generateEQFilter(API_DISCOVERY_STATE, "DISCOVERED"))
.addAllOrderBy(orderByExpressions)
.setIncludeNonLiveEntities(true)
.setLimit(limit)
.setOffset(offset)
.build();
EntitiesRequestContext entitiesRequestContext =
new EntitiesRequestContext(
forTenantId(tenantId), startTime, endTime, entityType.name(), "API.startTime");
Map<EntityKey, Builder> entityKeyBuilderResponseMap =
Map.of(
EntityKey.of("entity-id-0"),
Entity.newBuilder().putAttribute("API.name", getStringValue("entity-0")),
EntityKey.of("entity-id-1"),
Entity.newBuilder().putAttribute("API.name", getStringValue("entity-1")),
EntityKey.of("entity-id-2"),
Entity.newBuilder().putAttribute("API.name", getStringValue("entity-2")));

EntityFetcherResponse entityFetcherResponse =
new EntityFetcherResponse(entityKeyBuilderResponseMap);

when(expressionContext.getSourceToMetricExpressionMap())
.thenReturn(Map.of("QS", List.of(orderByExpressions.get(0).getExpression())));
when(expressionContext.getSourceToSelectionExpressionMap())
.thenReturn(Map.of("EDS", List.of(selectionExpression)));
when(expressionContext.getSourceToSelectionExpressionMap())
.thenReturn(Map.of("EDS", List.of(selectionExpression)));
when(executionContext.getEntitiesRequest()).thenReturn(entitiesRequest);
when(executionContext.getTenantId()).thenReturn(tenantId);
when(executionContext.getEntitiesRequestContext()).thenReturn(entitiesRequestContext);
when(executionContext.getTimestampAttributeId()).thenReturn("API.startTime");

Map<EntityKey, Builder> queryServiceEntityKeyBuilderResponseMap =
Map.of(
EntityKey.of("entity-id-0"),
Entity.newBuilder().putAttribute("API.numCalls", getLongValue(10)),
EntityKey.of("entity-id-1"),
Entity.newBuilder().putAttribute("API.numCalls", getLongValue(50)),
EntityKey.of("entity-id-3"),
Entity.newBuilder().putAttribute("API.numCalls", getLongValue(100)));

when(queryServiceEntityFetcher.getEntities(eq(entitiesRequestContext), any()))
.thenReturn(new EntityFetcherResponse(queryServiceEntityKeyBuilderResponseMap));

NoOpNode mockChildNode = mock(NoOpNode.class);
when(mockChildNode.acceptVisitor(any()))
.thenReturn(new EntityResponse(entityFetcherResponse, 100));
JoinNode joinNode =
new JoinNode.Builder(mockChildNode)
.setAttrSelectionSources(Collections.emptySet())
.setAggMetricSelectionSources(Set.of(QS_SOURCE))
.build();

Map<EntityKey, Builder> mergedEntityKeyBuilderResponseMap =
Map.of(
EntityKey.of("entity-id-0"),
Entity.newBuilder()
.putAttribute("API.name", getStringValue("entity-0"))
.putAttribute("API.numCalls", getLongValue(10)),
EntityKey.of("entity-id-1"),
Entity.newBuilder()
.putAttribute("API.name", getStringValue("entity-1"))
.putAttribute("API.numCalls", getLongValue(50)),
EntityKey.of("entity-id-2"),
Entity.newBuilder().putAttribute("API.name", getStringValue("entity-2")));

compareEntityResponses(
new EntityResponse(new EntityFetcherResponse(mergedEntityKeyBuilderResponseMap), 100L),
executionVisitor.visit(joinNode));
}

@Test
public void test_visitDataFetcherNodeQs() {
List<OrderByExpression> orderByExpressions = List.of(buildOrderByExpression(API_ID_ATTR));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -47,10 +47,6 @@ timestamp.config = [
{
scope = LOG_EVENT
timestamp = timestamp
},
{
scope = ENVIRONMENT
timestamp = lastSeen
}
]

Expand All @@ -71,10 +67,6 @@ entity.idcolumn.config = [
scope = BACKEND
key = id
},
{
scope = ENVIRONMENT
key = id
},
]

filter.entity = {
Expand Down

0 comments on commit b16a950

Please sign in to comment.