Examples with Document org.hypertrace.core.documentstore.Document used on opensource projects

Example 21 with Document

use of org.hypertrace.core.documentstore.Document in project entity-service by hypertrace.

the class EntityQueryServiceImplTest method testExecute_withAliases.

@Test
@Disabled("Disabled until we enable querying based on the new Query DTO")
public void testExecute_withAliases() throws Exception {
    Collection mockEntitiesCollection = mock(Collection.class);
    List<Document> docs = List.of(new JSONDocument("{\n" + "    \"tenantId\": \"tenant-1\",\n" + "    \"entityId\": \"" + UUID.randomUUID() + "\",\n" + "    \"entityType\": \"" + TEST_ENTITY_TYPE + "\",\n" + "    \"entityName\": \"Test entity 1\",\n" + "    \"col1\": \"col1-value\",\n" + "    \"Entity\": \n" + "    {\n" + "        \"status\":\n" + "        {\n" + "            \"value\":\n" + "            {\n" + "                \"string\": \"col2-value\"\n" + "            }\n" + "        }\n" + "    }\n" + "}"));
    when(mockEntitiesCollection.aggregate(any())).thenReturn(convertToCloseableIterator(docs.iterator()));
    when(mockEntitiesCollection.search(any())).thenReturn(convertToCloseableIterator(docs.iterator()));
    EntityQueryRequest request = EntityQueryRequest.newBuilder().setEntityType(TEST_ENTITY_TYPE).addSelection(Expression.newBuilder().setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(ATTRIBUTE_ID1).setAlias("col1"))).addSelection(Expression.newBuilder().setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(ATTRIBUTE_ID2))).build();
    StreamObserver<ResultSetChunk> mockResponseObserver = mock(StreamObserver.class);
    Context.current().withValue(RequestContext.CURRENT, mockRequestContextWithTenantId()).run(() -> {
        EntityQueryServiceImpl eqs = new EntityQueryServiceImpl(mockEntitiesCollection, mockMappingForAttributes1And2(), 100, false);
        eqs.execute(request, mockResponseObserver);
    });
    ResultSetChunk expectedResponse = ResultSetChunk.newBuilder().setIsLastChunk(true).setResultSetMetadata(ResultSetMetadata.newBuilder().addColumnMetadata(ColumnMetadata.newBuilder().setColumnName("col1")).addColumnMetadata(ColumnMetadata.newBuilder().setColumnName(ATTRIBUTE_ID2))).addRow(Row.newBuilder().addColumn(Value.newBuilder().setString("col1-value")).addColumn(Value.newBuilder().setString("col2-value")).build()).build();
    verify(mockResponseObserver, times(1)).onNext(eq(expectedResponse));
    verify(mockResponseObserver, times(1)).onCompleted();
}

Example 22 with Document

use of org.hypertrace.core.documentstore.Document in project entity-service by hypertrace.

the class EntityQueryServiceImplTest method testExecute_success.

