CR updates

achals · achals · commit d8ea76facdae · 2021-12-14T15:23:17.000-08:00
Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;
diff --git a/sdk/python/feast/diff/FcoDiff.py b/sdk/python/feast/diff/FcoDiff.py
@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, List, Set, Tuple, Union, TypeVar
+from typing import Any, List, Set, Tuple, Union, TypeVar, Iterable
 from unittest.mock import Base
 
 from feast.base_feature_view import BaseFeatureView
@@ -9,6 +9,8 @@
 from feast.feature_table import FeatureTable
 from feast.feature_view import FeatureView
 from feast.on_demand_feature_view import OnDemandFeatureView
+from feast.protos.feast.core.Entity_pb2 import Entity as EntityProto
+from feast.protos.feast.core.FeatureView_pb2 import FeatureView as FeatureViewProto
 from feast.request_feature_view import RequestFeatureView
 
 
@@ -35,6 +37,7 @@ class FcoDiff:
     transition_type: TransitionType
 
 
+@dataclass
 class RegistryDiff:
     fco_diffs: List[FcoDiff]
 
@@ -49,8 +52,8 @@ def add_fco_diff(self, fco_diff: FcoDiff):
 
 
 def tag_objects_for_keep_delete_add(
-    existing_objs: Set[T], desired_objs: Set[T]
-) -> Tuple[Set[T], Set[T], Set[T]]:
+    existing_objs: Iterable[T], desired_objs: Iterable[T]
+) -> Tuple[Iterable[T], Iterable[T], Iterable[T]]:
     existing_obj_names = {e.name for e in existing_objs}
     desired_obj_names = {e.name for e in desired_objs}
 
@@ -63,3 +66,23 @@ def tag_objects_for_keep_delete_add(
     }
 
     return objs_to_keep, objs_to_delete, objs_to_add
+
+
+U = TypeVar('U', EntityProto, FeatureViewProto)
+
+
+def tag_proto_objects_for_keep_delete_add(
+        existing_objs: Iterable[U], desired_objs: Iterable[U]
+) -> Tuple[Iterable[U], Iterable[U], Iterable[U]]:
+    existing_obj_names = {e.spec.name for e in existing_objs}
+    desired_obj_names = {e.spec.name for e in desired_objs}
+
+    objs_to_add = [
+        e for e in desired_objs if e.spec.name not in existing_obj_names
+    ]
+    objs_to_keep = [e for e in desired_objs if e.spec.name in existing_obj_names]
+    objs_to_delete = [
+        e for e in existing_objs if e.spec.name not in desired_obj_names
+    ]
+
+    return objs_to_keep, objs_to_delete, objs_to_add
diff --git a/sdk/python/feast/registry.py b/sdk/python/feast/registry.py
@@ -28,7 +28,7 @@
     FcoDiff,
     RegistryDiff,
     TransitionType,
-    tag_objects_for_keep_delete_add,
+    tag_proto_objects_for_keep_delete_add,
 )
 from feast.entity import Entity
 from feast.errors import (
@@ -130,19 +130,10 @@ def __init__(
                 else 0
             )
 
-    @classmethod
-    def from_proto(cls, regsitry_proto: RegistryProto):
-        registry = cls(None, None)
-        registry.cached_registry_proto = regsitry_proto
-        registry.cached_registry_proto_created = datetime.utcnow()
-        registry.cached_registry_proto_ttl = timedelta(days=1)
-        registry.cache_being_updated = True
-        return registry
-
     # TODO(achals): This method needs to be filled out and used in the feast plan/apply methods.
     @staticmethod
     def diff_between(
-        project: str, current_registry: "Registry", new_registry: "Registry"
+        current_registry: RegistryProto, new_registry: RegistryProto
     ) -> RegistryDiff:
         diff = RegistryDiff()
 
@@ -151,9 +142,9 @@ def diff_between(
             entities_to_keep,
             entities_to_delete,
             entities_to_add,
-        ) = tag_objects_for_keep_delete_add(
-            set(current_registry.list_entities(project=project, allow_cache=True)),
-            set(new_registry.list_entities(project=project, allow_cache=True)),
+        ) = tag_proto_objects_for_keep_delete_add(
+            current_registry.entities,
+            new_registry.entities,
         )
 
         for e in entities_to_add:
