Search in sources :

Example 1 with BulkArrayValueUpdateRequest

use of org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest in project document-store by hypertrace.

the class MongoDocStoreTest method test_bulkOperationOnArrayValue_addOperation.

@Test
public void test_bulkOperationOnArrayValue_addOperation() throws Exception {
    datastore.createCollection(COLLECTION_NAME, null);
    Collection collection = datastore.getCollection(COLLECTION_NAME);
    Key key1 = new SingleValueKey("default", "testKey1");
    Document key1InsertedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey1"), ImmutablePair.of("attributes", Map.of("name", "testKey1", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label1"))))))));
    Document key1ExpectedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey1"), ImmutablePair.of("attributes", Map.of("name", "testKey1", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label1")), ImmutablePair.of("value", Map.of("string", "Label2")), ImmutablePair.of("value", Map.of("string", "Label3"))))))));
    collection.upsert(key1, key1InsertedDocument);
    Key key2 = new SingleValueKey("default", "testKey2");
    Document key2InsertedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey2"), ImmutablePair.of("attributes", Map.of("name", "testKey2", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label2"))))))));
    Document key2ExpectedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey2"), ImmutablePair.of("attributes", Map.of("name", "testKey2", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label2")), ImmutablePair.of("value", Map.of("string", "Label3"))))))));
    collection.upsert(key2, key2InsertedDocument);
    Key key3 = new SingleValueKey("default", "testKey3");
    Document key3InsertedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey3"), ImmutablePair.of("attributes", Map.of("name", "testKey3")));
    Document key3ExpectedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey3"), ImmutablePair.of("attributes", Map.of("name", "testKey3", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label2")), ImmutablePair.of("value", Map.of("string", "Label3"))))))));
    collection.upsert(key3, key3InsertedDocument);
    Key key4 = new SingleValueKey("default", "testKey4");
    Document key4InsertedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey4"), ImmutablePair.of("attributes", Map.of("name", "testKey4", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label1")), ImmutablePair.of("value", Map.of("string", "Label2")), ImmutablePair.of("value", Map.of("string", "Label3"))))))));
    Document key4ExpectedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey4"), ImmutablePair.of("attributes", Map.of("name", "testKey4", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label1")), ImmutablePair.of("value", Map.of("string", "Label2")), ImmutablePair.of("value", Map.of("string", "Label3"))))))));
    collection.upsert(key4, key4InsertedDocument);
    Document label2Document = Utils.createDocument(ImmutablePair.of("value", Map.of("string", "Label2")));
    Document label3Document = Utils.createDocument(ImmutablePair.of("value", Map.of("string", "Label3")));
    List<Document> subDocuments = List.of(label2Document, label3Document);
    BulkArrayValueUpdateRequest bulkArrayValueUpdateRequest = new BulkArrayValueUpdateRequest(Set.of(key1, key2, key3, key4), "attributes.labels.valueList.values", ADD, subDocuments);
    BulkUpdateResult bulkUpdateResult = collection.bulkOperationOnArrayValue(bulkArrayValueUpdateRequest);
    assertEquals(4, bulkUpdateResult.getUpdatedCount());
    // get all documents
    Query query = new Query();
    Iterator<Document> results = collection.search(query);
    List<Document> documents = new ArrayList<>();
    while (results.hasNext()) {
        documents.add(results.next());
    }
    assertEquals(4, documents.size());
    Map<String, JsonNode> actualDocs = convertToMap(documents, "id");
    Map<String, JsonNode> expectedDocs = convertToMap(List.of(key1ExpectedDocument, key2ExpectedDocument, key3ExpectedDocument, key4ExpectedDocument), "id");
    // Verify that the documents returned are as expected
    for (Map.Entry<String, JsonNode> entry : actualDocs.entrySet()) {
        String key = entry.getKey();
        JsonNode attributesJsonNode = entry.getValue().get("attributes");
        JsonNode expectedAttributesJsonNode = expectedDocs.get(key).get("attributes");
        assertEquals(expectedAttributesJsonNode, attributesJsonNode);
    }
}
Also used : Query(org.hypertrace.core.documentstore.Query) BulkUpdateResult(org.hypertrace.core.documentstore.BulkUpdateResult) ArrayList(java.util.ArrayList) JsonNode(com.fasterxml.jackson.databind.JsonNode) Document(org.hypertrace.core.documentstore.Document) JSONDocument(org.hypertrace.core.documentstore.JSONDocument) BulkArrayValueUpdateRequest(org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest) SingleValueKey(org.hypertrace.core.documentstore.SingleValueKey) MongoCollection(com.mongodb.client.MongoCollection) Collection(org.hypertrace.core.documentstore.Collection) Map(java.util.Map) HashMap(java.util.HashMap) SingleValueKey(org.hypertrace.core.documentstore.SingleValueKey) Key(org.hypertrace.core.documentstore.Key) Test(org.junit.jupiter.api.Test)

Example 2 with BulkArrayValueUpdateRequest

use of org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest in project document-store by hypertrace.

the class MongoDocStoreTest method test_bulkOperationOnArrayValue_removeOperation.

@Test
public void test_bulkOperationOnArrayValue_removeOperation() throws Exception {
    datastore.createCollection(COLLECTION_NAME, null);
    Collection collection = datastore.getCollection(COLLECTION_NAME);
    Key key1 = new SingleValueKey("default", "testKey1");
    Document key1InsertedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey1"), ImmutablePair.of("attributes", Map.of("name", "testKey1", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label1"))))))));
    Document key1ExpectedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey1"), ImmutablePair.of("attributes", Map.of("name", "testKey1", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label1"))))))));
    collection.upsert(key1, key1InsertedDocument);
    Key key2 = new SingleValueKey("default", "testKey2");
    Document key2InsertedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey2"), ImmutablePair.of("attributes", Map.of("name", "testKey2", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label1")), ImmutablePair.of("value", Map.of("string", "Label2"))))))));
    Document key2ExpectedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey2"), ImmutablePair.of("attributes", Map.of("name", "testKey2", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label1"))))))));
    collection.upsert(key2, key2InsertedDocument);
    Key key3 = new SingleValueKey("default", "testKey3");
    Document key3InsertedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey3"), ImmutablePair.of("attributes", Map.of("name", "testKey3")));
    Document key3ExpectedDocument = key3InsertedDocument;
    collection.upsert(key3, key3InsertedDocument);
    Key key4 = new SingleValueKey("default", "testKey4");
    Document key4InsertedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey4"), ImmutablePair.of("attributes", Map.of("name", "testKey4", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label1")), ImmutablePair.of("value", Map.of("string", "Label2")), ImmutablePair.of("value", Map.of("string", "Label3"))))))));
    Document key4ExpectedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey4"), ImmutablePair.of("attributes", Map.of("name", "testKey4", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label1"))))))));
    collection.upsert(key4, key4InsertedDocument);
    Document label2Document = Utils.createDocument(ImmutablePair.of("value", Map.of("string", "Label2")));
    Document label3Document = Utils.createDocument(ImmutablePair.of("value", Map.of("string", "Label3")));
    List<Document> subDocuments = List.of(label2Document, label3Document);
    BulkArrayValueUpdateRequest bulkArrayValueUpdateRequest = new BulkArrayValueUpdateRequest(Set.of(key1, key2, key3, key4), "attributes.labels.valueList.values", REMOVE, subDocuments);
    BulkUpdateResult bulkUpdateResult = collection.bulkOperationOnArrayValue(bulkArrayValueUpdateRequest);
    assertEquals(4, bulkUpdateResult.getUpdatedCount());
    // get all documents
    Query query = new Query();
    Iterator<Document> results = collection.search(query);
    List<Document> documents = new ArrayList<>();
    while (results.hasNext()) {
        documents.add(results.next());
    }
    assertEquals(4, documents.size());
    Map<String, JsonNode> actualDocs = convertToMap(documents, "id");
    Map<String, JsonNode> expectedDocs = convertToMap(List.of(key1ExpectedDocument, key2ExpectedDocument, key3ExpectedDocument, key4ExpectedDocument), "id");
    // Verify that the documents returned are as expected
    for (Map.Entry<String, JsonNode> entry : actualDocs.entrySet()) {
        String key = entry.getKey();
        JsonNode attributesJsonNode = entry.getValue().get("attributes");
        JsonNode expectedAttributesJsonNode = expectedDocs.get(key).get("attributes");
        assertEquals(expectedAttributesJsonNode, attributesJsonNode);
    }
}
Also used : Query(org.hypertrace.core.documentstore.Query) BulkUpdateResult(org.hypertrace.core.documentstore.BulkUpdateResult) ArrayList(java.util.ArrayList) JsonNode(com.fasterxml.jackson.databind.JsonNode) Document(org.hypertrace.core.documentstore.Document) JSONDocument(org.hypertrace.core.documentstore.JSONDocument) BulkArrayValueUpdateRequest(org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest) SingleValueKey(org.hypertrace.core.documentstore.SingleValueKey) MongoCollection(com.mongodb.client.MongoCollection) Collection(org.hypertrace.core.documentstore.Collection) Map(java.util.Map) HashMap(java.util.HashMap) SingleValueKey(org.hypertrace.core.documentstore.SingleValueKey) Key(org.hypertrace.core.documentstore.Key) Test(org.junit.jupiter.api.Test)

Example 3 with BulkArrayValueUpdateRequest

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

the class EntityQueryServiceImplTest method testBulkUpdateEntityArrayAttribute.

@Test
void testBulkUpdateEntityArrayAttribute() throws Exception {
    List<String> entityIds = IntStream.rangeClosed(1, 5).mapToObj(i -> UUID.randomUUID().toString()).collect(Collectors.toList());
    BulkEntityArrayAttributeUpdateRequest request = BulkEntityArrayAttributeUpdateRequest.newBuilder().addAllEntityIds(entityIds).setAttribute(ColumnIdentifier.newBuilder().setColumnName(ATTRIBUTE_ID3).build()).setEntityType(TEST_ENTITY_TYPE).setOperation(BulkEntityArrayAttributeUpdateRequest.Operation.OPERATION_ADD).addAllValues(List.of(LiteralConstant.newBuilder().setValue(Value.newBuilder().setString("Label1")).build(), LiteralConstant.newBuilder().setValue(Value.newBuilder().setString("Label2")).build())).build();
    when(mockAttributeMapping.getDocStorePathByAttributeId(requestContext, ATTRIBUTE_ID3)).thenReturn(Optional.of(EDS_COLUMN_NAME3));
    StreamObserver<BulkEntityArrayAttributeUpdateResponse> mockResponseObserver = mock(StreamObserver.class);
    Context.current().withValue(RequestContext.CURRENT, mockRequestContextWithTenantId()).call(() -> {
        EntityQueryServiceImpl eqs = new EntityQueryServiceImpl(entitiesCollection, mockAttributeMapping, 1, false);
        eqs.bulkUpdateEntityArrayAttribute(request, mockResponseObserver);
        return null;
    });
    ArgumentCaptor<BulkArrayValueUpdateRequest> argumentCaptor = ArgumentCaptor.forClass(BulkArrayValueUpdateRequest.class);
    verify(entitiesCollection, times(1)).bulkOperationOnArrayValue(argumentCaptor.capture());
    BulkArrayValueUpdateRequest bulkArrayValueUpdateRequest = argumentCaptor.getValue();
    assertEquals(entityIds.stream().map(entityId -> new SingleValueKey("tenant1", entityId)).collect(Collectors.toCollection(LinkedHashSet::new)), bulkArrayValueUpdateRequest.getKeys());
    assertEquals(BulkArrayValueUpdateRequest.Operation.ADD, bulkArrayValueUpdateRequest.getOperation());
    assertEquals(EDS_COLUMN_NAME3 + ".valueList.values", bulkArrayValueUpdateRequest.getSubDocPath());
    List<Document> subDocuments = bulkArrayValueUpdateRequest.getSubDocuments();
    assertEquals(2, subDocuments.size());
    assertEquals("{\"value\":{\"string\":\"Label1\"}}", subDocuments.get(0).toString());
    assertEquals("{\"value\":{\"string\":\"Label2\"}}", subDocuments.get(1).toString());
}
Also used : Context(io.grpc.Context) ArgumentMatchers.argThat(org.mockito.ArgumentMatchers.argThat) Builder(org.hypertrace.entity.query.service.v1.LiteralConstant.Builder) ResultSetChunk(org.hypertrace.entity.query.service.v1.ResultSetChunk) ArgumentMatchers.eq(org.mockito.ArgumentMatchers.eq) UpdateOperation(org.hypertrace.entity.query.service.v1.UpdateOperation) ResultSetMetadata(org.hypertrace.entity.query.service.v1.ResultSetMetadata) Query(org.hypertrace.core.documentstore.Query) Disabled(org.junit.jupiter.api.Disabled) ArgumentMatcher(org.mockito.ArgumentMatcher) StreamObserver(io.grpc.stub.StreamObserver) ExtendWith(org.junit.jupiter.api.extension.ExtendWith) Map(java.util.Map) BulkEntityArrayAttributeUpdateResponse(org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateResponse) BulkEntityArrayAttributeUpdateRequest(org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateRequest) SingleValueKey(org.hypertrace.core.documentstore.SingleValueKey) Value(org.hypertrace.entity.query.service.v1.Value) MockitoExtension(org.mockito.junit.jupiter.MockitoExtension) BulkEntityUpdateRequest(org.hypertrace.entity.query.service.v1.BulkEntityUpdateRequest) UUID(java.util.UUID) Collectors(java.util.stream.Collectors) ColumnMetadata(org.hypertrace.entity.query.service.v1.ColumnMetadata) Test(org.junit.jupiter.api.Test) List(java.util.List) JsonFormat(com.google.protobuf.util.JsonFormat) TotalEntitiesRequest(org.hypertrace.entity.query.service.v1.TotalEntitiesRequest) Optional(java.util.Optional) EntityQueryRequest(org.hypertrace.entity.query.service.v1.EntityQueryRequest) LiteralConstant(org.hypertrace.entity.query.service.v1.LiteralConstant) Mockito.mock(org.mockito.Mockito.mock) SetAttribute(org.hypertrace.entity.query.service.v1.SetAttribute) ArgumentMatchers.any(org.mockito.ArgumentMatchers.any) IntStream(java.util.stream.IntStream) Row(org.hypertrace.entity.query.service.v1.Row) Document(org.hypertrace.core.documentstore.Document) RequestContext(org.hypertrace.core.grpcutils.context.RequestContext) DocStoreJsonFormat(org.hypertrace.entity.service.util.DocStoreJsonFormat) MockitoSettings(org.mockito.junit.jupiter.MockitoSettings) Filter(org.hypertrace.core.documentstore.Filter) Mock(org.mockito.Mock) EntityUpdateRequest(org.hypertrace.entity.query.service.v1.EntityUpdateRequest) LENIENT(org.mockito.quality.Strictness.LENIENT) Nested(org.junit.jupiter.api.Nested) Collection(org.hypertrace.core.documentstore.Collection) ArgumentCaptor(org.mockito.ArgumentCaptor) EntityUpdateInfo(org.hypertrace.entity.query.service.v1.BulkEntityUpdateRequest.EntityUpdateInfo) AttributeValue(org.hypertrace.entity.data.service.v1.AttributeValue) Entity(org.hypertrace.entity.data.service.v1.Entity) Assertions.assertEquals(org.junit.jupiter.api.Assertions.assertEquals) LinkedHashSet(java.util.LinkedHashSet) TotalEntitiesResponse(org.hypertrace.entity.query.service.v1.TotalEntitiesResponse) BulkArrayValueUpdateRequest(org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest) ColumnIdentifier(org.hypertrace.entity.query.service.v1.ColumnIdentifier) BulkUpdateResult(org.hypertrace.core.documentstore.BulkUpdateResult) JSONDocument(org.hypertrace.core.documentstore.JSONDocument) Mockito.times(org.mockito.Mockito.times) Mockito.when(org.mockito.Mockito.when) Mockito.verify(org.mockito.Mockito.verify) DisplayName(org.junit.jupiter.api.DisplayName) Mockito(org.mockito.Mockito) TestUtils.convertToCloseableIterator(org.hypertrace.entity.TestUtils.convertToCloseableIterator) Expression(org.hypertrace.entity.query.service.v1.Expression) ValueType(org.hypertrace.entity.query.service.v1.ValueType) OrderByExpression(org.hypertrace.entity.query.service.v1.OrderByExpression) LinkedHashSet(java.util.LinkedHashSet) BulkEntityArrayAttributeUpdateRequest(org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateRequest) BulkEntityArrayAttributeUpdateResponse(org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateResponse) Document(org.hypertrace.core.documentstore.Document) JSONDocument(org.hypertrace.core.documentstore.JSONDocument) BulkArrayValueUpdateRequest(org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest) SingleValueKey(org.hypertrace.core.documentstore.SingleValueKey) Test(org.junit.jupiter.api.Test)

