Examples with Entities org.hypertrace.entity.data.service.v1.Entities used on opensource projects

Example 11 with Entities

use of org.hypertrace.entity.data.service.v1.Entities in project entity-service by hypertrace.

the class EntityDataServiceImpl method searchByIdAndStreamSingleResponse.

private <T extends Message> void searchByIdAndStreamSingleResponse(String tenantId, String entityId, String entityType, Collection collection, Message.Builder builder, StreamObserver<T> responseObserver) {
    org.hypertrace.core.documentstore.Query query = new org.hypertrace.core.documentstore.Query();
    String docId = this.entityNormalizer.getEntityDocKey(tenantId, entityType, entityId).toString();
    query.setFilter(new Filter(Filter.Op.EQ, EntityServiceConstants.ID, docId));
    Iterator<Document> result = collection.search(query);
    List<T> entities = new ArrayList<>();
    while (result.hasNext()) {
        PARSER.<T>parseOrLog(result.next(), builder.clone()).map(entity -> {
            // Populate the tenant id field with the tenant id that's received for backward
            // compatibility.
            Descriptors.FieldDescriptor fieldDescriptor = entity.getDescriptorForType().findFieldByName("tenant_id");
            if (fieldDescriptor != null) {
                return (T) entity.toBuilder().setField(fieldDescriptor, tenantId).build();
            }
            return entity;
        }).ifPresent(entities::add);
    }
    if (LOG.isDebugEnabled()) {
        LOG.debug("Docstore query has returned the result: {}", entities);
    }
    if (entities.size() == 1) {
        responseObserver.onNext(entities.get(0));
        responseObserver.onCompleted();
    } else if (entities.size() > 1) {
        responseObserver.onError(new IllegalStateException("Multiple entities with same id are found."));
    } else {
        // When there is no result, we should return the default instance, which is a way
        // of saying it's null.
        // TODO : Not convinced with the default instance
        responseObserver.onNext((T) builder.build());
        responseObserver.onCompleted();
    }
}

Example 12 with Entities

use of org.hypertrace.entity.data.service.v1.Entities in project entity-service by hypertrace.

the class EntityDataServiceTest method testEntityQuery.

@Test
public void testEntityQuery() {
    Entity entity1 = Entity.newBuilder().setTenantId(TENANT_ID).setEntityType(EntityType.K8S_POD.name()).setEntityName("Some Service 1").putIdentifyingAttributes(EntityConstants.getValue(CommonAttribute.COMMON_ATTRIBUTE_EXTERNAL_ID), generateRandomUUIDAttrValue()).build();
    Entity createdEntity1 = entityDataServiceClient.upsert(entity1);
    assertNotNull(createdEntity1);
    assertNotNull(createdEntity1.getEntityId().trim());
    Entity entity2 = Entity.newBuilder().setTenantId(TENANT_ID).setEntityType(EntityType.K8S_POD.name()).setEntityName("Some Service 2").putIdentifyingAttributes(EntityConstants.getValue(CommonAttribute.COMMON_ATTRIBUTE_EXTERNAL_ID), generateRandomUUIDAttrValue()).build();
    Entity createdEntity2 = entityDataServiceClient.upsert(entity2);
    assertNotNull(createdEntity2);
    assertNotNull(createdEntity2.getEntityId().trim());
    Query entityTypeQuery = Query.newBuilder().setEntityType(EntityType.K8S_POD.name()).build();
    // Query all entities of type SERVICE
    List<Entity> entitiesList = entityDataServiceClient.query(TENANT_ID, entityTypeQuery);
    assertTrue(entitiesList.size() > 1);
    Query entityLimitQuery = Query.newBuilder().setEntityType(EntityType.K8S_POD.name()).setLimit(1).build();
    entitiesList = entityDataServiceClient.query(TENANT_ID, entityLimitQuery);
    assertEquals(1, entitiesList.size());
    Query entityOffsetQuery = Query.newBuilder().setEntityType(EntityType.K8S_POD.name()).setLimit(1).setOffset(1).build();
    List<Entity> entityWithOffset = entityDataServiceClient.query(TENANT_ID, entityOffsetQuery);
    assertEquals(1, entityWithOffset.size());
    assertNotEquals(entitiesList.get(0).getEntityId(), entityWithOffset.get(0).getEntityId());
    // Query specific entity
    Query entityTypeAndIdQuery = Query.newBuilder().addEntityId(createdEntity1.getEntityId()).setEntityType(EntityType.K8S_POD.name()).build();
    entitiesList = entityDataServiceClient.query(TENANT_ID, entityTypeAndIdQuery);
    Entity foundEntity = entitiesList.get(0);
    assertEquals(createdEntity1, foundEntity);
    // Query for entity that doesn't exist
    Query randomEntityIdQuery = Query.newBuilder().addEntityId(UUID.randomUUID().toString()).setEntityType(EntityType.K8S_POD.name()).build();
    entitiesList = entityDataServiceClient.query(TENANT_ID, randomEntityIdQuery);
    assertTrue(entitiesList.isEmpty());
}

Example 13 with Entities

use of org.hypertrace.entity.data.service.v1.Entities in project entity-service by hypertrace.

the class EntityDataServiceTest method whenNEntitiesAreUpdated_thenExpectThemToBeUpdated.

@Test
public void whenNEntitiesAreUpdated_thenExpectThemToBeUpdated() {
    int N = 5;
    Map<String, Entity> externalIdToEntity = new HashMap<>();
    Map<String, Entity> externalIdToNewEntity = new HashMap<>();
    Map<String, AttributeValue> externalIdToAV = new HashMap<>();
    // Create N entities first with an attribute
    for (int i = 0; i < N; i++) {
        AttributeValue entityId = generateRandomUUIDAttrValue();
        Entity entity = Entity.newBuilder().setTenantId(TENANT_ID).setEntityType(EntityType.K8S_POD.name()).setEntityName("Some Service").putIdentifyingAttributes(EntityConstants.getValue(CommonAttribute.COMMON_ATTRIBUTE_EXTERNAL_ID), entityId).putAttributes("foo.key", AttributeValue.newBuilder().setValue(Value.newBuilder().setString("foo.value.old").build()).build()).build();
        externalIdToEntity.put(entityId.getValue().getString(), entity);
        externalIdToAV.put(entityId.getValue().getString(), entityId);
    }
    entityDataServiceClient.bulkUpsert(TENANT_ID, externalIdToEntity.values());
    for (String id : externalIdToEntity.keySet()) {
        List<Entity> readEntity = entityDataServiceClient.getEntitiesWithGivenAttribute(TENANT_ID, EntityType.K8S_POD.name(), EntityConstants.getValue(CommonAttribute.COMMON_ATTRIBUTE_EXTERNAL_ID), externalIdToAV.get(id));
        // exactly one entity exists
        assertEquals(1, readEntity.size());
        System.out.println("Id = " + id + ", old attribute_val = " + readEntity.get(0).getAttributesMap().get("foo.key").getValue().getString());
        // check that they contain old value
        Assertions.assertEquals("foo.value.old", readEntity.get(0).getAttributesMap().get("foo.key").getValue().getString());
    }
    // Now update the attribute
    for (Map.Entry<String, Entity> entry : externalIdToEntity.entrySet()) {
        Entity entity = entry.getValue();
        String entityId = entry.getKey();
        Entity newEntity = Entity.newBuilder().setEntityType(entity.getEntityType()).setEntityName(entity.getEntityName()).putIdentifyingAttributes(EntityConstants.getValue(CommonAttribute.COMMON_ATTRIBUTE_EXTERNAL_ID), entity.getIdentifyingAttributesMap().get(EntityConstants.getValue(CommonAttribute.COMMON_ATTRIBUTE_EXTERNAL_ID))).putAttributes("foo.key", AttributeValue.newBuilder().setValue(Value.newBuilder().setString("foo.value.new").build()).build()).build();
        externalIdToNewEntity.put(entityId, newEntity);
    }
    // upsert
    entityDataServiceClient.bulkUpsert(TENANT_ID, externalIdToNewEntity.values());
    // verify if the entities have been updated
    for (String id : externalIdToEntity.keySet()) {
        List<Entity> readEntity = entityDataServiceClient.getEntitiesWithGivenAttribute(TENANT_ID, EntityType.K8S_POD.name(), EntityConstants.getValue(CommonAttribute.COMMON_ATTRIBUTE_EXTERNAL_ID), externalIdToAV.get(id));
        // exactly one entity exists
        assertEquals(1, readEntity.size());
        System.out.println("Id = " + id + ", new attribute_val = " + readEntity.get(0).getAttributesMap().get("foo.key").getValue().getString());
        // check that they contain old value
        Assertions.assertEquals("foo.value.new", readEntity.get(0).getAttributesMap().get("foo.key").getValue().getString());
    }
}

Example 14 with Entities

use of org.hypertrace.entity.data.service.v1.Entities in project entity-service by hypertrace.

the class EntityDataServiceTest method whenNNewEntitiesAreUpserted_thenExpectNNewEntities.

@Test
public void whenNNewEntitiesAreUpserted_thenExpectNNewEntities() {
    int N = 5;
    Map<String, Entity> externalIdToEntity = new HashMap<>();
    Map<String, AttributeValue> externalIdToAV = new HashMap<>();
    for (int i = 0; i < N; i++) {
        AttributeValue entityId = generateRandomUUIDAttrValue();
        Entity entity = Entity.newBuilder().setTenantId(TENANT_ID).setEntityType(EntityType.K8S_POD.name()).setEntityName("Some Service").putIdentifyingAttributes(EntityConstants.getValue(CommonAttribute.COMMON_ATTRIBUTE_EXTERNAL_ID), entityId).build();
        externalIdToEntity.put(entityId.getValue().getString(), entity);
        externalIdToAV.put(entityId.getValue().getString(), entityId);
    }
    entityDataServiceClient.bulkUpsert(TENANT_ID, externalIdToEntity.values());
    // all N entities should have been created
    Map<String, Entity> entityMap = new HashMap<>();
    for (String id : externalIdToEntity.keySet()) {
        List<Entity> readEntity = entityDataServiceClient.getEntitiesWithGivenAttribute(TENANT_ID, EntityType.K8S_POD.name(), EntityConstants.getValue(CommonAttribute.COMMON_ATTRIBUTE_EXTERNAL_ID), externalIdToAV.get(id));
        // exactly one entity exists
        assertEquals(1, readEntity.size());
        assertNotNull(readEntity.get(0).getEntityId());
        entityMap.put(readEntity.get(0).getEntityId(), readEntity.get(0));
    }
    // Try getAndBulkUpsert, verify that the returned entities were in previous state.
    Iterator<Entity> iterator = entityDataServiceClient.getAndBulkUpsert(TENANT_ID, externalIdToEntity.values());
    while (iterator.hasNext()) {
        Entity entity = iterator.next();
        assertNotNull(entityMap.get(entity.getEntityId()));
        assertEquals(entityMap.get(entity.getEntityId()), entity);
    }
}

Example 15 with Entities

use of org.hypertrace.entity.data.service.v1.Entities in project entity-service by hypertrace.

the class EntityQueryServiceTest method testExecuteWithAggregations.

