Examples with ResultSetChunk org.hypertrace.entity.query.service.v1.ResultSetChunk used on opensource projects

Example 1 with ResultSetChunk

use of org.hypertrace.entity.query.service.v1.ResultSetChunk in project gateway-service by hypertrace.

the class EntityDataServiceEntityFetcher method getEntities.

@Override
public EntityFetcherResponse getEntities(EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest) {
    List<String> entityIdAttributeIds = AttributeMetadataUtil.getIdAttributeIds(attributeMetadataProvider, entityIdColumnsConfigs, requestContext, entitiesRequest.getEntityType());
    Map<String, List<String>> requestedAliasesByEntityIdAttributeIds = getExpectedResultNamesForEachAttributeId(entitiesRequest.getSelectionList(), entityIdAttributeIds);
    EntityQueryRequest.Builder builder = EntityQueryRequest.newBuilder().setEntityType(entitiesRequest.getEntityType()).setFilter(EntityServiceAndGatewayServiceConverter.convertToEntityServiceFilter(entitiesRequest.getFilter())).addAllSelection(entityIdAttributeIds.stream().map(entityIdAttr -> EntityServiceAndGatewayServiceConverter.createColumnExpression(entityIdAttr).build()).collect(Collectors.toList()));
    // add time filter for supported scope
    if (!entitiesRequest.getIncludeNonLiveEntities()) {
        EntityServiceAndGatewayServiceConverter.addBetweenTimeFilter(entitiesRequest.getStartTimeMillis(), entitiesRequest.getEndTimeMillis(), attributeMetadataProvider, entitiesRequest, builder, requestContext);
    }
    // Add all expressions in the select that are already not part of the EntityID attributes
    entitiesRequest.getSelectionList().stream().filter(ExpressionReader::isAttributeSelection).filter(expression -> ExpressionReader.getAttributeIdFromAttributeSelection(expression).map(attributeId -> !entityIdAttributeIds.contains(attributeId)).orElse(true)).forEach(expression -> builder.addSelection(EntityServiceAndGatewayServiceConverter.convertToEntityServiceExpression(expression)));
    int limit = entitiesRequest.getLimit();
    if (limit > 0) {
        builder.setLimit(limit);
    }
    int offset = entitiesRequest.getOffset();
    if (offset > 0) {
        builder.setOffset(offset);
    }
    if (!entitiesRequest.getOrderByList().isEmpty()) {
        builder.addAllOrderBy(EntityServiceAndGatewayServiceConverter.convertToOrderByExpressions(entitiesRequest.getOrderByList()));
    }
    EntityQueryRequest entityQueryRequest = builder.build();
    LOG.debug("Sending Query to EDS  ======== \n {}", entityQueryRequest);
    Iterator<ResultSetChunk> resultSetChunkIterator = entityQueryServiceClient.execute(builder.build(), requestContext.getHeaders());
    Map<String, AttributeMetadata> resultMetadataMap = this.getAttributeMetadataByAlias(requestContext, entitiesRequest);
    // We want to retain the order as returned from the respective source. Hence using a
    // LinkedHashMap
    Map<EntityKey, Builder> entityBuilders = new LinkedHashMap<>();
    while (resultSetChunkIterator.hasNext()) {
        ResultSetChunk chunk = resultSetChunkIterator.next();
        LOG.debug("Received chunk: {}", chunk);
        if (chunk.getRowCount() < 1) {
            break;
        }
        for (Row row : chunk.getRowList()) {
            // Construct the entity id from the entityIdAttributes columns
            EntityKey entityKey = EntityKey.of(IntStream.range(0, entityIdAttributeIds.size()).mapToObj(value -> row.getColumn(value).getString()).toArray(String[]::new));
            Builder entityBuilder = entityBuilders.computeIfAbsent(entityKey, k -> Entity.newBuilder());
            entityBuilder.setEntityType(entitiesRequest.getEntityType());
            entityBuilder.setId(entityKey.toString());
            // as post processing.
            for (int i = 0; i < entityIdAttributeIds.size(); i++) {
                entityBuilder.putAttribute(entityIdAttributeIds.get(i), Value.newBuilder().setString(entityKey.getAttributes().get(i)).setValueType(ValueType.STRING).build());
            }
            requestedAliasesByEntityIdAttributeIds.forEach((attributeId, requestedAliasList) -> requestedAliasList.forEach(requestedAlias -> entityBuilder.putAttribute(requestedAlias, entityBuilder.getAttributeOrThrow(attributeId))));
            for (int i = entityIdAttributeIds.size(); i < chunk.getResultSetMetadata().getColumnMetadataCount(); i++) {
                String resultName = chunk.getResultSetMetadata().getColumnMetadata(i).getColumnName();
                AttributeMetadata attributeMetadata = resultMetadataMap.get(resultName);
                entityBuilder.putAttribute(resultName, EntityServiceAndGatewayServiceConverter.convertQueryValueToGatewayValue(row.getColumn(i), attributeMetadata));
            }
        }
    }
    return new EntityFetcherResponse(entityBuilders);
}

Example 2 with ResultSetChunk

use of org.hypertrace.entity.query.service.v1.ResultSetChunk in project gateway-service by hypertrace.

the class EdsEntityUpdaterTest method nonExistentEntity.

@Test
public void nonExistentEntity() {
    EntityQueryServiceClient mockEqsClient = mock(EntityQueryServiceClient.class);
    EdsEntityUpdater entityUpdater = new EdsEntityUpdater(mockEqsClient);
    // mock empty ResultSetChunk
    UpdateEntityRequest updateEntityRequest = createResolveTestRequest("non-existent-id");
    ResultSetChunk emptyResult = ResultSetChunk.newBuilder().build();
    when(mockEqsClient.update(any(), any())).thenReturn(List.of(emptyResult).iterator());
    UpdateExecutionContext updateExecutionContext = new UpdateExecutionContext(null, mockTestMetadata());
    UpdateEntityResponse response = entityUpdater.update(updateEntityRequest, updateExecutionContext).build();
    assertFalse(response.hasEntity());
}

