Add InfraDiff class for feast plan (feast-dev#2190)

felixwang9817 · web-flow · commit 6c09bc4e7e92 · 2022-01-05T12:18:24.000-08:00
* Add InfraDiff and InfraObjectDiff classes

Signed-off-by: Felix Wang &lt;wangfelix98@gmail.com&gt;

* Add support for Infra objects in the registry

Signed-off-by: Felix Wang &lt;wangfelix98@gmail.com&gt;

* Rename utils to property_diff

Signed-off-by: Felix Wang &lt;wangfelix98@gmail.com&gt;
diff --git a/protos/feast/core/Registry.proto b/protos/feast/core/Registry.proto
@@ -25,6 +25,7 @@ import "feast/core/Entity.proto";
 import "feast/core/FeatureService.proto";
 import "feast/core/FeatureTable.proto";
 import "feast/core/FeatureView.proto";
+import "feast/core/InfraObject.proto";
 import "feast/core/OnDemandFeatureView.proto";
 import "feast/core/RequestFeatureView.proto";
 import "google/protobuf/timestamp.proto";
@@ -36,6 +37,7 @@ message Registry {
     repeated OnDemandFeatureView on_demand_feature_views = 8;
     repeated RequestFeatureView request_feature_views = 9;
     repeated FeatureService feature_services = 7;
+    Infra infra = 10;
 
     string registry_schema_version = 3; // to support migrations; incremented when schema is changed
     string version_id = 4; // version id, random string generated on each update of the data; now used only for debugging purposes
diff --git a/sdk/python/feast/diff/FcoDiff.py b/sdk/python/feast/diff/FcoDiff.py
@@ -1,29 +1,14 @@
 from dataclasses import dataclass
-from enum import Enum
 from typing import Any, Iterable, List, Set, Tuple, TypeVar
 
 from feast.base_feature_view import BaseFeatureView
+from feast.diff.property_diff import PropertyDiff, TransitionType
 from feast.entity import Entity
 from feast.feature_service import FeatureService
 from feast.protos.feast.core.Entity_pb2 import Entity as EntityProto
 from feast.protos.feast.core.FeatureView_pb2 import FeatureView as FeatureViewProto
 
 
-@dataclass
-class PropertyDiff:
-    property_name: str
-    val_existing: str
-    val_declared: str
-
-
-class TransitionType(Enum):
-    UNKNOWN = 0
-    CREATE = 1
-    DELETE = 2
-    UPDATE = 3
-    UNCHANGED = 4
-
-
 @dataclass
 class FcoDiff:
     name: str
diff --git a/sdk/python/feast/diff/infra_diff.py b/sdk/python/feast/diff/infra_diff.py
@@ -0,0 +1,28 @@
+from dataclasses import dataclass
+from typing import Any, List
+
+from feast.diff.property_diff import PropertyDiff, TransitionType
+
+
+@dataclass
+class InfraObjectDiff:
+    name: str
+    infra_object_type: str
+    current_fco: Any
+    new_fco: Any
+    fco_property_diffs: List[PropertyDiff]
+    transition_type: TransitionType
+
+
+@dataclass
+class InfraDiff:
+    infra_object_diffs: List[InfraObjectDiff]
+
+    def __init__(self):
+        self.infra_object_diffs = []
+
+    def update(self):
+        pass
+
+    def to_string(self):
+        pass
diff --git a/sdk/python/feast/diff/property_diff.py b/sdk/python/feast/diff/property_diff.py
@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+from enum import Enum
+
+
+@dataclass
+class PropertyDiff:
+    property_name: str
+    val_existing: str
+    val_declared: str
+
+
+class TransitionType(Enum):
+    UNKNOWN = 0
+    CREATE = 1
+    DELETE = 2
+    UPDATE = 3
+    UNCHANGED = 4
diff --git a/sdk/python/feast/registry.py b/sdk/python/feast/registry.py
@@ -42,6 +42,7 @@
 )
 from feast.feature_service import FeatureService
 from feast.feature_view import FeatureView
+from feast.infra.infra_object import Infra
 from feast.on_demand_feature_view import OnDemandFeatureView
 from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
 from feast.registry_store import NoopRegistryStore
@@ -222,6 +223,36 @@ def _initialize_registry(self):
             registry_proto.registry_schema_version = REGISTRY_SCHEMA_VERSION
             self._registry_store.update_registry_proto(registry_proto)
 
+    def update_infra(self, infra: Infra, project: str, commit: bool = True):
+        """
+        Updates the stored Infra object.
+
+        Args:
+            infra: The new Infra object to be stored.
+            project: Feast project that the Infra object refers to
+            commit: Whether the change should be persisted immediately
+        """
+        self._prepare_registry_for_changes()
+        assert self.cached_registry_proto
+
+        self.cached_registry_proto.infra.CopyFrom(infra.to_proto())
+        if commit:
+            self.commit()
+
+    def get_infra(self, project: str, allow_cache: bool = False) -> Infra:
+        """
+        Retrieves the stored Infra object.
+
+        Args:
+            project: Feast project that the Infra object refers to
+            allow_cache: Whether to allow returning this entity from a cached registry
+
+        Returns:
+            The stored Infra object.
+        """
+        registry_proto = self._get_registry_proto(allow_cache=allow_cache)
+        return Infra.from_proto(registry_proto.infra)
+
     def apply_entity(self, entity: Entity, project: str, commit: bool = True):
         """
         Registers a single entity with Feast
