diff --git a/entity-service-api/src/main/java/org/hypertrace/entity/service/constants/EntityCollectionConstants.java b/entity-service-api/src/main/java/org/hypertrace/entity/service/constants/EntityCollectionConstants.java index 2d245201..22978d23 100644 --- a/entity-service-api/src/main/java/org/hypertrace/entity/service/constants/EntityCollectionConstants.java +++ b/entity-service-api/src/main/java/org/hypertrace/entity/service/constants/EntityCollectionConstants.java @@ -10,4 +10,5 @@ public class EntityCollectionConstants { public static final String RAW_ENTITIES_COLLECTION = "raw_entities"; public static final String ENRICHED_ENTITIES_COLLECTION = "enriched_entities"; public static final String ENTITY_RELATIONSHIPS_COLLECTION = "entity_relationships"; + public static final String ENTITY_LABELS_COLLECTION = "entity_labels"; } diff --git a/entity-service-api/src/main/proto/org/hypertrace/entity/data/service/v1/entity_label.proto b/entity-service-api/src/main/proto/org/hypertrace/entity/data/service/v1/entity_label.proto index 1dd888e4..815a3565 100644 --- a/entity-service-api/src/main/proto/org/hypertrace/entity/data/service/v1/entity_label.proto +++ b/entity-service-api/src/main/proto/org/hypertrace/entity/data/service/v1/entity_label.proto @@ -9,6 +9,8 @@ message EntityLabel { string name = 2; // TODO: This is up for debate. string color = 3; + // TODO: I don't like having this here in the data transfer object. Will deal with it after. + string tenant_id = 4; } message EntityLabelByIdRequest { diff --git a/entity-service-impl/src/main/java/org/hypertrace/entity/query/service/EntityQueryServiceImpl.java b/entity-service-impl/src/main/java/org/hypertrace/entity/query/service/EntityQueryServiceImpl.java index 5b088d5f..b24e9102 100644 --- a/entity-service-impl/src/main/java/org/hypertrace/entity/query/service/EntityQueryServiceImpl.java +++ b/entity-service-impl/src/main/java/org/hypertrace/entity/query/service/EntityQueryServiceImpl.java @@ -1,9 +1,13 @@ package org.hypertrace.entity.query.service; import static java.util.stream.Collectors.toUnmodifiableMap; +import static org.hypertrace.entity.service.constants.EntityCollectionConstants.ENTITY_LABELS_COLLECTION; import static org.hypertrace.entity.service.constants.EntityCollectionConstants.RAW_ENTITIES_COLLECTION; +import com.google.protobuf.Descriptors; +import com.google.protobuf.GeneratedMessageV3; import com.google.protobuf.InvalidProtocolBufferException; +import com.google.protobuf.Message; import com.google.protobuf.ServiceException; import com.typesafe.config.Config; import io.grpc.stub.StreamObserver; @@ -17,11 +21,13 @@ import org.hypertrace.core.documentstore.Collection; import org.hypertrace.core.documentstore.Datastore; import org.hypertrace.core.documentstore.Document; +import org.hypertrace.core.documentstore.Filter; import org.hypertrace.core.documentstore.JSONDocument; import org.hypertrace.core.documentstore.SingleValueKey; import org.hypertrace.core.grpcutils.context.RequestContext; import org.hypertrace.entity.data.service.v1.AttributeValue; import org.hypertrace.entity.data.service.v1.Entity; +import org.hypertrace.entity.data.service.v1.EntityLabel; import org.hypertrace.entity.data.service.v1.Query; import org.hypertrace.entity.query.service.v1.ColumnIdentifier; import org.hypertrace.entity.query.service.v1.ColumnMetadata; @@ -41,6 +47,7 @@ import org.hypertrace.entity.service.util.DocStoreJsonFormat; import org.hypertrace.entity.service.util.DocStoreJsonFormat.Parser; import org.hypertrace.entity.service.util.StringUtils; +import org.hypertrace.entity.service.util.UUIDGenerator; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -51,11 +58,13 @@ public class EntityQueryServiceImpl extends EntityQueryServiceImplBase { private static final Parser PARSER = DocStoreJsonFormat.parser().ignoringUnknownFields(); private final Collection entitiesCollection; + private final Collection entityLabelsCollection; private final Map> attrNameToEDSAttrMap; public EntityQueryServiceImpl(Datastore datastore, Config config) { this( datastore.getCollection(RAW_ENTITIES_COLLECTION), + datastore.getCollection(ENTITY_LABELS_COLLECTION), config.getConfigList(ATTRIBUTE_MAP_CONFIG_PATH) .stream() .collect(toUnmodifiableMap( @@ -70,9 +79,10 @@ public EntityQueryServiceImpl(Datastore datastore, Config config) { ))); } - public EntityQueryServiceImpl(Collection entitiesCollection, + public EntityQueryServiceImpl(Collection entitiesCollection, Collection entityLabelsCollection, Map> attrNameToEDSAttrMap) { this.entitiesCollection = entitiesCollection; + this.entityLabelsCollection = entityLabelsCollection; this.attrNameToEDSAttrMap = attrNameToEDSAttrMap; } @@ -232,4 +242,223 @@ key, subDocPath, new JSONDocument(jsonValue))) { } } } + + @Override + public void createEntityLabel(EntityLabel request, io.grpc.stub.StreamObserver responseObserver) { + Optional tenantId = RequestContext.CURRENT.get().getTenantId(); + if (tenantId.isEmpty()) { + responseObserver.onError(new ServiceException("Tenant id is missing in the request.")); + return; + } + + String entityLabelId = generateEntityLabelId(tenantId.get(), request.getName()); + EntityLabel entityLabel = EntityLabel.newBuilder(request) + .setTenantId(tenantId.get()) + .setId(entityLabelId) + .build(); + + try { + Document document = convertProtoMessageToDocument(entityLabel); + entityLabelsCollection.upsert(new SingleValueKey(tenantId.get(), entityLabelId), document); + // Will also invoke the sending of the object in the response + searchByIdAndStreamSingleResponse(tenantId.get(), entityLabelId, + entityLabelsCollection, EntityLabel.newBuilder(), responseObserver); + } catch (IOException e) { + responseObserver.onError(new RuntimeException("Could not create entity label.", e)); + } + + //TODO: Optimize this later. For now converting to EDS Query and then again to DocStore Query. +// Query query = EntityQueryConverter +// .convertToEDSQuery(request, attrNameToEDSAttrMap.get(request.getEntityType())); +// Iterator documentIterator = entitiesCollection.search( +// DocStoreConverter.transform(tenantId.get(), query)); +// List entities = convertDocsToEntities(documentIterator); + + //Build result + //TODO : chunk response. For now sending everything in one chunk +// responseObserver.onNext(); +// responseObserver.onCompleted(); + } + + @Override + public void getEntityLabel(org.hypertrace.entity.data.service.v1.EntityLabelByIdRequest request, io.grpc.stub.StreamObserver responseObserver) { + Optional tenantId = RequestContext.CURRENT.get().getTenantId(); + if (tenantId.isEmpty()) { + responseObserver.onError(new ServiceException("Tenant id is missing in the request.")); + return; + } + +// String entityLabelId = generateEntityLabelId(tenantId.get(), request.getName()); +// EntityLabel entityLabel = EntityLabel.newBuilder(request) +// .setTenantId(tenantId.get()) +// .setId(entityLabelId) +// .build(); + searchByIdAndStreamSingleResponse(tenantId.get(), request.getId(), + entityLabelsCollection, EntityLabel.newBuilder(), responseObserver); + +// try { +//// Document document = convertProtoMessageToDocument(entityLabel); +//// entityLabelsCollection.upsert(new SingleValueKey(tenantId.get(), entityLabelId), document); +// // Will also invoke the sending of the object in the response +// searchByIdAndStreamSingleResponse(tenantId.get(), request.getId(), +// entityLabelsCollection, EntityLabel.newBuilder(), responseObserver); +// } catch (IOException e) { +// responseObserver.onError(new RuntimeException("Could not create entity label.", e)); +// } + } + + @Override + public void updateEntityLabel(org.hypertrace.entity.data.service.v1.EntityLabel request, io.grpc.stub.StreamObserver responseObserver) { + super.updateEntityLabel(request, responseObserver); + } + + @Override + public void getAllEntityLabels(org.hypertrace.entity.data.service.v1.Empty request, io.grpc.stub.StreamObserver responseObserver) { + Optional tenantId = RequestContext.CURRENT.get().getTenantId(); + if (tenantId.isEmpty()) { + responseObserver.onError(new ServiceException("Tenant id is missing in the request.")); + return; + } + +// String entityLabelId = generateEntityLabelId(tenantId.get(), request.getName()); +// EntityLabel entityLabel = EntityLabel.newBuilder(request) +// .setTenantId(tenantId.get()) +// .setId(entityLabelId) +// .build(); + searchAndStreamEntityLabels(tenantId.get(), entityLabelsCollection, EntityLabel.newBuilder(), + responseObserver); +// searchByIdAndStreamSingleResponse(tenantId.get(), request.getId(), +// entityLabelsCollection, EntityLabel.newBuilder(), responseObserver); + } + + @Override + public void addEntityLabelToEntity(org.hypertrace.entity.query.service.v1.EntityIdAndLabelId request, io.grpc.stub.StreamObserver responseObserver) { + super.addEntityLabelToEntity(request, responseObserver); + } + + @Override + public void removeEntityLabelFromEntity(org.hypertrace.entity.query.service.v1.EntityIdAndLabelId request, io.grpc.stub.StreamObserver responseObserver) { + super.removeEntityLabelFromEntity(request, responseObserver); + } + + @Override + public void getEntitiesByLabel(org.hypertrace.entity.data.service.v1.Empty request, io.grpc.stub.StreamObserver responseObserver) { + super.getEntitiesByLabel(request, responseObserver); + } + + private String generateEntityLabelId(String tenantId, String labelName) { + // TODO: Not considering Label attributes such as tenant id and name for now. + return UUIDGenerator.getRandomUUID(); + } + + // TODO: Copied from EntityDataServiceImpl + private JSONDocument convertProtoMessageToDocument(T message) + throws IOException { + try { + return DocStoreConverter.transform(message); + } catch (IOException e) { + LOG.error("Could not covert the attributes into JSON doc.", e); + throw e; + } + } + + // TODO: Copied from EntityDataServiceImpl. Should make this generic and stop using entity or entities as variable names + private void searchByIdAndStreamSingleResponse( + String tenantId, String entityLabelId, Collection collection, Message.Builder builder, + StreamObserver responseObserver) { + org.hypertrace.core.documentstore.Query query = new org.hypertrace.core.documentstore.Query(); + String docId = new SingleValueKey(tenantId, entityLabelId).toString(); + query.setFilter(new Filter(Filter.Op.EQ, EntityServiceConstants.ID, docId)); + + Iterator result = collection.search(query); + List entities = new ArrayList<>(); + while (result.hasNext()) { + Document next = result.next(); + Message.Builder b = builder.clone(); + try { + PARSER.merge(next.toJson(), b); + +// // Populate the tenant id field with the tenant id that's received for backward +// // compatibility. +// Descriptors.FieldDescriptor fieldDescriptor = +// b.getDescriptorForType().findFieldByName("tenant_id"); +// if (fieldDescriptor != null) { +// b.setField(fieldDescriptor, tenantId); +// } + } catch (InvalidProtocolBufferException e) { + LOG.error("Could not deserialize the document into an entity label.", e); + } + + entities.add((T) b.build()); + } + + if (LOG.isDebugEnabled()) { + LOG.debug("MongoDB 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(); + } + } + + private void searchAndStreamEntityLabels( + String tenantId, Collection collection, Message.Builder builder, + StreamObserver responseObserver) { + org.hypertrace.core.documentstore.Query query = new org.hypertrace.core.documentstore.Query(); + //String docId = new SingleValueKey(tenantId, entityLabelId).toString(); + query.setFilter(new Filter(Filter.Op.EQ, "tenant_id", tenantId)); + + Iterator result = collection.search(query); + //List entities = new ArrayList<>(); + while (result.hasNext()) { + Document next = result.next(); + Message.Builder b = builder.clone(); + try { + PARSER.merge(next.toJson(), b); + + responseObserver.onNext((T) b.build()); +// // Populate the tenant id field with the tenant id that's received for backward +// // compatibility. +// Descriptors.FieldDescriptor fieldDescriptor = +// b.getDescriptorForType().findFieldByName("tenant_id"); +// if (fieldDescriptor != null) { +// b.setField(fieldDescriptor, tenantId); +// } + } catch (InvalidProtocolBufferException e) { + LOG.error("Could not deserialize the document into an entity label.", e); + } + + //entities.add((T) b.build()); + } + + responseObserver.onCompleted(); + +// if (LOG.isDebugEnabled()) { +// LOG.debug("MongoDB 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(); +// } + } } diff --git a/entity-service-impl/src/main/java/org/hypertrace/entity/service/util/DocStoreConverter.java b/entity-service-impl/src/main/java/org/hypertrace/entity/service/util/DocStoreConverter.java index 0eae7f38..0e784788 100644 --- a/entity-service-impl/src/main/java/org/hypertrace/entity/service/util/DocStoreConverter.java +++ b/entity-service-impl/src/main/java/org/hypertrace/entity/service/util/DocStoreConverter.java @@ -30,16 +30,16 @@ public class DocStoreConverter { private static DocStoreJsonFormat.Printer JSONFORMAT_PRINTER = DocStoreJsonFormat.printer(); /** - * Transforms entity to JSONDocument + * Transforms protoMessage to JSONDocument * - * @param entity + * @param protoMessage * @return */ - public static JSONDocument transform(T entity) throws IOException { + public static JSONDocument transform(T protoMessage) throws IOException { // We need to use patched json converter because // the one from protobuf serializes 64 bit numbers into strings. // See https://github.com/protocolbuffers/protobuf/issues/1823 - String json = JSONFORMAT_PRINTER.print(entity); + String json = JSONFORMAT_PRINTER.print(protoMessage); return new JSONDocument(json); } diff --git a/entity-service-impl/src/main/java/org/hypertrace/entity/service/util/UUIDGenerator.java b/entity-service-impl/src/main/java/org/hypertrace/entity/service/util/UUIDGenerator.java index 7e4c1f94..15f502f3 100644 --- a/entity-service-impl/src/main/java/org/hypertrace/entity/service/util/UUIDGenerator.java +++ b/entity-service-impl/src/main/java/org/hypertrace/entity/service/util/UUIDGenerator.java @@ -38,6 +38,10 @@ public static String generateUUID(Map attributes) { return getUUIDWithVersion3(transform(attributes).toString()).toString(); } + public static String getRandomUUID() { + return java.util.UUID.randomUUID().toString(); + } + /** * Explicitly set the version of UUID to UUIDv3 *