Example 4 with BulkArrayValueUpdateRequest

use of org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest in project document-store by hypertrace.

the class MongoDocStoreTest method test_bulkOperationOnArrayValue_setOperation.

@Test
public void test_bulkOperationOnArrayValue_setOperation() throws Exception {
    datastore.createCollection(COLLECTION_NAME, null);
    Collection collection = datastore.getCollection(COLLECTION_NAME);
    Key key1 = new SingleValueKey("default", "testKey1");
    Document key1InsertedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey1"), ImmutablePair.of("attributes", Map.of("name", "testKey1", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label1"))))))));
    Document key1ExpectedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey1"), ImmutablePair.of("attributes", Map.of("name", "testKey1", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label2")), ImmutablePair.of("value", Map.of("string", "Label3"))))))));
    collection.upsert(key1, key1InsertedDocument);
    Key key2 = new SingleValueKey("default", "testKey2");
    Document key2InsertedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey2"), ImmutablePair.of("attributes", Map.of("name", "testKey2", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label2"))))))));
    Document key2ExpectedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey2"), ImmutablePair.of("attributes", Map.of("name", "testKey2", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label2")), ImmutablePair.of("value", Map.of("string", "Label3"))))))));
    collection.upsert(key2, key2InsertedDocument);
    Key key3 = new SingleValueKey("default", "testKey3");
    Document key3InsertedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey3"), ImmutablePair.of("attributes", Map.of("name", "testKey3")));
    Document key3ExpectedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey3"), ImmutablePair.of("attributes", Map.of("name", "testKey3", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label2")), ImmutablePair.of("value", Map.of("string", "Label3"))))))));
    collection.upsert(key3, key3InsertedDocument);
    Key key4 = new SingleValueKey("default", "testKey4");
    Document key4InsertedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey4"), ImmutablePair.of("attributes", Map.of("name", "testKey4", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label1")), ImmutablePair.of("value", Map.of("string", "Label2")), ImmutablePair.of("value", Map.of("string", "Label3"))))))));
    Document key4ExpectedDocument = Utils.createDocument(ImmutablePair.of("id", "testKey4"), ImmutablePair.of("attributes", Map.of("name", "testKey4", "labels", ImmutablePair.of("valueList", ImmutablePair.of("values", List.of(ImmutablePair.of("value", Map.of("string", "Label2")), ImmutablePair.of("value", Map.of("string", "Label3"))))))));
    collection.upsert(key4, key4InsertedDocument);
    Document label2Document = Utils.createDocument(ImmutablePair.of("value", Map.of("string", "Label2")));
    Document label3Document = Utils.createDocument(ImmutablePair.of("value", Map.of("string", "Label3")));
    List<Document> subDocuments = List.of(label2Document, label3Document);
    BulkArrayValueUpdateRequest bulkArrayValueUpdateRequest = new BulkArrayValueUpdateRequest(Set.of(key1, key2, key3, key4), "attributes.labels.valueList.values", SET, subDocuments);
    BulkUpdateResult bulkUpdateResult = collection.bulkOperationOnArrayValue(bulkArrayValueUpdateRequest);
    assertEquals(4, bulkUpdateResult.getUpdatedCount());
    // get all documents
    Query query = new Query();
    Iterator<Document> results = collection.search(query);
    List<Document> documents = new ArrayList<>();
    while (results.hasNext()) {
        documents.add(results.next());
    }
    assertEquals(4, documents.size());
    Map<String, JsonNode> actualDocs = convertToMap(documents, "id");
    Map<String, JsonNode> expectedDocs = convertToMap(List.of(key1ExpectedDocument, key2ExpectedDocument, key3ExpectedDocument, key4ExpectedDocument), "id");
    // Verify that the documents returned are as expected
    for (Map.Entry<String, JsonNode> entry : actualDocs.entrySet()) {
        String key = entry.getKey();
        JsonNode attributesJsonNode = entry.getValue().get("attributes");
        JsonNode expectedAttributesJsonNode = expectedDocs.get(key).get("attributes");
        assertEquals(expectedAttributesJsonNode, attributesJsonNode);
    }
}
Also used : Query(org.hypertrace.core.documentstore.Query) BulkUpdateResult(org.hypertrace.core.documentstore.BulkUpdateResult) ArrayList(java.util.ArrayList) JsonNode(com.fasterxml.jackson.databind.JsonNode) Document(org.hypertrace.core.documentstore.Document) JSONDocument(org.hypertrace.core.documentstore.JSONDocument) BulkArrayValueUpdateRequest(org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest) SingleValueKey(org.hypertrace.core.documentstore.SingleValueKey) MongoCollection(com.mongodb.client.MongoCollection) Collection(org.hypertrace.core.documentstore.Collection) Map(java.util.Map) HashMap(java.util.HashMap) SingleValueKey(org.hypertrace.core.documentstore.SingleValueKey) Key(org.hypertrace.core.documentstore.Key) Test(org.junit.jupiter.api.Test)

