diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index cd3dce4959..a8a192dc83 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -176,7 +176,9 @@ def delete_feature_view(self, name: str): return self._registry.delete_feature_view(name, self.project) - def apply(self, objects: List[Union[FeatureView, Entity]]): + def apply( + self, objects: Union[Entity, FeatureView, List[Union[FeatureView, Entity]]] + ): """Register objects to metadata store and update related infrastructure. The apply method registers one or more definitions (e.g., Entity, FeatureView) and registers or updates these @@ -210,6 +212,8 @@ def apply(self, objects: List[Union[FeatureView, Entity]]): # TODO: Add locking # TODO: Optimize by only making a single call (read/write) + if isinstance(objects, Entity) or isinstance(objects, FeatureView): + objects = [objects] views_to_update = [] for ob in objects: if isinstance(ob, FeatureView): diff --git a/sdk/python/tests/test_feature_store.py b/sdk/python/tests/test_feature_store.py index cd60002811..198c287a72 100644 --- a/sdk/python/tests/test_feature_store.py +++ b/sdk/python/tests/test_feature_store.py @@ -79,8 +79,8 @@ def test_apply_entity_success(test_feature_store): labels={"team": "matchmaking"}, ) - # Register Entity with Core - test_feature_store.apply([entity]) + # Register Entity + test_feature_store.apply(entity) entities = test_feature_store.list_entities() @@ -107,7 +107,7 @@ def test_apply_entity_integration(test_feature_store): labels={"team": "matchmaking"}, ) - # Register Entity with Core + # Register Entity test_feature_store.apply([entity]) entities = test_feature_store.list_entities()