Examples with RequestContext org.hypertrace.core.grpcutils.context.RequestContext used on opensource projects

Example 6 with RequestContext

use of org.hypertrace.core.grpcutils.context.RequestContext in project entity-service by hypertrace.

the class EntityQueryServiceImpl method execute.

@Override
public void execute(EntityQueryRequest request, StreamObserver<ResultSetChunk> responseObserver) {
    RequestContext requestContext = RequestContext.CURRENT.get();
    Optional<String> tenantId = requestContext.getTenantId();
    if (tenantId.isEmpty()) {
        responseObserver.onError(new ServiceException("Tenant id is missing in the request."));
        return;
    }
    final Iterator<Document> documentIterator;
    if (queryAggregationEnabled) {
        final Converter<EntityQueryRequest, org.hypertrace.core.documentstore.query.Query> queryConverter = getQueryConverter();
        final org.hypertrace.core.documentstore.query.Query query;
        try {
            query = queryConverter.convert(request, requestContext);
            documentIterator = entitiesCollection.aggregate(query);
        } catch (final Exception e) {
            responseObserver.onError(new ServiceException(e));
            return;
        }
    } else {
        // TODO: Optimize this later. For now converting to EDS Query and then again to DocStore
        // Query.
        Query query = entityQueryConverter.convertToEDSQuery(requestContext, request);
        /**
         * {@link EntityQueryRequest} selections need to treated differently, since they don't
         * transform one to one to {@link org.hypertrace.entity.data.service.v1.EntityDataRequest}
         * selections
         */
        List<String> docStoreSelections = entityQueryConverter.convertSelectionsToDocStoreSelections(requestContext, request.getSelectionList());
        documentIterator = entitiesCollection.search(DocStoreConverter.transform(tenantId.get(), query, docStoreSelections));
    }
    final DocumentConverter rowConverter = injector.getInstance(DocumentConverter.class);
    ResultSetMetadata resultSetMetadata;
    try {
        resultSetMetadata = this.buildMetadataForSelections(request.getSelectionList());
    } catch (final ConversionException e) {
        responseObserver.onError(new ServiceException(e));
        return;
    }
    if (!documentIterator.hasNext()) {
        ResultSetChunk.Builder resultBuilder = ResultSetChunk.newBuilder();
        resultBuilder.setResultSetMetadata(resultSetMetadata);
        resultBuilder.setIsLastChunk(true);
        resultBuilder.setChunkId(0);
        responseObserver.onNext(resultBuilder.build());
        responseObserver.onCompleted();
        return;
    }
    boolean isNewChunk = true;
    int chunkId = 0, rowCount = 0;
    ResultSetChunk.Builder resultBuilder = ResultSetChunk.newBuilder();
    while (documentIterator.hasNext()) {
        // Set metadata for new chunk
        if (isNewChunk) {
            resultBuilder.setResultSetMetadata(resultSetMetadata);
            isNewChunk = false;
        }
        try {
            final Row row;
            if (queryAggregationEnabled) {
                row = rowConverter.convertToRow(documentIterator.next(), resultSetMetadata);
                resultBuilder.addRow(row);
                rowCount++;
            } else {
                Optional<Entity> entity = DOCUMENT_PARSER.parseOrLog(documentIterator.next(), Entity.newBuilder());
                if (entity.isPresent()) {
                    row = convertToEntityQueryResult(requestContext, entity.get(), request.getSelectionList());
                    resultBuilder.addRow(row);
                    rowCount++;
                }
            }
        } catch (final Exception e) {
            responseObserver.onError(new ServiceException(e));
            return;
        }
        // current chunk is complete
        if (rowCount >= CHUNK_SIZE || !documentIterator.hasNext()) {
            resultBuilder.setChunkId(chunkId++);
            resultBuilder.setIsLastChunk(!documentIterator.hasNext());
            responseObserver.onNext(resultBuilder.build());
            resultBuilder = ResultSetChunk.newBuilder();
            isNewChunk = true;
            rowCount = 0;
        }
    }
    responseObserver.onCompleted();
}

Example 7 with RequestContext

use of org.hypertrace.core.grpcutils.context.RequestContext in project entity-service by hypertrace.

the class EntityQueryServiceImpl method total.

@Override
public void total(TotalEntitiesRequest request, StreamObserver<TotalEntitiesResponse> responseObserver) {
    RequestContext requestContext = RequestContext.CURRENT.get();
    Optional<String> tenantId = requestContext.getTenantId();
    if (tenantId.isEmpty()) {
        responseObserver.onError(new ServiceException("Tenant id is missing in the request."));
        return;
    }
    // converting total entities request to entity query request
    EntityQueryRequest entityQueryRequest = EntityQueryRequest.newBuilder().setEntityType(request.getEntityType()).setFilter(request.getFilter()).build();
    // converting entity query request to entity data service query
    Query query = entityQueryConverter.convertToEDSQuery(requestContext, entityQueryRequest);
    // TODO: Replace to use the new org.hypertrace.core.documentstore.query.Query DTO
    long total = entitiesCollection.total(DocStoreConverter.transform(tenantId.get(), query, emptyList()));
    responseObserver.onNext(TotalEntitiesResponse.newBuilder().setTotal(total).build());
    responseObserver.onCompleted();
}

