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

Example 16 with Document

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

the class CreateUpdateTestThread method updateRun.

private void updateRun() throws IOException {
    Document document = Utils.createDocument(ImmutablePair.of("id", documentKey.getValue()), ImmutablePair.of("name", String.format("thread-%s", this.testValue)), ImmutablePair.of("size", this.testValue));
    // do conditional update
    Filter condition = new Filter(Op.EQ, "size", this.testValue);
    updateTestResult = collection.update(documentKey, document, condition);
}

Example 17 with Document

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

the class PostgresCollection method bulkUpsertImpl.

private int[] bulkUpsertImpl(Map<Key, Document> documents) throws SQLException, IOException {
    try (PreparedStatement preparedStatement = client.prepareStatement(getUpsertSQL(), Statement.RETURN_GENERATED_KEYS)) {
        for (Map.Entry<Key, Document> entry : documents.entrySet()) {
            Key key = entry.getKey();
            String jsonString = prepareDocument(key, entry.getValue());
            preparedStatement.setString(1, key.toString());
            preparedStatement.setString(2, jsonString);
            preparedStatement.setString(3, jsonString);
            preparedStatement.addBatch();
        }
        return preparedStatement.executeBatch();
    }
}

Example 18 with Document

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

the class PostgresCollection method bulkUpdateSubDocs.

@Override
public BulkUpdateResult bulkUpdateSubDocs(Map<Key, Map<String, Document>> documents) throws Exception {
    String updateSubDocSQL = String.format("UPDATE %s SET %s=jsonb_set(%s, ?::text[], ?::jsonb) WHERE %s = ?", collectionName, DOCUMENT, DOCUMENT, ID);
    try {
        PreparedStatement preparedStatement = client.prepareStatement(updateSubDocSQL);
        for (Key key : documents.keySet()) {
            Map<String, Document> subDocuments = documents.get(key);
            for (String subDocPath : subDocuments.keySet()) {
                Document subDocument = subDocuments.get(subDocPath);
                String jsonSubDocPath = getJsonSubDocPath(subDocPath);
                String jsonString = subDocument.toJson();
                preparedStatement.setString(1, jsonSubDocPath);
                preparedStatement.setString(2, jsonString);
                preparedStatement.setString(3, key.toString());
                preparedStatement.addBatch();
            }
        }
        int[] updateCounts = preparedStatement.executeBatch();
        if (LOGGER.isDebugEnabled()) {
            LOGGER.debug("Write result: {}", updateCounts);
        }
        int totalUpdateCount = 0;
        for (int update : updateCounts) {
            totalUpdateCount += update;
        }
        return new BulkUpdateResult(totalUpdateCount);
    } catch (SQLException e) {
        LOGGER.error("SQLException updating sub document.", e);
        throw e;
    }
}

Example 19 with Document

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

the class EntityQueryServiceImplTest method testExecute_success_chunksize_2.

@Test
public void testExecute_success_chunksize_2() 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();
    Entity entity3 = Entity.newBuilder().setTenantId("tenant-1").setEntityType(TEST_ENTITY_TYPE).setEntityId(UUID.randomUUID().toString()).setEntityName("Test entity 3").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)), new JSONDocument(JsonFormat.printer().print(entity3)), // 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(), 2, false);
        eqs.execute(request, mockResponseObserver);
        return null;
    });
    verify(mockEntitiesCollection, times(0)).aggregate(any());
    verify(mockEntitiesCollection, times(1)).search(any());
    verify(mockResponseObserver, times(2)).onNext(any());
    verify(mockResponseObserver, times(1)).onCompleted();
}

Example 20 with Document

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