@Test
@Disabled("Disabled until we enable querying based on the new Query DTO")
void testExecuteWithAggregations() {
    // create and upsert some entities
    Entity entity1 = Entity.newBuilder().setTenantId(TENANT_ID).setEntityType(EntityType.SERVICE.name()).setEntityName("Some Service 1").putAttributes(EntityConstants.getValue(SERVICE_ATTRIBUTE_SERVICE_TYPE), generateAttrValue(JAEGER_SERVICE.name())).putIdentifyingAttributes(EntityConstants.getValue(CommonAttribute.COMMON_ATTRIBUTE_FQN), generateRandomUUIDAttrValue()).build();
    Entity createdEntity1 = entityDataServiceClient.upsert(entity1);
    assertNotNull(createdEntity1);
    assertFalse(createdEntity1.getEntityId().trim().isEmpty());
    Entity entity2 = Entity.newBuilder().setTenantId(TENANT_ID).setEntityType(EntityType.SERVICE.name()).setEntityName("Some Service 2").putAttributes(EntityConstants.getValue(SERVICE_ATTRIBUTE_SERVICE_TYPE), generateAttrValue(HOSTNAME.name())).putIdentifyingAttributes(EntityConstants.getValue(CommonAttribute.COMMON_ATTRIBUTE_FQN), generateRandomUUIDAttrValue()).build();
    Entity createdEntity2 = entityDataServiceClient.upsert(entity2);
    assertNotNull(createdEntity2);
    assertFalse(createdEntity2.getEntityId().trim().isEmpty());
    Entity entity3 = Entity.newBuilder().setTenantId(TENANT_ID).setEntityType(EntityType.SERVICE.name()).setEntityName("Some Service 3").putAttributes(EntityConstants.getValue(SERVICE_ATTRIBUTE_SERVICE_TYPE), generateAttrValue(JAEGER_SERVICE.name())).putIdentifyingAttributes(EntityConstants.getValue(CommonAttribute.COMMON_ATTRIBUTE_FQN), generateRandomUUIDAttrValue()).build();
    Entity createdEntity3 = entityDataServiceClient.upsert(entity3);
    assertNotNull(createdEntity3);
    assertFalse(createdEntity3.getEntityId().trim().isEmpty());
    Entity entity4 = Entity.newBuilder().setTenantId(TENANT_ID).setEntityType(EntityType.SERVICE.name()).setEntityName("Some Service 4").putAttributes(EntityConstants.getValue(SERVICE_ATTRIBUTE_SERVICE_TYPE), generateAttrValue(K8S_SERVICE.name())).putIdentifyingAttributes(EntityConstants.getValue(CommonAttribute.COMMON_ATTRIBUTE_FQN), generateRandomUUIDAttrValue()).build();
    Entity createdEntity4 = entityDataServiceClient.upsert(entity4);
    assertNotNull(createdEntity4);
    assertFalse(createdEntity4.getEntityId().trim().isEmpty());
    Entity entity5 = Entity.newBuilder().setTenantId(TENANT_ID).setEntityType(EntityType.SERVICE.name()).setEntityName("Some Service 5").putAttributes(EntityConstants.getValue(SERVICE_ATTRIBUTE_SERVICE_TYPE), generateAttrValue(JAEGER_SERVICE.name())).putIdentifyingAttributes(EntityConstants.getValue(CommonAttribute.COMMON_ATTRIBUTE_FQN), generateRandomUUIDAttrValue()).build();
    Entity createdEntity5 = entityDataServiceClient.upsert(entity5);
    assertNotNull(createdEntity5);
    assertFalse(createdEntity5.getEntityId().trim().isEmpty());
    EntityQueryRequest queryRequest = EntityQueryRequest.newBuilder().setEntityType(EntityType.SERVICE.name()).setFilter(Filter.newBuilder().setOperatorValue(Operator.NEQ.getNumber()).setLhs(Expression.newBuilder().setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(SERVICE_NAME_ATTR))).setRhs(Expression.newBuilder().setLiteral(LiteralConstant.newBuilder().setValue(org.hypertrace.entity.query.service.v1.Value.newBuilder().setString("Some Service 6").setValueType(STRING))))).addSelection(Expression.newBuilder().setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(SERVICE_TYPE_ATTR)).build()).addSelection(Expression.newBuilder().setFunction(Function.newBuilder().setFunctionName("COUNT").addArguments(Expression.newBuilder().setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(SERVICE_ID_ATTR))))).addGroupBy(Expression.newBuilder().setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(SERVICE_TYPE_ATTR))).addOrderBy(OrderByExpression.newBuilder().setOrder(ASC).setExpression(Expression.newBuilder().setColumnIdentifier(ColumnIdentifier.newBuilder().setColumnName(SERVICE_TYPE_ATTR)))).build();
    // this entity will be filtered out
    Entity entity6 = Entity.newBuilder().setTenantId(TENANT_ID).setEntityType(EntityType.SERVICE.name()).setEntityName("Some Service 6").putAttributes(EntityConstants.getValue(SERVICE_ATTRIBUTE_SERVICE_TYPE), generateAttrValue(JAEGER_SERVICE.name())).putIdentifyingAttributes(EntityConstants.getValue(CommonAttribute.COMMON_ATTRIBUTE_FQN), generateRandomUUIDAttrValue()).build();
    Entity createdEntity6 = entityDataServiceClient.upsert(entity6);
    assertNotNull(createdEntity6);
    assertFalse(createdEntity6.getEntityId().trim().isEmpty());
    Iterator<ResultSetChunk> resultSetChunkIterator = GrpcClientRequestContextUtil.executeWithHeadersContext(HEADERS, () -> entityQueryServiceClient.execute(queryRequest));
    List<ResultSetChunk> list = Lists.newArrayList(resultSetChunkIterator);
    assertEquals(2, list.size());
    assertEquals(2, list.get(0).getRowCount());
    assertEquals(0, list.get(0).getChunkId());
    assertFalse(list.get(0).getIsLastChunk());
    assertEquals(1, list.get(1).getRowCount());
    assertEquals(1, list.get(1).getChunkId());
    assertTrue(list.get(1).getIsLastChunk());
    assertEquals(HOSTNAME.name(), list.get(0).getRow(0).getColumn(0).getString());
    assertEquals(1, list.get(0).getRow(0).getColumn(1).getInt());
    assertEquals(JAEGER_SERVICE.name(), list.get(0).getRow(1).getColumn(0).getString());
    assertEquals(3, list.get(0).getRow(1).getColumn(1).getInt());
    assertEquals(K8S_SERVICE.name(), list.get(1).getRow(0).getColumn(0).getString());
    assertEquals(1, list.get(1).getRow(0).getColumn(1).getInt());
    ResultSetMetadata resultSetMetadata = ResultSetMetadata.newBuilder().addColumnMetadata(ColumnMetadata.newBuilder().setColumnName(SERVICE_TYPE_ATTR).build()).addColumnMetadata(ColumnMetadata.newBuilder().setColumnName("COUNT_" + SERVICE_ID_ATTR).build()).build();
    // verify metadata sent for each chunk
    assertEquals(resultSetMetadata, list.get(0).getResultSetMetadata());
    assertEquals(resultSetMetadata, list.get(1).getResultSetMetadata());
}