metadata methods

achals · achals · commit 4618566f9dbc · 2022-06-06T11:19:43.000-07:00
Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;
diff --git a/sdk/python/feast/infra/registry_stores/sql.py b/sdk/python/feast/infra/registry_stores/sql.py
@@ -140,6 +140,7 @@
 
 
 class SqlRegistry(BaseRegistry):
+
     def __init__(
         self, registry_config: Optional[RegistryConfig], repo_path: Optional[Path]
     ):
@@ -474,6 +475,59 @@ def update_infra(self, infra: Infra, project: str, commit: bool = True):
     def get_infra(self, project: str, allow_cache: bool = False) -> Infra:
         pass
 
+    def apply_user_metadata(self, project: str, feature_view: BaseFeatureView, metadata_bytes: Optional[bytes]):
+        if isinstance(feature_view, FeatureView):
+            table = feature_views
+        elif isinstance(feature_view, OnDemandFeatureView):
+            table = on_demand_feature_views
+        elif isinstance(feature_view, RequestFeatureView):
+            table = request_feature_views
+        elif isinstance(feature_view, StreamFeatureView):
+            table = streaming_feature_views
+        else:
+            raise ValueError(f"Unexpected feature view type: {type(feature_view)}")
+
+        name = feature_view.name
+        with self.engine.connect() as conn:
+            stmt = select(table).where(getattr(table.c, "feature_view_name") == name)
+            row = conn.execute(stmt).first()
+            update_datetime = datetime.utcnow()
+            update_time = int(update_datetime.timestamp())
+            if row:
+                values = {
+                    "user_metadata": metadata_bytes,
+                    "last_updated_timestamp": update_time,
+                }
+                update_stmt = (
+                    update(table)
+                        .where(getattr(table.c, "feature_view_name") == name)
+                        .values(values,)
+                )
+                conn.execute(update_stmt)
+            else:
+                raise FeatureViewNotFoundException(feature_view.name, project=project)
+
+    def get_user_metadata(self, project: str, feature_view: BaseFeatureView) -> Optional[bytes]:
+        if isinstance(feature_view, FeatureView):
+            table = feature_views
+        elif isinstance(feature_view, OnDemandFeatureView):
+            table = on_demand_feature_views
+        elif isinstance(feature_view, RequestFeatureView):
+            table = request_feature_views
+        elif isinstance(feature_view, StreamFeatureView):
+            table = streaming_feature_views
+        else:
+            raise ValueError(f"Unexpected feature view type: {type(feature_view)}")
+
+        name = feature_view.name
+        with self.engine.connect() as conn:
+            stmt = select(table).where(getattr(table.c, "feature_view_name") == name)
+            row = conn.execute(stmt).first()
+            if row:
+                return row["user_metadata"]
+            else:
+                raise FeatureViewNotFoundException(feature_view.name, project=project)
+
     def proto(self) -> RegistryProto:
         r = RegistryProto()
         project = ""
@@ -508,6 +562,7 @@ def _apply_object(
             update_time = int(update_datetime.timestamp())
             if hasattr(obj, "last_updated_timestamp"):
                 obj.last_updated_timestamp = update_datetime
+
             if row:
                 values = {
                     proto_field_name: obj.to_proto().SerializeToString(),
@@ -528,6 +583,7 @@ def _apply_object(
                 insert_stmt = insert(table).values(values,)
                 conn.execute(insert_stmt)
 
+
     def _delete_object(self, table, name, project, id_field_name, not_found_exception):
         with self.engine.connect() as conn:
             stmt = delete(table).where(getattr(table.c, id_field_name) == name)
diff --git a/sdk/python/feast/registry.py b/sdk/python/feast/registry.py
@@ -631,6 +631,14 @@ def get_infra(self, project: str, allow_cache: bool = False) -> Infra:
             The stored Infra object.
         """
 
+    @abstractmethod
+    def apply_user_metadata(self, project: str, feature_view: BaseFeatureView, metadata_bytes: Optional[bytes]):
+        ...
+
+    @abstractmethod
+    def get_user_metadata(self, project: str, feature_view: BaseFeatureView) -> Optional[bytes]:
+        ...
+
     @abstractmethod
     def proto(self) -> RegistryProto:
         """
@@ -654,6 +662,12 @@ class Registry(BaseRegistry):
     Registry: A registry allows for the management and persistence of feature definitions and related metadata.
     """
 
+    def apply_user_metadata(self, project: str, feature_view: BaseFeatureView, metadata_bytes: Optional[bytes]):
+        pass
+
+    def get_user_metadata(self, project: str, feature_view: BaseFeatureView) -> Optional[bytes]:
+        pass
+
     # The cached_registry_proto object is used for both reads and writes. In particular,
     # all write operations refresh the cache and modify it in memory; the write must
     # then be persisted to the underlying RegistryStore with a call to commit().