@Test
public void testExecute_success() throws Exception {
    Collection mockEntitiesCollection = mock(Collection.class);
    Entity entity1 = Entity.newBuilder().setTenantId("tenant-1").setEntityType(TEST_ENTITY_TYPE).setEntityId(UUID.randomUUID().toString()).setEntityName("Test entity 1").putAttributes(EDS_COLUMN_NAME1, AttributeValue.newBuilder().setValue(org.hypertrace.entity.data.service.v1.Value.newBuilder().setString("foo1")).build()).build();
    Entity entity2 = Entity.newBuilder().setTenantId("tenant-1").setEntityType(TEST_ENTITY_TYPE).setEntityId(UUID.randomUUID().toString()).setEntityName("Test entity 2").putAttributes(EDS_COLUMN_NAME1, AttributeValue.newBuilder().setValue(org.hypertrace.entity.data.service.v1.Value.newBuilder().setString("foo2")).build()).build();
    List<Document> docs = List.of(new JSONDocument(JsonFormat.printer().print(entity1)), new JSONDocument(JsonFormat.printer().print(entity2)), // this doc will result in parsing error
    new JSONDocument("{\"entityId\": [1, 2]}"));
    when(mockEntitiesCollection.aggregate(any())).thenReturn(convertToCloseableIterator(docs.iterator()));
    when(mockEntitiesCollection.search(any())).thenReturn(convertToCloseableIterator(docs.iterator()));
    EntityQueryRequest request = EntityQueryRequest.newBuilder().setEntityType(TEST_ENTITY_TYPE).addOrderBy(OrderByExpression.newBuilder().setExpression(Expression.newBuilder().setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(ATTRIBUTE_ID1).build()))).build();
    StreamObserver<ResultSetChunk> mockResponseObserver = mock(StreamObserver.class);
    Context.current().withValue(RequestContext.CURRENT, mockRequestContextWithTenantId()).call(() -> {
        EntityQueryServiceImpl eqs = new EntityQueryServiceImpl(mockEntitiesCollection, mockMappingForAttributes1And2(), 1, false);
        eqs.execute(request, mockResponseObserver);
        return null;
    });
    verify(mockEntitiesCollection, times(0)).aggregate(any());
    verify(mockEntitiesCollection, times(1)).search(any());
    verify(mockResponseObserver, times(3)).onNext(any());
    verify(mockResponseObserver, times(1)).onCompleted();
}

Example 23 with Document

use of org.hypertrace.core.documentstore.Document in project entity-service by hypertrace.

the class EntityTypeServiceImpl method queryEntityTypes.

@Override
public void queryEntityTypes(EntityTypeFilter request, StreamObserver<EntityType> responseObserver) {
    Optional<String> tenantId = RequestContext.CURRENT.get().getTenantId();
    if (tenantId.isEmpty()) {
        responseObserver.onError(new ServiceException("Tenant id is missing in the request."));
        return;
    }
    Iterator<Document> entityTypes = entityTypeCol.search(transform(tenantId.get(), request, true));
    while (entityTypes.hasNext()) {
        Document entityType = entityTypes.next();
        EntityType.Builder builder = EntityType.newBuilder();
        try {
            PARSER.merge(entityType.toJson(), builder);
            // Populate the tenant id field with the tenant id that's received for backward
            // compatibility.
            builder.setTenantId(tenantId.get());
        } catch (InvalidProtocolBufferException e) {
            LOG.error("Error processing entityType: {}", entityType.toJson(), e);
            responseObserver.onError(e);
            return;
        }
        responseObserver.onNext(builder.build());
    }
    responseObserver.onCompleted();
}

Example 24 with Document

use of org.hypertrace.core.documentstore.Document in project entity-service by hypertrace.

the class EntityTypeServiceImpl method queryEntityTypes.

@Override
public void queryEntityTypes(org.hypertrace.entity.type.service.v2.QueryEntityTypesRequest request, io.grpc.stub.StreamObserver<org.hypertrace.entity.type.service.v2.QueryEntityTypesResponse> responseObserver) {
    Optional<String> tenantId = RequestContext.CURRENT.get().getTenantId();
    if (tenantId.isEmpty()) {
        responseObserver.onError(new ServiceException("Tenant id is missing in the request."));
        return;
    }
    Iterator<Document> entityTypes = entityTypeCollection.search(transform(tenantId.get(), request.getNameList()));
    List<EntityType> entityTypeList = new ArrayList<>();
    while (entityTypes.hasNext()) {
        String json = entityTypes.next().toJson();
        try {
            entityTypeList.add(EntityTypeDocument.fromJson(json).toProto());
        } catch (JsonProcessingException e) {
            LOGGER.warn("Failed to parse the EntityType json: {}", json, e);
            responseObserver.onError(new RuntimeException("Error querying the entity types."));
            return;
        }
    }
    responseObserver.onNext(QueryEntityTypesResponse.newBuilder().addAllEntityType(entityTypeList).build());
    responseObserver.onCompleted();
}

Example 25 with Document

use of org.hypertrace.core.documentstore.Document 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();
}