From 4ec680de24f4cd0dbc4f565354c2f368683cd07f Mon Sep 17 00:00:00 2001 From: Danny Chiao Date: Wed, 14 Dec 2022 20:18:19 -0500 Subject: [PATCH] fix: Fix bug where SQL registry was incorrectly writing info around sqlite online store Signed-off-by: Danny Chiao --- sdk/python/feast/infra/registry/sql.py | 178 +++++++++++---------- sdk/python/tests/unit/test_sql_registry.py | 28 ++++ 2 files changed, 122 insertions(+), 84 deletions(-) diff --git a/sdk/python/feast/infra/registry/sql.py b/sdk/python/feast/infra/registry/sql.py index 7867cdff4c3..5e980e90b35 100644 --- a/sdk/python/feast/infra/registry/sql.py +++ b/sdk/python/feast/infra/registry/sql.py @@ -207,14 +207,14 @@ def get_stream_feature_view( self, name: str, project: str, allow_cache: bool = False ): return self._get_object( - stream_feature_views, - name, - project, - StreamFeatureViewProto, - StreamFeatureView, - "feature_view_name", - "feature_view_proto", - FeatureViewNotFoundException, + table=stream_feature_views, + name=name, + project=project, + proto_class=StreamFeatureViewProto, + python_class=StreamFeatureView, + id_field_name="feature_view_name", + proto_field_name="feature_view_proto", + not_found_exception=FeatureViewNotFoundException, ) def list_stream_feature_views( @@ -230,101 +230,105 @@ def list_stream_feature_views( def apply_entity(self, entity: Entity, project: str, commit: bool = True): return self._apply_object( - entities, project, "entity_name", entity, "entity_proto" + table=entities, + project=project, + id_field_name="entity_name", + obj=entity, + proto_field_name="entity_proto", ) def get_entity(self, name: str, project: str, allow_cache: bool = False) -> Entity: return self._get_object( - entities, - name, - project, - EntityProto, - Entity, - "entity_name", - "entity_proto", - EntityNotFoundException, + table=entities, + name=name, + project=project, + proto_class=EntityProto, + python_class=Entity, + id_field_name="entity_name", + proto_field_name="entity_proto", + not_found_exception=EntityNotFoundException, ) def get_feature_view( self, name: str, project: str, allow_cache: bool = False ) -> FeatureView: return self._get_object( - feature_views, - name, - project, - FeatureViewProto, - FeatureView, - "feature_view_name", - "feature_view_proto", - FeatureViewNotFoundException, + table=feature_views, + name=name, + project=project, + proto_class=FeatureViewProto, + python_class=FeatureView, + id_field_name="feature_view_name", + proto_field_name="feature_view_proto", + not_found_exception=FeatureViewNotFoundException, ) def get_on_demand_feature_view( self, name: str, project: str, allow_cache: bool = False ) -> OnDemandFeatureView: return self._get_object( - on_demand_feature_views, - name, - project, - OnDemandFeatureViewProto, - OnDemandFeatureView, - "feature_view_name", - "feature_view_proto", - FeatureViewNotFoundException, + table=on_demand_feature_views, + name=name, + project=project, + proto_class=OnDemandFeatureViewProto, + python_class=OnDemandFeatureView, + id_field_name="feature_view_name", + proto_field_name="feature_view_proto", + not_found_exception=FeatureViewNotFoundException, ) def get_request_feature_view(self, name: str, project: str): return self._get_object( - request_feature_views, - name, - project, - RequestFeatureViewProto, - RequestFeatureView, - "feature_view_name", - "feature_view_proto", - FeatureViewNotFoundException, + table=request_feature_views, + name=name, + project=project, + proto_class=RequestFeatureViewProto, + python_class=RequestFeatureView, + id_field_name="feature_view_name", + proto_field_name="feature_view_proto", + not_found_exception=FeatureViewNotFoundException, ) def get_feature_service( self, name: str, project: str, allow_cache: bool = False ) -> FeatureService: return self._get_object( - feature_services, - name, - project, - FeatureServiceProto, - FeatureService, - "feature_service_name", - "feature_service_proto", - FeatureServiceNotFoundException, + table=feature_services, + name=name, + project=project, + proto_class=FeatureServiceProto, + python_class=FeatureService, + id_field_name="feature_service_name", + proto_field_name="feature_service_proto", + not_found_exception=FeatureServiceNotFoundException, ) def get_saved_dataset( self, name: str, project: str, allow_cache: bool = False ) -> SavedDataset: return self._get_object( - saved_datasets, - name, - project, - SavedDatasetProto, - SavedDataset, - "saved_dataset_name", - "saved_dataset_proto", - SavedDatasetNotFound, + table=saved_datasets, + name=name, + project=project, + proto_class=SavedDatasetProto, + python_class=SavedDataset, + id_field_name="saved_dataset_name", + proto_field_name="saved_dataset_proto", + not_found_exception=SavedDatasetNotFound, ) def get_validation_reference( self, name: str, project: str, allow_cache: bool = False ) -> ValidationReference: return self._get_object( - validation_references, - name, - project, - ValidationReferenceProto, - ValidationReference, - "validation_reference_name", - "validation_reference_proto", - ValidationReferenceNotFound, + table=validation_references, + name=name, + project=project, + proto_class=ValidationReferenceProto, + python_class=ValidationReference, + id_field_name="validation_reference_name", + proto_field_name="validation_reference_proto", + not_found_exception=ValidationReferenceNotFound, ) def list_entities(self, project: str, allow_cache: bool = False) -> List[Entity]: @@ -364,14 +368,14 @@ def get_data_source( self, name: str, project: str, allow_cache: bool = False ) -> DataSource: return self._get_object( - data_sources, - name, - project, - DataSourceProto, - DataSource, - "data_source_name", - "data_source_proto", - DataSourceObjectNotFoundException, + table=data_sources, + name=name, + project=project, + proto_class=DataSourceProto, + python_class=DataSource, + id_field_name="data_source_name", + proto_field_name="data_source_proto", + not_found_exception=DataSourceObjectNotFoundException, ) def list_data_sources( @@ -556,22 +560,28 @@ def delete_validation_reference(self, name: str, project: str, commit: bool = Tr def update_infra(self, infra: Infra, project: str, commit: bool = True): self._apply_object( - managed_infra, project, "infra_name", infra, "infra_proto", name="infra_obj" + table=managed_infra, + project=project, + id_field_name="infra_name", + obj=infra, + proto_field_name="infra_proto", + name="infra_obj", ) def get_infra(self, project: str, allow_cache: bool = False) -> Infra: infra_object = self._get_object( - managed_infra, - "infra_obj", - project, - InfraProto, - Infra, - "infra_name", - "infra_proto", - None, + table=managed_infra, + name="infra_obj", + project=project, + proto_class=InfraProto, + python_class=Infra, + id_field_name="infra_name", + proto_field_name="infra_proto", + not_found_exception=None, ) - infra_object = infra_object or InfraProto() - return Infra.from_proto(infra_object) + if infra_object: + return infra_object + return Infra() def apply_user_metadata( self, diff --git a/sdk/python/tests/unit/test_sql_registry.py b/sdk/python/tests/unit/test_sql_registry.py index 877811d0f32..51cb430c9eb 100644 --- a/sdk/python/tests/unit/test_sql_registry.py +++ b/sdk/python/tests/unit/test_sql_registry.py @@ -28,6 +28,8 @@ from feast.errors import FeatureViewNotFoundException from feast.feature_view import FeatureView from feast.field import Field +from feast.infra.infra_object import Infra +from feast.infra.online_stores.sqlite import SqliteTable from feast.infra.registry.sql import SqlRegistry from feast.on_demand_feature_view import on_demand_feature_view from feast.repo_config import RegistryConfig @@ -258,10 +260,20 @@ def test_apply_feature_view_success(sql_registry): and feature_view.features[3].dtype == Array(Bytes) and feature_view.entities[0] == "fs1_my_entity_1" ) + assert feature_view.ttl == timedelta(minutes=5) # After the first apply, the created_timestamp should be the same as the last_update_timestamp. assert feature_view.created_timestamp == feature_view.last_updated_timestamp + # Modify the feature view and apply again to test if diffing the online store table works + fv1.ttl = timedelta(minutes=6) + sql_registry.apply_feature_view(fv1, project) + feature_views = sql_registry.list_feature_views(project) + assert len(feature_views) == 1 + feature_view = sql_registry.get_feature_view("my_feature_view_1", project) + assert feature_view.ttl == timedelta(minutes=6) + + # Delete feature view sql_registry.delete_feature_view("my_feature_view_1", project) feature_views = sql_registry.list_feature_views(project) assert len(feature_views) == 0 @@ -570,6 +582,22 @@ def test_update_infra(sql_registry): project = "project" infra = sql_registry.get_infra(project=project) + assert len(infra.infra_objects) == 0 + # Should run update infra successfully sql_registry.update_infra(infra, project) + + # Should run update infra successfully when adding + new_infra = Infra() + new_infra.infra_objects.append( + SqliteTable( + path="/tmp/my_path.db", + name="my_table", + ) + ) + sql_registry.update_infra(new_infra, project) + infra = sql_registry.get_infra(project=project) + assert len(infra.infra_objects) == 1 + + # Try again since second time, infra should be not-empty sql_registry.teardown()