feast-dev · feast-ci-bot · Sep 28, 2020 · Sep 23, 2020 · Sep 23, 2020 · Sep 23, 2020
diff --git a/common-test/src/main/java/feast/common/it/DataGenerator.java b/common-test/src/main/java/feast/common/it/DataGenerator.java
@@ -17,6 +17,7 @@
 package feast.common.it;
 
 import com.google.common.collect.ImmutableList;
+import feast.proto.core.EntityProto;
 import feast.proto.core.FeatureSetProto;
 import feast.proto.core.SourceProto;
 import feast.proto.core.StoreProto;
@@ -111,11 +112,24 @@ public static FeatureSetProto.FeatureSpec createFeature(
         .build();
   }
 
-  public static FeatureSetProto.EntitySpec createEntity(
+  public static FeatureSetProto.EntitySpec createEntitySpec(
       String name, ValueProto.ValueType.Enum valueType) {
     return FeatureSetProto.EntitySpec.newBuilder().setName(name).setValueType(valueType).build();
   }
 
+  public static EntityProto.EntitySpecV2 createEntitySpecV2(
+      String name,
+      String description,
+      ValueProto.ValueType.Enum valueType,
+      Map<String, String> labels) {
+    return EntityProto.EntitySpecV2.newBuilder()
+        .setName(name)
+        .setDescription(description)
+        .setValueType(valueType)
+        .putAllLabels(labels)
+        .build();
+  }
+
   public static FeatureSetProto.FeatureSet createFeatureSet(
       SourceProto.Source source,
       String projectName,
@@ -152,7 +166,7 @@ public static FeatureSetProto.FeatureSet createFeatureSet(
                 .putAllLabels(labels)
                 .addAllEntities(
                     entities.entrySet().stream()
-                        .map(entry -> createEntity(entry.getKey(), entry.getValue()))
+                        .map(entry -> createEntitySpec(entry.getKey(), entry.getValue()))
                         .collect(Collectors.toList()))
                 .addAllFeatures(
                     features.entrySet().stream()

diff --git a/common-test/src/main/java/feast/common/it/SimpleCoreClient.java b/common-test/src/main/java/feast/common/it/SimpleCoreClient.java
@@ -35,6 +35,46 @@ public CoreServiceProto.ApplyFeatureSetResponse simpleApplyFeatureSet(
         CoreServiceProto.ApplyFeatureSetRequest.newBuilder().setFeatureSet(featureSet).build());
   }
 
+  public CoreServiceProto.ApplyEntityResponse simpleApplyEntity(
+      String projectName, EntityProto.EntitySpecV2 spec) {
+    return stub.applyEntity(
+        CoreServiceProto.ApplyEntityRequest.newBuilder()
+            .setProject(projectName)
+            .setSpec(spec)
+            .build());
+  }
+
+  public List<EntityProto.Entity> simpleListEntities(String projectName) {
+    return stub.listEntities(
+            CoreServiceProto.ListEntitiesRequest.newBuilder()
+                .setFilter(
+                    CoreServiceProto.ListEntitiesRequest.Filter.newBuilder()
+                        .setProject(projectName)
+                        .build())
+                .build())
+        .getEntitiesList();
+  }
+
+  public List<EntityProto.Entity> simpleListEntities(
+      String projectName, Map<String, String> labels) {
+    return stub.listEntities(
+            CoreServiceProto.ListEntitiesRequest.newBuilder()
+                .setFilter(
+                    CoreServiceProto.ListEntitiesRequest.Filter.newBuilder()
+                        .setProject(projectName)
+                        .putAllLabels(labels)
+                        .build())
+                .build())
+        .getEntitiesList();
+  }
+
+  public List<EntityProto.Entity> simpleListEntities(
+      CoreServiceProto.ListEntitiesRequest.Filter filter) {
+    return stub.listEntities(
+            CoreServiceProto.ListEntitiesRequest.newBuilder().setFilter(filter).build())
+        .getEntitiesList();
+  }
+
   public List<FeatureSetProto.FeatureSet> simpleListFeatureSets(
       String projectName, String featureSetName, Map<String, String> labels) {
     return stub.listFeatureSets(
@@ -82,6 +122,15 @@ public FeatureSetProto.FeatureSet simpleGetFeatureSet(String projectName, String
         .getFeatureSet();
   }
 
+  public EntityProto.Entity simpleGetEntity(String projectName, String name) {
+    return stub.getEntity(
+            CoreServiceProto.GetEntityRequest.newBuilder()
+                .setName(name)
+                .setProject(projectName)
+                .build())
+        .getEntity();
+  }
+
   public void updateFeatureSetStatus(
       String projectName, String name, FeatureSetProto.FeatureSetStatus status) {
     stub.updateFeatureSetStatus(

diff --git a/core/src/main/java/feast/core/dao/EntityRepository.java b/core/src/main/java/feast/core/dao/EntityRepository.java
@@ -0,0 +1,33 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright 2018-2020 The Feast Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package feast.core.dao;
+
+import feast.core.model.EntityV2;
+import java.util.List;
+import org.springframework.data.jpa.repository.JpaRepository;
+
+/** JPA repository supplying EntityV2 objects keyed by id. */
+public interface EntityRepository extends JpaRepository<EntityV2, String> {
+
+  long count();
+
+  // Find all EntityV2s by project
+  List<EntityV2> findAllByProject_Name(String project);
+
+  // Find single EntityV2 by project and name
+  EntityV2 findEntityByNameAndProject_Name(String name, String project);
+}
diff --git a/core/src/main/java/feast/core/grpc/CoreServiceImpl.java b/core/src/main/java/feast/core/grpc/CoreServiceImpl.java
@@ -28,6 +28,7 @@
 import feast.core.service.StatsService;
 import feast.proto.core.CoreServiceGrpc.CoreServiceImplBase;
 import feast.proto.core.CoreServiceProto.*;
+import feast.proto.core.EntityProto.EntitySpecV2;
 import feast.proto.core.FeatureSetProto.FeatureSet;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
@@ -95,6 +96,31 @@ public void getFeatureSet(
     }
   }
 
+  @Override
+  public void getEntity(
+      GetEntityRequest request, StreamObserver<GetEntityResponse> responseObserver) {
+    try {
+      GetEntityResponse response = specService.getEntity(request);
+      responseObserver.onNext(response);
+      responseObserver.onCompleted();
+    } catch (RetrievalException e) {
+      log.error("Unable to fetch entity requested in GetEntity method: ", e);
+      responseObserver.onError(
+          Status.NOT_FOUND.withDescription(e.getMessage()).withCause(e).asRuntimeException());
+    } catch (IllegalArgumentException e) {
+      log.error("Illegal arguments provided to GetEntity method: ", e);
+      responseObserver.onError(
+          Status.INVALID_ARGUMENT
+              .withDescription(e.getMessage())
+              .withCause(e)
+              .asRuntimeException());
+    } catch (Exception e) {
+      log.error("Exception has occurred in GetEntity method: ", e);
+      responseObserver.onError(
+          Status.INTERNAL.withDescription(e.getMessage()).withCause(e).asRuntimeException());
+    }
+  }
+
   @Override
   public void listFeatureSets(
       ListFeatureSetsRequest request, StreamObserver<ListFeatureSetsResponse> responseObserver) {
@@ -135,6 +161,32 @@ public void listFeatures(
     }
   }
 
+  /** Retrieve a list of entities */
+  @Override
+  public void listEntities(
+      ListEntitiesRequest request, StreamObserver<ListEntitiesResponse> responseObserver) {
+    try {
+      ListEntitiesResponse response = specService.listEntities(request.getFilter());
+      responseObserver.onNext(response);
+      responseObserver.onCompleted();
+    } catch (IllegalArgumentException e) {
+      log.error("Illegal arguments provided to ListEntities method: ", e);
+      responseObserver.onError(
+          Status.INVALID_ARGUMENT
+              .withDescription(e.getMessage())
+              .withCause(e)
+              .asRuntimeException());
+    } catch (RetrievalException e) {
+      log.error("Unable to fetch entities requested in ListEntities method: ", e);
+      responseObserver.onError(
+          Status.NOT_FOUND.withDescription(e.getMessage()).withCause(e).asRuntimeException());
+    } catch (Exception e) {
+      log.error("Exception has occurred in ListEntities method: ", e);
+      responseObserver.onError(
+          Status.INTERNAL.withDescription(e.getMessage()).withCause(e).asRuntimeException());
+    }
+  }
+
   @Override
   public void getFeatureStatistics(
       GetFeatureStatisticsRequest request,
@@ -191,6 +243,41 @@ public void listStores(
     }
   }
 
+  /* Registers an entity to Feast Core */
+  @Override
+  public void applyEntity(
+      ApplyEntityRequest request, StreamObserver<ApplyEntityResponse> responseObserver) {
+
+    String projectId = null;
+
+    try {
+      EntitySpecV2 spec = request.getSpec();
+      projectId = request.getProject();
+      authorizationService.authorizeRequest(SecurityContextHolder.getContext(), projectId);
+      ApplyEntityResponse response = specService.applyEntity(spec, projectId);
+      responseObserver.onNext(response);
+      responseObserver.onCompleted();
+    } catch (org.hibernate.exception.ConstraintViolationException e) {
+      log.error(
+          "Unable to persist this entity due to a constraint violation. Please ensure that"
+              + " field names are unique within the project namespace: ",
+          e);
+      responseObserver.onError(
+          Status.ALREADY_EXISTS.withDescription(e.getMessage()).withCause(e).asRuntimeException());
+    } catch (AccessDeniedException e) {
+      log.info(String.format("User prevented from accessing project: %s", projectId));
+      responseObserver.onError(
+          Status.PERMISSION_DENIED
+              .withDescription(e.getMessage())
+              .withCause(e)
+              .asRuntimeException());
+    } catch (Exception e) {
+      log.error("Exception has occurred in ApplyEntity method: ", e);
+      responseObserver.onError(
+          Status.INTERNAL.withDescription(e.getMessage()).withCause(e).asRuntimeException());
+    }
+  }
+
   @Override
   public void applyFeatureSet(
       ApplyFeatureSetRequest request, StreamObserver<ApplyFeatureSetResponse> responseObserver) {