Example 8 with RequestContext

use of org.hypertrace.core.grpcutils.context.RequestContext in project entity-service by hypertrace.

the class EntityQueryServiceImpl method bulkUpdate.

@Override
public void bulkUpdate(BulkEntityUpdateRequest request, StreamObserver<ResultSetChunk> responseObserver) {
    // Validations
    RequestContext requestContext = RequestContext.CURRENT.get();
    Optional<String> maybeTenantId = requestContext.getTenantId();
    if (maybeTenantId.isEmpty()) {
        responseObserver.onError(new ServiceException("Tenant id is missing in the request."));
        return;
    }
    if (StringUtils.isEmpty(request.getEntityType())) {
        responseObserver.onError(new ServiceException("Entity type is missing in the request."));
        return;
    }
    if (request.getEntitiesCount() == 0) {
        responseObserver.onError(new ServiceException("Entities are missing in the request."));
    }
    Map<String, EntityUpdateInfo> entitiesMap = request.getEntitiesMap();
    try {
        doBulkUpdate(requestContext, entitiesMap);
        responseObserver.onCompleted();
    } catch (Exception e) {
        responseObserver.onError(new ServiceException("Error occurred while executing " + request, e));
    }
}

Example 9 with RequestContext

use of org.hypertrace.core.grpcutils.context.RequestContext in project entity-service by hypertrace.

the class EntityQueryServiceImpl method bulkUpdateEntityArrayAttribute.

@Override
public void bulkUpdateEntityArrayAttribute(BulkEntityArrayAttributeUpdateRequest request, StreamObserver<BulkEntityArrayAttributeUpdateResponse> responseObserver) {
    RequestContext requestContext = RequestContext.CURRENT.get();
    String tenantId = requestContext.getTenantId().orElse(null);
    if (isNull(tenantId)) {
        responseObserver.onError(new ServiceException("Tenant id is missing in the request."));
        return;
    }
    try {
        Set<Key> keys = request.getEntityIdsList().stream().map(entityId -> new SingleValueKey(tenantId, entityId)).collect(Collectors.toCollection(LinkedHashSet::new));
        String attributeId = request.getAttribute().getColumnName();
        String subDocPath = entityAttributeMapping.getDocStorePathByAttributeId(requestContext, attributeId).orElseThrow(() -> new IllegalArgumentException("Unknown attribute " + attributeId));
        List<Document> subDocuments = request.getValuesList().stream().map(this::convertToJsonDocument).collect(toUnmodifiableList());
        BulkArrayValueUpdateRequest bulkArrayValueUpdateRequest = new BulkArrayValueUpdateRequest(keys, subDocPath + ARRAY_VALUE_PATH_SUFFIX, getMatchingOperation(request.getOperation()), subDocuments);
        entitiesCollection.bulkOperationOnArrayValue(bulkArrayValueUpdateRequest);
        responseObserver.onNext(BulkEntityArrayAttributeUpdateResponse.newBuilder().build());
        responseObserver.onCompleted();
    } catch (Exception e) {
        responseObserver.onError(e);
    }
}

Example 10 with RequestContext

use of org.hypertrace.core.grpcutils.context.RequestContext in project entity-service by hypertrace.

the class EntityQueryServiceImpl method update.

@Override
public void update(EntityUpdateRequest request, StreamObserver<ResultSetChunk> responseObserver) {
    // Validations
    RequestContext requestContext = RequestContext.CURRENT.get();
    Optional<String> maybeTenantId = requestContext.getTenantId();
    if (maybeTenantId.isEmpty()) {
        responseObserver.onError(new ServiceException("Tenant id is missing in the request."));
        return;
    }
    if (StringUtils.isEmpty(request.getEntityType())) {
        responseObserver.onError(new ServiceException("Entity type is missing in the request."));
        return;
    }
    if (request.getEntityIdsCount() == 0) {
        responseObserver.onError(new ServiceException("Entity IDs are missing in the request."));
    }
    if (!request.hasOperation()) {
        responseObserver.onError(new ServiceException("Operation is missing in the request."));
    }
    try {
        // Execute the update
        doUpdate(requestContext, request);
        // Finally, return the selections
        List<Document> documents = getProjectedDocuments(request.getEntityIdsList(), request.getSelectionList(), requestContext);
        responseObserver.onNext(convertDocumentsToResultSetChunk(documents, request.getSelectionList()));
        responseObserver.onCompleted();
    } catch (Exception e) {
        responseObserver.onError(new ServiceException("Error occurred while executing " + request, e));
    }
}