Pass entities information to Provider (feast-dev#1498)

Signed-off-by: Matt Delacour <[email protected]>

Author: MattDelac

sdk/python/feast/feature_store.py

@@ -217,12 +217,14 @@ def apply(
         if isinstance(objects, Entity) or isinstance(objects, FeatureView):
             objects = [objects]
         views_to_update = []
+        entities_to_update = []
         for ob in objects:
             if isinstance(ob, FeatureView):
                 self._registry.apply_feature_view(ob, project=self.project)
                 views_to_update.append(ob)
             elif isinstance(ob, Entity):
                 self._registry.apply_entity(ob, project=self.project)
+                entities_to_update.append(ob)
             else:
                 raise ValueError(
                     f"Unknown object type ({type(ob)}) provided as part of apply() call"
@@ -231,6 +233,8 @@ def apply(
             project=self.project,
             tables_to_delete=[],
             tables_to_keep=views_to_update,
+            entities_to_delete=[],
+            entities_to_keep=entities_to_update,
             partial=True,
         )
 

sdk/python/feast/infra/gcp.py

@@ -9,6 +9,7 @@
 from google.auth.exceptions import DefaultCredentialsError
 
 from feast import FeatureTable, utils
+from feast.entity import Entity
 from feast.errors import FeastProviderLoginError
 from feast.feature_view import FeatureView
 from feast.infra.key_encoding_utils import serialize_entity_key
@@ -55,6 +56,8 @@ def update_infra(
         project: str,
         tables_to_delete: Sequence[Union[FeatureTable, FeatureView]],
         tables_to_keep: Sequence[Union[FeatureTable, FeatureView]],
+        entities_to_delete: Sequence[Entity],
+        entities_to_keep: Sequence[Entity],
         partial: bool,
     ):
         from google.cloud import datastore
@@ -77,7 +80,10 @@ def update_infra(
             client.delete(key)
 
     def teardown_infra(
-        self, project: str, tables: Sequence[Union[FeatureTable, FeatureView]]
+        self,
+        project: str,
+        tables: Sequence[Union[FeatureTable, FeatureView]],
+        entities: Sequence[Entity],
     ) -> None:
         client = self._initialize_client()
 

sdk/python/feast/infra/local.py

@@ -8,6 +8,7 @@
 import pytz
 
 from feast import FeatureTable, utils
+from feast.entity import Entity
 from feast.feature_view import FeatureView
 from feast.infra.key_encoding_utils import serialize_entity_key
 from feast.infra.offline_stores.helpers import get_offline_store_from_sources
@@ -50,6 +51,8 @@ def update_infra(
         project: str,
         tables_to_delete: Sequence[Union[FeatureTable, FeatureView]],
         tables_to_keep: Sequence[Union[FeatureTable, FeatureView]],
+        entities_to_delete: Sequence[Entity],
+        entities_to_keep: Sequence[Entity],
         partial: bool,
     ):
         conn = self._get_conn()
@@ -65,7 +68,10 @@ def update_infra(
             conn.execute(f"DROP TABLE IF EXISTS {_table_id(project, table)}")
 
     def teardown_infra(
-        self, project: str, tables: Sequence[Union[FeatureTable, FeatureView]]
+        self,
+        project: str,
+        tables: Sequence[Union[FeatureTable, FeatureView]],
+        entities: Sequence[Entity],
     ) -> None:
         os.unlink(self._db_path)
 

sdk/python/feast/infra/provider.py

@@ -26,6 +26,8 @@ def update_infra(
         project: str,
         tables_to_delete: Sequence[Union[FeatureTable, FeatureView]],
         tables_to_keep: Sequence[Union[FeatureTable, FeatureView]],
+        entities_to_delete: Sequence[Entity],
+        entities_to_keep: Sequence[Entity],
         partial: bool,
     ):
         """
@@ -37,21 +39,29 @@ def update_infra(
                 clean up the corresponding cloud resources.
             tables_to_keep: Tables that are still in the feature repo. Depending on implementation,
                 provider may or may not need to update the corresponding resources.
+            entities_to_delete: Entities that were deleted from the feature repo, so provider needs to
+                clean up the corresponding cloud resources.
+            entities_to_keep: Entities that are still in the feature repo. Depending on implementation,
+                provider may or may not need to update the corresponding resources.
             partial: if true, then tables_to_delete and tables_to_keep are *not* exhaustive lists.
                 There may be other tables that are not touched by this update.
         """
         ...
 
     @abc.abstractmethod
     def teardown_infra(
-        self, project: str, tables: Sequence[Union[FeatureTable, FeatureView]]
+        self,
+        project: str,
+        tables: Sequence[Union[FeatureTable, FeatureView]],
+        entities: Sequence[Entity],
     ):
         """
         Tear down all cloud resources for a repo.
 
         Args:
             project: Feast project to which tables belong
             tables: Tables that are declared in the feature repo.
+            entities: Entities that are declared in the feature repo.
         """
         ...
 

sdk/python/feast/repo_operations.py

@@ -129,6 +129,14 @@ def apply_total(repo_config: RepoConfig, repo_path: Path):
     all_to_keep.extend(repo.feature_tables)
     all_to_keep.extend(repo.feature_views)
 
+    entities_to_delete: List[Entity] = []
+    repo_entities_names = set([e.name for e in repo.entities])
+    for registry_entity in registry.list_entities(project=project):
+        if registry_entity.name not in repo_entities_names:
+            entities_to_delete.append(registry_entity)
+
+    entities_to_keep: List[Entity] = repo.entities
+
     for name in [view.name for view in repo.feature_tables] + [
         table.name for table in repo.feature_views
     ]:
@@ -141,10 +149,13 @@ def apply_total(repo_config: RepoConfig, repo_path: Path):
         click.echo(
             f"Removing infrastructure for {Style.BRIGHT + Fore.GREEN}{name}{Style.RESET_ALL}"
         )
+
     infra_provider.update_infra(
         project,
         tables_to_delete=all_to_delete,
         tables_to_keep=all_to_keep,
+        entities_to_delete=entities_to_delete,
+        entities_to_keep=entities_to_keep,
         partial=False,
     )
 
@@ -160,8 +171,13 @@ def teardown(repo_config: RepoConfig, repo_path: Path):
     registry_tables: List[Union[FeatureTable, FeatureView]] = []
     registry_tables.extend(registry.list_feature_tables(project=project))
     registry_tables.extend(registry.list_feature_views(project=project))
+
+    registry_entities: List[Entity] = registry.list_entities(project=project)
+
     infra_provider = get_provider(repo_config, repo_path)
-    infra_provider.teardown_infra(project, tables=registry_tables)
+    infra_provider.teardown_infra(
+        project, tables=registry_tables, entities=registry_entities
+    )
 
 
 def registry_dump(repo_config: RepoConfig, repo_path: Path):

sdk/python/tests/online_write_benchmark.py

@@ -89,7 +89,10 @@ def benchmark_writes():
             )
 
         registry_tables = store.list_feature_views()
-        provider.teardown_infra(store.project, tables=registry_tables)
+        registry_entities = store.list_entities()
+        provider.teardown_infra(
+            store.project, tables=registry_tables, entities=registry_entities
+        )
 
 
 if __name__ == "__main__":