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

Example 6 with JSONDocument

use of org.hypertrace.core.documentstore.JSONDocument 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 7 with JSONDocument

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

the class EntityTypeServiceImpl method upsertEntityRelationshipType.

@Override
public void upsertEntityRelationshipType(EntityRelationshipType request, StreamObserver<EntityRelationshipType> responseObserver) {
    Optional<String> tenantId = RequestContext.CURRENT.get().getTenantId();
    if (tenantId.isEmpty()) {
        responseObserver.onError(new ServiceException("Tenant id is missing in the request."));
        return;
    }
    try {
        // TODO: Since we currently use the same object that's received in the request while
        // persisting in the doc store, we need to add tenantId explicitly here.
        // We need to split these objects later.
        EntityRelationshipType newRequest = EntityRelationshipType.newBuilder(request).setTenantId(tenantId.get()).build();
        entityTypeRelationsCol.upsert(new SingleValueKey(tenantId.get(), request.getName()), new JSONDocument(PROTO_PRINTER.print(newRequest)));
        responseObserver.onNext(newRequest);
        responseObserver.onCompleted();
    } catch (IOException ioe) {
        responseObserver.onError(new RuntimeException("Error upserting EntityType:" + request, ioe));
    }
}

Example 8 with JSONDocument

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

the class DocStoreConverterTest method entityToJSONDocumentConversion64BitNumberFixup.

@Test
public void entityToJSONDocumentConversion64BitNumberFixup() throws IOException {
    // Verify that timestamp attribute which is a long number won't get
    // converted to string
    Entity testEntity = Entity.newBuilder().setTenantId("tenant1").setEntityName("myentity1").putIdentifyingAttributes("entity_id", AttributeValue.newBuilder().setValue(Value.newBuilder().setString("my-entity-id-1")).build()).putAttributes("timestamp", AttributeValue.newBuilder().setValue(Value.newBuilder().setLong(1584055141072L)).build()).build();
    JSONDocument jsonDocument = DocStoreConverter.transform(testEntity);
    Assertions.assertEquals("{\"tenantId\":\"tenant1\",\"entityName\":\"myentity1\",\"identifyingAttributes\":{\"entity_id\":{\"value\":{\"string\":\"my-entity-id-1\"}}},\"attributes\":{\"timestamp\":{\"value\":{\"long\":1584055141072}}}}", jsonDocument.toJson());
}

Example 9 with JSONDocument

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

the class EntityQueryServiceImpl method doUpdate.

private void doUpdate(RequestContext requestContext, EntityUpdateRequest request) throws Exception {
    if (request.getOperation().hasSetAttribute()) {
        SetAttribute setAttribute = request.getOperation().getSetAttribute();
        String attributeId = setAttribute.getAttribute().getColumnName();
        String subDocPath = entityAttributeMapping.getDocStorePathByAttributeId(requestContext, attributeId).orElseThrow(() -> new IllegalArgumentException("Unknown attribute FQN " + attributeId));
        JSONDocument jsonDocument = convertToJsonDocument(setAttribute.getValue());
        Map<Key, Map<String, Document>> entitiesUpdateMap = new HashMap<>();
        for (String entityId : request.getEntityIdsList()) {
            SingleValueKey key = new SingleValueKey(requestContext.getTenantId().orElseThrow(), entityId);
            if (entitiesUpdateMap.containsKey(key)) {
                entitiesUpdateMap.get(key).put(subDocPath, jsonDocument);
            } else {
                Map<String, Document> subDocument = new HashMap<>();
                subDocument.put(subDocPath, jsonDocument);
                entitiesUpdateMap.put(key, subDocument);
            }
        }
        try {
            entitiesCollection.bulkUpdateSubDocs(entitiesUpdateMap);
        } catch (Exception e) {
            LOG.error("Failed to update entities {}, subDocPath {}, with new doc {}.", entitiesUpdateMap, subDocPath, jsonDocument, e);
            throw e;
        }
    }
}

Example 10 with JSONDocument

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

the class EntityQueryServiceImpl method convertToJsonDocument.

@SneakyThrows
private JSONDocument convertToJsonDocument(LiteralConstant literalConstant) {
    // Convert setAttribute LiteralConstant to AttributeValue. Need to be able to store an array
    // literal constant as an array
    AttributeValue attributeValue = EntityQueryConverter.convertToAttributeValue(literalConstant).build();
    String jsonValue = PRINTER.print(attributeValue);
    return new JSONDocument(jsonValue);
}