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

Example 6 with Collection

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

the class PostgresDocStoreTest method testUpsertAndReturn.

@Test
public void testUpsertAndReturn() throws IOException {
    Collection collection = datastore.getCollection(COLLECTION_NAME);
    Document document = Utils.createDocument("foo1", "bar1");
    Document resultDocument = collection.upsertAndReturn(new SingleValueKey("default", "testKey"), document);
    Assertions.assertEquals(document.toJson(), resultDocument.toJson());
}

Example 7 with Collection

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

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

the class EntityQueryServiceImplTest method testUpdate_success.

@Test
public void testUpdate_success() throws Exception {
    Collection mockEntitiesCollection = mockEntitiesCollection();
    Builder newStatus = LiteralConstant.newBuilder().setValue(Value.newBuilder().setValueType(ValueType.STRING).setString("NEW_STATUS"));
    EntityUpdateRequest updateRequest = EntityUpdateRequest.newBuilder().setEntityType(TEST_ENTITY_TYPE).addEntityIds("entity-id-1").setOperation(UpdateOperation.newBuilder().setSetAttribute(SetAttribute.newBuilder().setAttribute(ColumnIdentifier.newBuilder().setColumnName(ATTRIBUTE_ID2)).setValue(newStatus))).addSelection(Expression.newBuilder().setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(ATTRIBUTE_ID1))).addSelection(Expression.newBuilder().setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(ATTRIBUTE_ID2))).build();
    StreamObserver<ResultSetChunk> mockResponseObserver = mock(StreamObserver.class);
    Context.current().withValue(RequestContext.CURRENT, mockRequestContextWithTenantId()).call(() -> {
        EntityQueryServiceImpl eqs = new EntityQueryServiceImpl(mockEntitiesCollection, mockMappingForAttributes1And2(), 1, false);
        eqs.update(updateRequest, mockResponseObserver);
        return null;
    });
    verify(mockEntitiesCollection, times(1)).bulkUpdateSubDocs(eq(Map.of(new SingleValueKey("tenant1", "entity-id-1"), Map.of("attributes.status", new JSONDocument(DocStoreJsonFormat.printer().print(newStatus))))));
}

Example 9 with Collection

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

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