Example 5 with BulkArrayValueUpdateRequest

use of org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest 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);
    }
}
Also used : DocumentParser(org.hypertrace.entity.data.service.DocumentParser) Function(org.hypertrace.entity.query.service.v1.Function) EntityQueryServiceImplBase(org.hypertrace.entity.query.service.v1.EntityQueryServiceGrpc.EntityQueryServiceImplBase) SneakyThrows(lombok.SneakyThrows) ResultSetChunk(org.hypertrace.entity.query.service.v1.ResultSetChunk) UpdateOperation(org.hypertrace.entity.query.service.v1.UpdateOperation) ServiceException(com.google.protobuf.ServiceException) ResultSetMetadata(org.hypertrace.entity.query.service.v1.ResultSetMetadata) LoggerFactory(org.slf4j.LoggerFactory) ConversionException(org.hypertrace.entity.query.service.converter.ConversionException) StreamObserver(io.grpc.stub.StreamObserver) Map(java.util.Map) RelationalExpression(org.hypertrace.core.documentstore.expression.impl.RelationalExpression) BulkEntityArrayAttributeUpdateResponse(org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateResponse) BulkEntityArrayAttributeUpdateRequest(org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateRequest) Objects.isNull(java.util.Objects.isNull) SingleValueKey(org.hypertrace.core.documentstore.SingleValueKey) DocStoreConverter(org.hypertrace.entity.service.util.DocStoreConverter) Filter(org.hypertrace.core.documentstore.query.Filter) Value(org.hypertrace.entity.query.service.v1.Value) Collections.emptyList(java.util.Collections.emptyList) BulkEntityUpdateRequest(org.hypertrace.entity.query.service.v1.BulkEntityUpdateRequest) Set(java.util.Set) StringUtils(org.hypertrace.entity.service.util.StringUtils) Collectors(java.util.stream.Collectors) Collectors.joining(java.util.stream.Collectors.joining) ColumnMetadata(org.hypertrace.entity.query.service.v1.ColumnMetadata) List(java.util.List) Lists.newArrayList(com.google.common.collect.Lists.newArrayList) Stream(java.util.stream.Stream) GrpcChannelRegistry(org.hypertrace.core.grpcutils.client.GrpcChannelRegistry) IN(org.hypertrace.core.documentstore.expression.operators.RelationalOperator.IN) TotalEntitiesRequest(org.hypertrace.entity.query.service.v1.TotalEntitiesRequest) Optional(java.util.Optional) EntityQueryRequest(org.hypertrace.entity.query.service.v1.EntityQueryRequest) TypeLiteral(com.google.inject.TypeLiteral) LiteralConstant(org.hypertrace.entity.query.service.v1.LiteralConstant) SetAttribute(org.hypertrace.entity.query.service.v1.SetAttribute) Row(org.hypertrace.entity.query.service.v1.Row) Document(org.hypertrace.core.documentstore.Document) RequestContext(org.hypertrace.core.grpcutils.context.RequestContext) DocStoreJsonFormat(org.hypertrace.entity.service.util.DocStoreJsonFormat) ENTITY_ATTRIBUTE_DOC_PREFIX(org.hypertrace.entity.query.service.EntityAttributeMapping.ENTITY_ATTRIBUTE_DOC_PREFIX) AliasProvider(org.hypertrace.entity.query.service.converter.AliasProvider) ConstantExpression(org.hypertrace.core.documentstore.expression.impl.ConstantExpression) EntityUpdateRequest(org.hypertrace.entity.query.service.v1.EntityUpdateRequest) HashMap(java.util.HashMap) Collection(org.hypertrace.core.documentstore.Collection) Collectors.toUnmodifiableList(java.util.stream.Collectors.toUnmodifiableList) ArrayList(java.util.ArrayList) Datastore(org.hypertrace.core.documentstore.Datastore) Converter(org.hypertrace.entity.query.service.converter.Converter) Key(org.hypertrace.core.documentstore.Key) ConverterModule(org.hypertrace.entity.query.service.converter.ConverterModule) EntityUpdateInfo(org.hypertrace.entity.query.service.v1.BulkEntityUpdateRequest.EntityUpdateInfo) AttributeValue(org.hypertrace.entity.data.service.v1.AttributeValue) EntityServiceConstants(org.hypertrace.entity.service.constants.EntityServiceConstants) VALUES_FIELD_NUMBER(org.hypertrace.entity.data.service.v1.AttributeValueList.VALUES_FIELD_NUMBER) Entity(org.hypertrace.entity.data.service.v1.Entity) DocumentConverter(org.hypertrace.entity.query.service.converter.response.DocumentConverter) LinkedHashSet(java.util.LinkedHashSet) TotalEntitiesResponse(org.hypertrace.entity.query.service.v1.TotalEntitiesResponse) BulkArrayValueUpdateRequest(org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest) IdentifierExpression(org.hypertrace.core.documentstore.expression.impl.IdentifierExpression) ColumnIdentifier(org.hypertrace.entity.query.service.v1.ColumnIdentifier) Logger(org.slf4j.Logger) Iterator(java.util.Iterator) Config(com.typesafe.config.Config) JSONDocument(org.hypertrace.core.documentstore.JSONDocument) Printer(org.hypertrace.entity.service.util.DocStoreJsonFormat.Printer) VALUE_LIST_FIELD_NUMBER(org.hypertrace.entity.data.service.v1.AttributeValue.VALUE_LIST_FIELD_NUMBER) AttributeValueList(org.hypertrace.entity.data.service.v1.AttributeValueList) Injector(com.google.inject.Injector) Selection(org.hypertrace.core.documentstore.query.Selection) Query(org.hypertrace.entity.data.service.v1.Query) Expression(org.hypertrace.entity.query.service.v1.Expression) RAW_ENTITIES_COLLECTION(org.hypertrace.entity.service.constants.EntityCollectionConstants.RAW_ENTITIES_COLLECTION) ValueType(org.hypertrace.entity.query.service.v1.ValueType) Guice(com.google.inject.Guice) VisibleForTesting(com.google.common.annotations.VisibleForTesting) ValueCase(org.hypertrace.entity.query.service.v1.Expression.ValueCase) Collections(java.util.Collections) SingleValueKey(org.hypertrace.core.documentstore.SingleValueKey) ServiceException(com.google.protobuf.ServiceException) RequestContext(org.hypertrace.core.grpcutils.context.RequestContext) Document(org.hypertrace.core.documentstore.Document) JSONDocument(org.hypertrace.core.documentstore.JSONDocument) BulkArrayValueUpdateRequest(org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest) SingleValueKey(org.hypertrace.core.documentstore.SingleValueKey) Key(org.hypertrace.core.documentstore.Key) ServiceException(com.google.protobuf.ServiceException) ConversionException(org.hypertrace.entity.query.service.converter.ConversionException)

Aggregations

Map (java.util.Map)5 BulkArrayValueUpdateRequest (org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest)5 Collection (org.hypertrace.core.documentstore.Collection)5 Document (org.hypertrace.core.documentstore.Document)5 JSONDocument (org.hypertrace.core.documentstore.JSONDocument)5 SingleValueKey (org.hypertrace.core.documentstore.SingleValueKey)5 ArrayList (java.util.ArrayList)4 HashMap (java.util.HashMap)4 BulkUpdateResult (org.hypertrace.core.documentstore.BulkUpdateResult)4 Key (org.hypertrace.core.documentstore.Key)4 Query (org.hypertrace.core.documentstore.Query)4 JsonNode (com.fasterxml.jackson.databind.JsonNode)3 MongoCollection (com.mongodb.client.MongoCollection)3 Test (org.junit.jupiter.api.Test)3 StreamObserver (io.grpc.stub.StreamObserver)2 LinkedHashSet (java.util.LinkedHashSet)2 List (java.util.List)2 Optional (java.util.Optional)2 Collectors (java.util.stream.Collectors)2 VisibleForTesting (com.google.common.annotations.VisibleForTesting)1