wip: independent labels impl

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";
 }

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 {

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<String, Map<String, String>> 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<String, Map<String, String>> 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<EntityLabel> responseObserver) {
+    Optional<String> 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<Document> documentIterator = entitiesCollection.search(
+//        DocStoreConverter.transform(tenantId.get(), query));
+//    List<Entity> 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<org.hypertrace.entity.data.service.v1.EntityLabel> responseObserver) {
+    Optional<String> 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<org.hypertrace.entity.data.service.v1.EntityLabel> responseObserver) {
+    super.updateEntityLabel(request, responseObserver);
+  }
+
+  @Override
+  public void getAllEntityLabels(org.hypertrace.entity.data.service.v1.Empty request, io.grpc.stub.StreamObserver<org.hypertrace.entity.data.service.v1.EntityLabel> responseObserver) {
+    Optional<String> 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<org.hypertrace.entity.data.service.v1.Entity> responseObserver) {
+    super.addEntityLabelToEntity(request, responseObserver);
+  }
+
+  @Override
+  public void removeEntityLabelFromEntity(org.hypertrace.entity.query.service.v1.EntityIdAndLabelId request, io.grpc.stub.StreamObserver<org.hypertrace.entity.data.service.v1.Entity> responseObserver) {
+    super.removeEntityLabelFromEntity(request, responseObserver);
+  }
+
+  @Override
+  public void getEntitiesByLabel(org.hypertrace.entity.data.service.v1.Empty request, io.grpc.stub.StreamObserver<org.hypertrace.entity.query.service.v1.EntitiesByLabel> 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 <T extends GeneratedMessageV3> 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 <T extends Message> void searchByIdAndStreamSingleResponse(
+      String tenantId, String entityLabelId, Collection collection, Message.Builder builder,
+      StreamObserver<T> 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<Document> result = collection.search(query);
+    List<T> 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 <T extends Message> void searchAndStreamEntityLabels(
+      String tenantId, Collection collection, Message.Builder builder,
+      StreamObserver<T> 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<Document> result = collection.search(query);
+    //List<T> 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();
+//    }
+  }
 }

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 <T extends GeneratedMessageV3> JSONDocument transform(T entity) throws IOException {
+  public static <T extends GeneratedMessageV3> 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);
   }

entity-service-impl/src/main/java/org/hypertrace/entity/service/util/UUIDGenerator.java

@@ -38,6 +38,10 @@ public static String generateUUID(Map<String, AttributeValue> 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
    * <p>