use of org.hypertrace.entity.data.service.v1.EntityRelationship in project entity-service by hypertrace.
the class EntityDataServiceImpl method getRelationships.
@Override
public void getRelationships(RelationshipsQuery query, StreamObserver<EntityRelationship> responseObserver) {
logQuery(query);
Optional<String> tenantId = RequestContext.CURRENT.get().getTenantId();
if (tenantId.isEmpty()) {
responseObserver.onError(new ServiceException("Tenant id is missing in the request."));
return;
}
org.hypertrace.core.documentstore.Query docStoreQuery = new org.hypertrace.core.documentstore.Query();
List<Filter> filters = new ArrayList<>();
filters.add(DocStoreConverter.getTenantIdEqFilter(tenantId.get()));
if (query.getEntityRelationshipCount() > 0) {
filters.add(new Filter(Filter.Op.IN, EntityServiceConstants.ENTITY_RELATIONSHIP_TYPE, query.getEntityRelationshipList()));
}
if (query.getFromEntityIdCount() > 0) {
filters.add(new Filter(Filter.Op.IN, EntityServiceConstants.FROM_ENTITY_ID, query.getFromEntityIdList()));
}
if (query.getToEntityIdCount() > 0) {
filters.add(new Filter(Filter.Op.IN, EntityServiceConstants.TO_ENTITY_ID, query.getToEntityIdList()));
}
if (!filters.isEmpty()) {
if (filters.size() == 1) {
docStoreQuery.setFilter(filters.get(0));
} else {
Filter f = new Filter();
f.setOp(Filter.Op.AND);
f.setChildFilters(filters.toArray(new Filter[] {}));
docStoreQuery.setFilter(f);
}
}
searchByQueryAndStreamRelationships(docStoreQuery, responseObserver, tenantId.get());
}
use of org.hypertrace.entity.data.service.v1.EntityRelationship in project entity-service by hypertrace.
the class EntityDataServiceImpl method upsertRelationships.
@Override
public void upsertRelationships(EntityRelationships request, StreamObserver<Empty> responseObserver) {
String tenantId = RequestContext.CURRENT.get().getTenantId().orElse(null);
if (tenantId == null) {
responseObserver.onError(new ServiceException("Tenant id is missing in the request."));
return;
}
Map<Key, Document> entityRelations = new HashMap<>();
try {
for (EntityRelationship relationship : request.getRelationshipList()) {
if (StringUtils.isEmpty(relationship.getFromEntityId()) || StringUtils.isEmpty(relationship.getToEntityId()) || StringUtils.isEmpty(relationship.getEntityRelationshipType())) {
LOG.warn("Invalid relationship upsert request:{}", relationship);
continue;
}
Document document = convertEntityRelationshipToDocument(tenantId, relationship);
Key key = new EntityRelationshipDocKey(tenantId, relationship.getEntityRelationshipType(), relationship.getFromEntityId(), relationship.getToEntityId());
entityRelations.put(key, document);
}
boolean status = relationshipsCollection.bulkUpsert(entityRelations);
if (status) {
responseObserver.onNext(Empty.newBuilder().build());
responseObserver.onCompleted();
} else {
responseObserver.onError(new RuntimeException("Could not bulk upsert relationships."));
}
} catch (IOException e) {
LOG.error("Failed to bulk upsert relationships.", e);
responseObserver.onError(e);
}
}
use of org.hypertrace.entity.data.service.v1.EntityRelationship in project entity-service by hypertrace.
the class EntityDataServiceRelationshipsTest method testCreateAndGetEntityRelationship.
@Test
public void testCreateAndGetEntityRelationship() {
// Create two different relationships.
EntityRelationship relationship1 = EntityRelationship.newBuilder().setFromEntityId(UUID.randomUUID().toString()).setToEntityId(// container id
UUID.randomUUID().toString()).setEntityRelationshipType("POD_CONTAINER").build();
entityDataServiceClient.upsertRelationships(TENANT_ID, EntityRelationships.newBuilder().addRelationship(relationship1).build());
EntityRelationship relationship2 = EntityRelationship.newBuilder().setFromEntityId(// replicaset id
UUID.randomUUID().toString()).setToEntityId(relationship1.getFromEntityId()).setEntityRelationshipType("REPLICASET_POD").build();
entityDataServiceClient.upsertRelationships(TENANT_ID, EntityRelationships.newBuilder().addRelationship(relationship2).build());
// Query all the relationships and verify.
List<EntityRelationship> relationships = new ArrayList<>();
entityDataServiceClient.getRelationships(TENANT_ID, null, null, null).forEachRemaining(relationships::add);
Assertions.assertEquals(2, relationships.size());
// Query only one relationship and verify.
relationships = new ArrayList<>();
entityDataServiceClient.getRelationships(TENANT_ID, Set.of(relationship1.getEntityRelationshipType()), null, null).forEachRemaining(relationships::add);
Assertions.assertEquals(1, relationships.size());
relationships = new ArrayList<>();
entityDataServiceClient.getRelationships(TENANT_ID, Set.of(relationship2.getEntityRelationshipType()), null, null).forEachRemaining(relationships::add);
Assertions.assertEquals(1, relationships.size());
relationships = new ArrayList<>();
entityDataServiceClient.getRelationships(TENANT_ID, null, Collections.singleton(relationship1.getFromEntityId()), null).forEachRemaining(relationships::add);
Assertions.assertEquals(1, relationships.size());
relationships = new ArrayList<>();
entityDataServiceClient.getRelationships(TENANT_ID, null, null, Collections.singleton(relationship1.getToEntityId())).forEachRemaining(relationships::add);
Assertions.assertEquals(1, relationships.size());
}
Aggregations