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

Example 1 with Builder

use of org.hypertrace.entity.query.service.v1.LiteralConstant.Builder 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 Builder

use of org.hypertrace.entity.query.service.v1.LiteralConstant.Builder 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;
}

Example 3 with Builder

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

the class EntityServiceAndGatewayServiceConverterTest method testAddBetweenFilter.

@Test
public void testAddBetweenFilter() {
    int startTimeMillis = 1;
    int endTimeMillis = 2;
    String timestamp = "lastActivity";
    String timestampAttributeName = BACKEND.name() + "." + timestamp;
    Expression.Builder expectedStartTimeConstant = Expression.newBuilder().setLiteral(LiteralConstant.newBuilder().setValue(Value.newBuilder().setValueType(ValueType.LONG).setLong(startTimeMillis)));
    Expression.Builder expectedEndTimeConstant = Expression.newBuilder().setLiteral(LiteralConstant.newBuilder().setValue(Value.newBuilder().setValueType(ValueType.LONG).setLong(endTimeMillis)));
    Filter.Builder expectedStartTimeFilterBuilder = Filter.newBuilder();
    expectedStartTimeFilterBuilder.setOperator(Operator.GE);
    expectedStartTimeFilterBuilder.setLhs(EntityServiceAndGatewayServiceConverter.createColumnExpression(timestampAttributeName));
    expectedStartTimeFilterBuilder.setRhs(expectedStartTimeConstant);
    Filter.Builder expectedEndTimeFilterBuilder = Filter.newBuilder();
    expectedEndTimeFilterBuilder.setOperator(Operator.LT);
    expectedEndTimeFilterBuilder.setLhs(EntityServiceAndGatewayServiceConverter.createColumnExpression(timestampAttributeName));
    expectedEndTimeFilterBuilder.setRhs(expectedEndTimeConstant);
    Filter.Builder expectedFilterBuilder = Filter.newBuilder();
    expectedFilterBuilder.setOperator(Operator.AND);
    expectedFilterBuilder.addChildFilter(expectedStartTimeFilterBuilder);
    expectedFilterBuilder.addChildFilter(expectedEndTimeFilterBuilder);
    AttributeMetadataProvider mockProvider = mock(AttributeMetadataProvider.class);
    when(mockProvider.getAttributeMetadata(any(EntitiesRequestContext.class), eq(BACKEND.name()), eq(timestamp))).thenReturn(Optional.of(AttributeMetadata.newBuilder().setId(timestampAttributeName).build()));
    EntityQueryRequest.Builder builder = EntityQueryRequest.newBuilder();
    EntitiesRequest request = EntitiesRequest.newBuilder().setEntityType(API.name()).build();
    EntityServiceAndGatewayServiceConverter.addBetweenTimeFilter(startTimeMillis, endTimeMillis, mockProvider, request, builder, mock(EntitiesRequestContext.class));
    // no filter added for unsupported entity types
    assertEquals(Filter.getDefaultInstance(), builder.getFilter());
    // if it's an backend, filters will be added
    request = EntitiesRequest.newBuilder().setEntityType(BACKEND.name()).build();
    EntityServiceAndGatewayServiceConverter.addBetweenTimeFilter(startTimeMillis, endTimeMillis, mockProvider, request, builder, mock(EntitiesRequestContext.class));
    assertEquals(expectedFilterBuilder.build(), builder.getFilter());
    // if there's an existing filter in EntitiesRequest, it will be preserved
    Expression.Builder scoreConstant = Expression.newBuilder().setLiteral(LiteralConstant.newBuilder().setValue(Value.newBuilder().setValueType(ValueType.INT).setInt(10)));
    Filter.Builder existingScoreFilter = Filter.newBuilder();
    existingScoreFilter.setOperator(Operator.GT);
    existingScoreFilter.setLhs(EntityServiceAndGatewayServiceConverter.createColumnExpression("TEST.score"));
    existingScoreFilter.setRhs(scoreConstant);
    expectedFilterBuilder.addChildFilter(existingScoreFilter);
    builder.setFilter(existingScoreFilter.build());
    EntityServiceAndGatewayServiceConverter.addBetweenTimeFilter(startTimeMillis, endTimeMillis, mockProvider, request, builder, mock(EntitiesRequestContext.class));
    assertEquals(expectedFilterBuilder.build(), builder.getFilter());
}