Example 3 with ResultSetChunk

use of org.hypertrace.entity.query.service.v1.ResultSetChunk in project gateway-service by hypertrace.

the class EntityDataServiceEntityFetcherTests method test_getEntities_WithoutPagination.

@Test
public void test_getEntities_WithoutPagination() {
    List<OrderByExpression> orderByExpressions = List.of(buildOrderByExpression(API_ID_ATTR));
    long startTime = 1L;
    long endTime = 10L;
    int limit = 0;
    int offset = 0;
    String tenantId = "TENANT_ID";
    Map<String, String> requestHeaders = Map.of("x-tenant-id", tenantId);
    AttributeScope entityType = AttributeScope.API;
    EntitiesRequest entitiesRequest = EntitiesRequest.newBuilder().setEntityType(entityType.name()).setStartTimeMillis(startTime).setEndTimeMillis(endTime).setFilter(Filter.newBuilder().setOperator(AND).addChildFilter(generateEQFilter(API_TYPE_ATTR, "HTTP")).addChildFilter(generateEQFilter(API_NAME_ATTR, "DISCOVERED"))).addAllOrderBy(orderByExpressions).setLimit(limit).setOffset(offset).build();
    EntitiesRequestContext entitiesRequestContext = new EntitiesRequestContext(tenantId, startTime, endTime, entityType.name(), "API.startTime", requestHeaders);
    EntityQueryRequest expectedQueryRequest = EntityQueryRequest.newBuilder().setEntityType("API").addSelection(convertToEntityServiceExpression(Expression.newBuilder().setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(API_ID_ATTR)).build())).setFilter(convertToEntityServiceFilter(entitiesRequest.getFilter())).addAllOrderBy(EntityServiceAndGatewayServiceConverter.convertToOrderByExpressions(orderByExpressions)).build();
    List<ResultSetChunk> resultSetChunks = List.of(getResultSetChunk(List.of("API.apiId"), new String[][] { { "apiId1" }, { "apiId2" } }));
    when(entityQueryServiceClient.execute(eq(expectedQueryRequest), eq(requestHeaders))).thenReturn(resultSetChunks.iterator());
    assertEquals(2, entityDataServiceEntityFetcher.getEntities(entitiesRequestContext, entitiesRequest).size());
}

Example 4 with ResultSetChunk

use of org.hypertrace.entity.query.service.v1.ResultSetChunk in project gateway-service by hypertrace.

the class EdsEntityUpdater method update.

public UpdateEntityResponse.Builder update(UpdateEntityRequest updateRequest, UpdateExecutionContext updateExecutionContext) {
    UpdateEntityResponse.Builder responseBuilder = UpdateEntityResponse.newBuilder();
    EntityUpdateRequest eqsUpdateRequest = convertToEqsUpdateRequest(updateRequest);
    if (LOG.isDebugEnabled()) {
        LOG.debug("Sending update request to EDS  ======== \n {}", eqsUpdateRequest);
    }
    Iterator<ResultSetChunk> resultSetChunkIterator = eqsClient.update(eqsUpdateRequest, updateExecutionContext.getRequestHeaders());
    if (!resultSetChunkIterator.hasNext()) {
        return responseBuilder;
    }
    ResultSetChunk chunk = resultSetChunkIterator.next();
    if (LOG.isDebugEnabled()) {
        LOG.debug("Received chunk: " + chunk.toString());
    }
    if (chunk.getRowCount() == 0) {
        return responseBuilder;
    }
    if (chunk.getRowCount() > 1) {
        LOG.warn("Received more than 1 row back. Only the first out of {} rows will be returned in the response", chunk.getRowCount());
    }
    Map<String, AttributeMetadata> resultAttributeMetadata = AttributeMetadataUtil.remapAttributeMetadataByResultKey(updateRequest.getSelectionList(), updateExecutionContext.getAttributeMetadata());
    Row row = chunk.getRow(0);
    Entity.Builder entityBuilder = Entity.newBuilder().setEntityType(updateRequest.getEntityType());
    for (int i = 0; i < chunk.getResultSetMetadata().getColumnMetadataCount(); i++) {
        String resultName = chunk.getResultSetMetadata().getColumnMetadata(i).getColumnName();
        entityBuilder.putAttribute(resultName, EntityServiceAndGatewayServiceConverter.convertQueryValueToGatewayValue(row.getColumn(i), resultAttributeMetadata.get(resultName)));
    }
    responseBuilder.setEntity(entityBuilder);
    return responseBuilder;
}

Example 5 with ResultSetChunk

use of org.hypertrace.entity.query.service.v1.ResultSetChunk in project gateway-service by hypertrace.

the class EntityRequestHandler method handleRequest.

@Override
public ExploreResponse.Builder handleRequest(ExploreRequestContext requestContext, ExploreRequest exploreRequest) {
    // ourselves.
    if (!exploreRequest.getGroupByList().isEmpty()) {
        requestContext.setHasGroupBy(true);
    }
    Set<String> entityIds = getEntityIds(requestContext, exploreRequest);
    Iterator<ResultSetChunk> resultSetChunkIterator = entityServiceEntityFetcher.getResults(requestContext, exploreRequest, entityIds);
    ExploreResponse.Builder builder = ExploreResponse.newBuilder();
    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));
    }
    // 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);
    }
    return builder;
}