@@ -19,35 +19,36 @@
 import feast.core.FeatureSetProto.EntitySpec;
 import feast.types.ValueProto.ValueType;
 import java.util.Objects;
-import javax.persistence.EmbeddedId;
-import javax.persistence.Table;
+import javax.persistence.*;
 import lombok.Getter;
 import lombok.Setter;
 
 /** Feast entity object. Contains name and type of the entity. */
 @Getter
 @Setter
 @javax.persistence.Entity
-@Table(name = "entities")
+@Table(
+    name = "entities",
+    uniqueConstraints = @UniqueConstraint(columnNames = {"name", "feature_set_id"}))
 public class Entity {
-  @EmbeddedId private EntityReference reference;
+
+  @Id @GeneratedValue private Long id;
+
+  private String name;
+
+  @ManyToOne(fetch = FetchType.LAZY)
+  private FeatureSet featureSet;
 
   /** Data type of the entity. String representation of {@link ValueType} * */
   private String type;
 
   public Entity() {}
 
   private Entity(String name, ValueType.Enum type) {
-    this.setReference(new EntityReference(name));
+    this.setName(name);
     this.setType(type.toString());
   }
 
-  public static Entity withRef(EntityReference entityRef) {
-    Entity entity = new Entity();
-    entity.setReference(entityRef);
-    return entity;
-  }
-
   public static Entity fromProto(EntitySpec entitySpec) {
     Entity entity = new Entity(entitySpec.getName(), entitySpec.getValueType());
     return entity;
@@ -61,12 +62,12 @@ public boolean equals(Object o) {
     if (o == null || getClass() != o.getClass()) {
       return false;
     }
-    Entity feature = (Entity) o;
-    return getReference().equals(feature.getReference()) && getType().equals(feature.getType());
+    Entity entity = (Entity) o;
+    return getName().equals(entity.getName()) && getType().equals(entity.getType());
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(super.hashCode(), getReference(), getType());
+    return Objects.hash(super.hashCode(), getName(), getType());
   }
 }
@@ -34,10 +34,17 @@
 @Getter
 @Setter
 @Entity
-@Table(name = "features")
+@Table(
+    name = "features",
+    uniqueConstraints = @UniqueConstraint(columnNames = {"name", "feature_set_id"}))
 public class Feature {
 
-  @EmbeddedId private FeatureReference reference;
+  @Id @GeneratedValue private Long id;
+
+  private String name;
+
+  @ManyToOne(fetch = FetchType.LAZY)
+  private FeatureSet featureSet;
 
   /** Data type of the feature. String representation of {@link ValueType} * */
   private String type;
@@ -74,16 +81,10 @@ public class Feature {
   public Feature() {}
 
   private Feature(String name, ValueType.Enum type) {
-    this.setReference(new FeatureReference(name));
+    this.setName(name);
     this.setType(type.toString());
   }
 
-  public static Feature withRef(FeatureReference featureRef) {
-    Feature feature = new Feature();
-    feature.setReference(featureRef);
-    return feature;
-  }
-
   public static Feature fromProto(FeatureSpec featureSpec) {
     Feature feature = new Feature(featureSpec.getName(), featureSpec.getValueType());
     feature.labels = TypeConversion.convertMapToJsonString(featureSpec.getLabelsMap());
@@ -166,7 +167,7 @@ public boolean equals(Object o) {
       return false;
     }
     Feature feature = (Feature) o;
-    return Objects.equals(getReference(), feature.getReference())
+    return Objects.equals(getName(), feature.getName())
         && Objects.equals(labels, feature.labels)
         && Arrays.equals(getPresence(), feature.getPresence())
         && Arrays.equals(getGroupPresence(), feature.getGroupPresence())
@@ -188,6 +189,6 @@ public boolean equals(Object o) {
 
   @Override
   public int hashCode() {
-    return Objects.hash(super.hashCode(), getReference(), getType(), labels);
+    return Objects.hash(super.hashCode(), getName(), getType(), getLabels());
   }
 }