Example 4 with Builder

use of org.hypertrace.entity.query.service.v1.LiteralConstant.Builder in project entity-service by hypertrace.

the class EntityQueryServiceImplTest method testUpdate_success.

@Test
public void testUpdate_success() throws Exception {
    Collection mockEntitiesCollection = mockEntitiesCollection();
    Builder newStatus = LiteralConstant.newBuilder().setValue(Value.newBuilder().setValueType(ValueType.STRING).setString("NEW_STATUS"));
    EntityUpdateRequest updateRequest = EntityUpdateRequest.newBuilder().setEntityType(TEST_ENTITY_TYPE).addEntityIds("entity-id-1").setOperation(UpdateOperation.newBuilder().setSetAttribute(SetAttribute.newBuilder().setAttribute(ColumnIdentifier.newBuilder().setColumnName(ATTRIBUTE_ID2)).setValue(newStatus))).addSelection(Expression.newBuilder().setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(ATTRIBUTE_ID1))).addSelection(Expression.newBuilder().setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(ATTRIBUTE_ID2))).build();
    StreamObserver<ResultSetChunk> mockResponseObserver = mock(StreamObserver.class);
    Context.current().withValue(RequestContext.CURRENT, mockRequestContextWithTenantId()).call(() -> {
        EntityQueryServiceImpl eqs = new EntityQueryServiceImpl(mockEntitiesCollection, mockMappingForAttributes1And2(), 1, false);
        eqs.update(updateRequest, mockResponseObserver);
        return null;
    });
    verify(mockEntitiesCollection, times(1)).bulkUpdateSubDocs(eq(Map.of(new SingleValueKey("tenant1", "entity-id-1"), Map.of("attributes.status", new JSONDocument(DocStoreJsonFormat.printer().print(newStatus))))));
}

Example 5 with Builder

use of org.hypertrace.entity.query.service.v1.LiteralConstant.Builder in project entity-service by hypertrace.

the class EntityQueryConverter method convertToAttributeFilter.

private AttributeFilter convertToAttributeFilter(RequestContext requestContext, Query.Builder queryBuilder, Filter filter) {
    if (filter == null || filter.equals(Filter.getDefaultInstance())) {
        return null;
    }
    AttributeFilter.Builder builder = null;
    if (filter.getChildFilterCount() == 0) {
        // Copy the lhs and rhs from the filter.
        AttributeColumnInformation attributeData = getAttributeColumnInfo(requestContext, filter.getLhs());
        String edsColumnName = attributeData.getColumnName();
        org.hypertrace.entity.query.service.v1.Value rhsValue = filter.getRhs().getLiteral().getValue();
        if (edsColumnName.equals(EntityServiceConstants.ENTITY_ID)) {
            if (rhsValue.getValueType() == ValueType.STRING) {
                queryBuilder.addEntityId(rhsValue.getString());
            } else if (rhsValue.getValueType() == ValueType.STRING_ARRAY) {
                queryBuilder.addAllEntityId(rhsValue.getStringArrayList());
            }
        } else if (edsColumnName.equals(EntityServiceConstants.ENTITY_NAME)) {
            queryBuilder.setEntityName(rhsValue.getString());
        } else {
            builder = AttributeFilter.newBuilder();
            builder.setOperator(convertOperator(filter.getOperator()));
            builder.setName(edsColumnName);
            builder.setIsMultiValued(attributeData.isMultiValued());
            builder.setAttributeValue(convertToAttributeValue(filter.getRhs()));
        }
    } else {
        builder = AttributeFilter.newBuilder();
        builder.setOperator(convertOperator(filter.getOperator()));
        for (Filter child : filter.getChildFilterList()) {
            AttributeFilter attributeFilter = convertToAttributeFilter(requestContext, queryBuilder, child);
            if (null != attributeFilter) {
                builder.addChildFilter(attributeFilter);
            }
        }
    }
    return (null != builder) ? builder.build() : null;
}