diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index 402cec758a2..2bcf939cadc 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -217,7 +217,6 @@ def apply( """ # TODO: Add locking - # TODO: Optimize by only making a single call (read/write) if isinstance(objects, Entity) or isinstance(objects, FeatureView): objects = [objects] @@ -242,9 +241,10 @@ def apply( raise ValueError("Unknown object type provided as part of apply() call") for view in views_to_update: - self._registry.apply_feature_view(view, project=self.project) + self._registry.apply_feature_view(view, project=self.project, commit=False) for ent in entities_to_update: - self._registry.apply_entity(ent, project=self.project) + self._registry.apply_entity(ent, project=self.project, commit=False) + self._registry.commit() self._get_provider().update_infra( project=self.project, diff --git a/sdk/python/feast/registry.py b/sdk/python/feast/registry.py index 53c3cae1e71..28601fbac9a 100644 --- a/sdk/python/feast/registry.py +++ b/sdk/python/feast/registry.py @@ -18,7 +18,7 @@ from datetime import datetime, timedelta from pathlib import Path from tempfile import TemporaryFile -from typing import Callable, List, Optional +from typing import List, Optional from urllib.parse import urlparse from feast.entity import Entity @@ -41,6 +41,9 @@ class Registry: Registry: A registry allows for the management and persistence of feature definitions and related metadata. """ + # 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(). cached_registry_proto: Optional[RegistryProto] = None cached_registry_proto_created: Optional[datetime] = None cached_registry_proto_ttl: timedelta @@ -73,34 +76,39 @@ def __init__(self, registry_path: str, repo_path: Path, cache_ttl: timedelta): return def _initialize_registry(self): - """Explicitly forces the initialization of a registry""" - self._registry_store.update_registry_proto(updater=None) + """Explicitly initializes the registry with an empty proto.""" + registry_proto = RegistryProto() + registry_proto.registry_schema_version = REGISTRY_SCHEMA_VERSION + self._registry_store.update_registry_proto(registry_proto) - def apply_entity(self, entity: Entity, project: str): + def apply_entity(self, entity: Entity, project: str, commit: bool = True): """ Registers a single entity with Feast Args: entity: Entity that will be registered project: Feast project that this entity belongs to + commit: Whether the change should be persisted immediately """ entity.is_valid() entity_proto = entity.to_proto() entity_proto.spec.project = project + self._prepare_registry_for_changes() + assert self.cached_registry_proto - def updater(registry_proto: RegistryProto): - for idx, existing_entity_proto in enumerate(registry_proto.entities): - if ( - existing_entity_proto.spec.name == entity_proto.spec.name - and existing_entity_proto.spec.project == project - ): - del registry_proto.entities[idx] - registry_proto.entities.append(entity_proto) - return registry_proto - registry_proto.entities.append(entity_proto) - return registry_proto + for idx, existing_entity_proto in enumerate( + self.cached_registry_proto.entities + ): + if ( + existing_entity_proto.spec.name == entity_proto.spec.name + and existing_entity_proto.spec.project == project + ): + del self.cached_registry_proto.entities[idx] + break - self._registry_store.update_registry_proto(updater) + self.cached_registry_proto.entities.append(entity_proto) + if commit: + self.commit() return def list_entities(self, project: str, allow_cache: bool = False) -> List[Entity]: @@ -139,71 +147,73 @@ def get_entity(self, name: str, project: str, allow_cache: bool = False) -> Enti return Entity.from_proto(entity_proto) raise EntityNotFoundException(name, project=project) - def apply_feature_table(self, feature_table: FeatureTable, project: str): + def apply_feature_table( + self, feature_table: FeatureTable, project: str, commit: bool = True + ): """ Registers a single feature table with Feast Args: feature_table: Feature table that will be registered project: Feast project that this feature table belongs to + commit: Whether the change should be persisted immediately """ feature_table.is_valid() feature_table_proto = feature_table.to_proto() feature_table_proto.spec.project = project + self._prepare_registry_for_changes() + assert self.cached_registry_proto - def updater(registry_proto: RegistryProto): - for idx, existing_feature_table_proto in enumerate( - registry_proto.feature_tables + for idx, existing_feature_table_proto in enumerate( + self.cached_registry_proto.feature_tables + ): + if ( + existing_feature_table_proto.spec.name == feature_table_proto.spec.name + and existing_feature_table_proto.spec.project == project ): - if ( - existing_feature_table_proto.spec.name - == feature_table_proto.spec.name - and existing_feature_table_proto.spec.project == project - ): - del registry_proto.feature_tables[idx] - registry_proto.feature_tables.append(feature_table_proto) - return registry_proto - registry_proto.feature_tables.append(feature_table_proto) - return registry_proto + del self.cached_registry_proto.feature_tables[idx] + break - self._registry_store.update_registry_proto(updater) + self.cached_registry_proto.feature_tables.append(feature_table_proto) + if commit: + self.commit() return - def apply_feature_view(self, feature_view: FeatureView, project: str): + def apply_feature_view( + self, feature_view: FeatureView, project: str, commit: bool = True + ): """ Registers a single feature view with Feast Args: feature_view: Feature view that will be registered project: Feast project that this feature view belongs to + commit: Whether the change should be persisted immediately """ feature_view.is_valid() feature_view_proto = feature_view.to_proto() feature_view_proto.spec.project = project + self._prepare_registry_for_changes() + assert self.cached_registry_proto - def updater(registry_proto: RegistryProto): - for idx, existing_feature_view_proto in enumerate( - registry_proto.feature_views + for idx, existing_feature_view_proto in enumerate( + self.cached_registry_proto.feature_views + ): + if ( + existing_feature_view_proto.spec.name == feature_view_proto.spec.name + and existing_feature_view_proto.spec.project == project ): - if ( - existing_feature_view_proto.spec.name - == feature_view_proto.spec.name - and existing_feature_view_proto.spec.project == project - ): - # do not update if feature view has not changed; updating will erase tracked materialization intervals - if ( - FeatureView.from_proto(existing_feature_view_proto) - == feature_view - ): - return registry_proto - else: - del registry_proto.feature_views[idx] - registry_proto.feature_views.append(feature_view_proto) - return registry_proto - registry_proto.feature_views.append(feature_view_proto) - return registry_proto - - self._registry_store.update_registry_proto(updater) + # do not update if feature view has not changed; updating will erase tracked materialization intervals + if FeatureView.from_proto(existing_feature_view_proto) == feature_view: + return + else: + del self.cached_registry_proto.feature_views[idx] + break + + self.cached_registry_proto.feature_views.append(feature_view_proto) + if commit: + self.commit() + return def apply_materialization( self, @@ -211,6 +221,7 @@ def apply_materialization( project: str, start_date: datetime, end_date: datetime, + commit: bool = True, ): """ Updates materialization intervals tracked for a single feature view in Feast @@ -220,30 +231,33 @@ def apply_materialization( project: Feast project that this feature view belongs to start_date (datetime): Start date of the materialization interval to track end_date (datetime): End date of the materialization interval to track + commit: Whether the change should be persisted immediately """ + self._prepare_registry_for_changes() + assert self.cached_registry_proto - def updater(registry_proto: RegistryProto): - for idx, existing_feature_view_proto in enumerate( - registry_proto.feature_views + for idx, existing_feature_view_proto in enumerate( + self.cached_registry_proto.feature_views + ): + if ( + existing_feature_view_proto.spec.name == feature_view.name + and existing_feature_view_proto.spec.project == project ): - if ( - existing_feature_view_proto.spec.name == feature_view.name - and existing_feature_view_proto.spec.project == project - ): - existing_feature_view = FeatureView.from_proto( - existing_feature_view_proto - ) - existing_feature_view.materialization_intervals.append( - (start_date, end_date) - ) - feature_view_proto = existing_feature_view.to_proto() - feature_view_proto.spec.project = project - del registry_proto.feature_views[idx] - registry_proto.feature_views.append(feature_view_proto) - return registry_proto - raise FeatureViewNotFoundException(feature_view.name, project) - - self._registry_store.update_registry_proto(updater) + existing_feature_view = FeatureView.from_proto( + existing_feature_view_proto + ) + existing_feature_view.materialization_intervals.append( + (start_date, end_date) + ) + feature_view_proto = existing_feature_view.to_proto() + feature_view_proto.spec.project = project + del self.cached_registry_proto.feature_views[idx] + self.cached_registry_proto.feature_views.append(feature_view_proto) + if commit: + self.commit() + return + + raise FeatureViewNotFoundException(feature_view.name, project) def list_feature_tables(self, project: str) -> List[FeatureTable]: """ @@ -324,57 +338,78 @@ def get_feature_view(self, name: str, project: str) -> FeatureView: return FeatureView.from_proto(feature_view_proto) raise FeatureViewNotFoundException(name, project) - def delete_feature_table(self, name: str, project: str): + def delete_feature_table(self, name: str, project: str, commit: bool = True): """ Deletes a feature table or raises an exception if not found. Args: name: Name of feature table project: Feast project that this feature table belongs to + commit: Whether the change should be persisted immediately """ + self._prepare_registry_for_changes() + assert self.cached_registry_proto - def updater(registry_proto: RegistryProto): - for idx, existing_feature_table_proto in enumerate( - registry_proto.feature_tables + for idx, existing_feature_table_proto in enumerate( + self.cached_registry_proto.feature_tables + ): + if ( + existing_feature_table_proto.spec.name == name + and existing_feature_table_proto.spec.project == project ): - if ( - existing_feature_table_proto.spec.name == name - and existing_feature_table_proto.spec.project == project - ): - del registry_proto.feature_tables[idx] - return registry_proto - raise FeatureTableNotFoundException(name, project) - - self._registry_store.update_registry_proto(updater) - return + del self.cached_registry_proto.feature_tables[idx] + if commit: + self.commit() + return - def delete_feature_view(self, name: str, project: str): + raise FeatureTableNotFoundException(name, project) + + def delete_feature_view(self, name: str, project: str, commit: bool = True): """ Deletes a feature view or raises an exception if not found. Args: name: Name of feature view project: Feast project that this feature view belongs to + commit: Whether the change should be persisted immediately """ + self._prepare_registry_for_changes() + assert self.cached_registry_proto - def updater(registry_proto: RegistryProto): - for idx, existing_feature_view_proto in enumerate( - registry_proto.feature_views + for idx, existing_feature_view_proto in enumerate( + self.cached_registry_proto.feature_views + ): + if ( + existing_feature_view_proto.spec.name == name + and existing_feature_view_proto.spec.project == project ): - if ( - existing_feature_view_proto.spec.name == name - and existing_feature_view_proto.spec.project == project - ): - del registry_proto.feature_views[idx] - return registry_proto - raise FeatureViewNotFoundException(name, project) + del self.cached_registry_proto.feature_views[idx] + if commit: + self.commit() + return - self._registry_store.update_registry_proto(updater) + raise FeatureViewNotFoundException(name, project) + + def commit(self): + """Commits the state of the registry cache to the remote registry store.""" + if self.cached_registry_proto: + self._registry_store.update_registry_proto(self.cached_registry_proto) def refresh(self): """Refreshes the state of the registry cache by fetching the registry state from the remote registry store.""" self._get_registry_proto(allow_cache=False) + def _prepare_registry_for_changes(self): + """Prepares the Registry for changes by refreshing the cache if necessary.""" + try: + self._get_registry_proto(allow_cache=True) + except FileNotFoundError: + registry_proto = RegistryProto() + registry_proto.registry_schema_version = REGISTRY_SCHEMA_VERSION + self.cached_registry_proto = registry_proto + self.cached_registry_proto_created = datetime.now() + return self.cached_registry_proto + def _get_registry_proto(self, allow_cache: bool = False) -> RegistryProto: """Returns the cached or remote registry state @@ -419,7 +454,7 @@ class RegistryStore(ABC): def get_registry_proto(self): """ Retrieves the registry proto from the registry path. If there is no file at that path, - returns an empty registry proto. + raises a FileNotFoundError. Returns: Returns either the registry proto stored at the registry path, or an empty registry proto. @@ -427,15 +462,13 @@ def get_registry_proto(self): pass @abstractmethod - def update_registry_proto( - self, updater: Optional[Callable[[RegistryProto], RegistryProto]] = None - ): + def update_registry_proto(self, registry_proto: RegistryProto): """ - Updates the registry using the function passed in. If the registry proto has not been created yet - this method will create it. This method writes to the registry path. + Overwrites the current registry proto with the proto passed in. This method + writes to the registry path. Args: - updater: function that takes in the current registry proto and outputs the desired registry proto + registry_proto: the new RegistryProto """ pass @@ -457,17 +490,7 @@ def get_registry_proto(self): f'Registry not found at path "{self._filepath}". Have you run "feast apply"?' ) - def update_registry_proto( - self, updater: Optional[Callable[[RegistryProto], RegistryProto]] = None - ): - try: - registry_proto = self.get_registry_proto() - except FileNotFoundError: - registry_proto = RegistryProto() - registry_proto.registry_schema_version = REGISTRY_SCHEMA_VERSION - - if updater: - registry_proto = updater(registry_proto) + def update_registry_proto(self, registry_proto: RegistryProto): self._write_registry(registry_proto) return @@ -518,16 +541,7 @@ def get_registry_proto(self): f'Registry not found at path "{self._uri.geturl()}". Have you run "feast apply"?' ) - def update_registry_proto( - self, updater: Optional[Callable[[RegistryProto], RegistryProto]] = None - ): - try: - registry_proto = self.get_registry_proto() - except FileNotFoundError: - registry_proto = RegistryProto() - registry_proto.registry_schema_version = REGISTRY_SCHEMA_VERSION - if updater: - registry_proto = updater(registry_proto) + def update_registry_proto(self, registry_proto: RegistryProto): self._write_registry(registry_proto) return @@ -592,17 +606,9 @@ def get_registry_proto(self): f"Error while trying to locate Registry at path {self._uri.geturl()}" ) from e - def update_registry_proto( - self, updater: Optional[Callable[[RegistryProto], RegistryProto]] = None - ): - try: - registry_proto = self.get_registry_proto() - except FileNotFoundError: - registry_proto = RegistryProto() - registry_proto.registry_schema_version = REGISTRY_SCHEMA_VERSION - if updater: - registry_proto = updater(registry_proto) + def update_registry_proto(self, registry_proto: RegistryProto): self._write_registry(registry_proto) + return def _write_registry(self, registry_proto: RegistryProto): registry_proto.version_id = str(uuid.uuid4()) diff --git a/sdk/python/feast/repo_operations.py b/sdk/python/feast/repo_operations.py index 91dbba0824c..6482451f714 100644 --- a/sdk/python/feast/repo_operations.py +++ b/sdk/python/feast/repo_operations.py @@ -146,13 +146,6 @@ def apply_total(repo_config: RepoConfig, repo_path: Path, skip_source_validation for view in repo.feature_views: view.infer_features_from_input_source(repo_config) - sys.dont_write_bytecode = False - for entity in repo.entities: - registry.apply_entity(entity, project=project) - click.echo( - f"Registered entity {Style.BRIGHT + Fore.GREEN}{entity.name}{Style.RESET_ALL}" - ) - repo_table_names = set(t.name for t in repo.feature_tables) for t in repo.feature_views: @@ -168,33 +161,43 @@ def apply_total(repo_config: RepoConfig, repo_path: Path, skip_source_validation if registry_view.name not in repo_table_names: views_to_delete.append(registry_view) + sys.dont_write_bytecode = False + for entity in repo.entities: + registry.apply_entity(entity, project=project, commit=False) + click.echo( + f"Registered entity {Style.BRIGHT + Fore.GREEN}{entity.name}{Style.RESET_ALL}" + ) + # Delete tables that should not exist for registry_table in tables_to_delete: - registry.delete_feature_table(registry_table.name, project=project) + registry.delete_feature_table( + registry_table.name, project=project, commit=False + ) click.echo( f"Deleted feature table {Style.BRIGHT + Fore.GREEN}{registry_table.name}{Style.RESET_ALL} from registry" ) # Create tables that should for table in repo.feature_tables: - registry.apply_feature_table(table, project) + registry.apply_feature_table(table, project, commit=False) click.echo( f"Registered feature table {Style.BRIGHT + Fore.GREEN}{table.name}{Style.RESET_ALL}" ) # Delete views that should not exist for registry_view in views_to_delete: - registry.delete_feature_view(registry_view.name, project=project) + registry.delete_feature_view(registry_view.name, project=project, commit=False) click.echo( f"Deleted feature view {Style.BRIGHT + Fore.GREEN}{registry_view.name}{Style.RESET_ALL} from registry" ) # Create views that should for view in repo.feature_views: - registry.apply_feature_view(view, project) + registry.apply_feature_view(view, project, commit=False) click.echo( f"Registered feature view {Style.BRIGHT + Fore.GREEN}{view.name}{Style.RESET_ALL}" ) + registry.commit() infra_provider = get_provider(repo_config, repo_path) diff --git a/sdk/python/tests/test_registry.py b/sdk/python/tests/test_registry.py new file mode 100644 index 00000000000..5ed5579b3ce --- /dev/null +++ b/sdk/python/tests/test_registry.py @@ -0,0 +1,347 @@ +# Copyright 2021 The Feast Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import time +from datetime import timedelta +from tempfile import mkstemp + +import pytest +from pytest_lazyfixture import lazy_fixture + +from feast.data_format import ParquetFormat +from feast.data_source import FileSource +from feast.entity import Entity +from feast.feature import Feature +from feast.feature_view import FeatureView +from feast.protos.feast.types import Value_pb2 as ValueProto +from feast.registry import Registry +from feast.value_type import ValueType + + +@pytest.fixture +def local_registry(): + fd, registry_path = mkstemp() + return Registry(registry_path, None, timedelta(600)) + + +@pytest.fixture +def gcs_registry(): + from google.cloud import storage + + storage_client = storage.Client() + bucket_name = f"feast-registry-test-{int(time.time() * 1000)}" + bucket = storage_client.bucket(bucket_name) + bucket = storage_client.create_bucket(bucket) + bucket.add_lifecycle_delete_rule( + age=14 + ) # delete buckets automatically after 14 days + bucket.patch() + bucket.blob("registry.db") + return Registry(f"gs://{bucket_name}/registry.db", None, timedelta(600)) + + +@pytest.fixture +def s3_registry(): + return Registry( + f"s3://feast-integration-tests/registries/{int(time.time() * 1000)}/registry.db", + None, + timedelta(600), + ) + + +@pytest.mark.parametrize( + "test_registry", [lazy_fixture("local_registry")], +) +def test_apply_entity_success(test_registry): + entity = Entity( + name="driver_car_id", + description="Car driver id", + value_type=ValueType.STRING, + labels={"team": "matchmaking"}, + ) + + project = "project" + + # Register Entity + test_registry.apply_entity(entity, project) + + entities = test_registry.list_entities(project) + + entity = entities[0] + assert ( + len(entities) == 1 + and entity.name == "driver_car_id" + and entity.value_type == ValueType(ValueProto.ValueType.STRING) + and entity.description == "Car driver id" + and "team" in entity.labels + and entity.labels["team"] == "matchmaking" + ) + + entity = test_registry.get_entity("driver_car_id", project) + assert ( + entity.name == "driver_car_id" + and entity.value_type == ValueType(ValueProto.ValueType.STRING) + and entity.description == "Car driver id" + and "team" in entity.labels + and entity.labels["team"] == "matchmaking" + ) + + +@pytest.mark.integration +@pytest.mark.parametrize( + "test_registry", [lazy_fixture("gcs_registry"), lazy_fixture("s3_registry")], +) +def test_apply_entity_integration(test_registry): + entity = Entity( + name="driver_car_id", + description="Car driver id", + value_type=ValueType.STRING, + labels={"team": "matchmaking"}, + ) + + project = "project" + + # Register Entity + test_registry.apply_entity(entity, project) + + entities = test_registry.list_entities(project) + + entity = entities[0] + assert ( + len(entities) == 1 + and entity.name == "driver_car_id" + and entity.value_type == ValueType(ValueProto.ValueType.STRING) + and entity.description == "Car driver id" + and "team" in entity.labels + and entity.labels["team"] == "matchmaking" + ) + + entity = test_registry.get_entity("driver_car_id", project) + assert ( + entity.name == "driver_car_id" + and entity.value_type == ValueType(ValueProto.ValueType.STRING) + and entity.description == "Car driver id" + and "team" in entity.labels + and entity.labels["team"] == "matchmaking" + ) + + +@pytest.mark.parametrize( + "test_registry", [lazy_fixture("local_registry")], +) +def test_apply_feature_view_success(test_registry): + # Create Feature Views + batch_source = FileSource( + file_format=ParquetFormat(), + file_url="file://feast/*", + event_timestamp_column="ts_col", + created_timestamp_column="timestamp", + date_partition_column="date_partition_col", + ) + + fv1 = FeatureView( + name="my_feature_view_1", + features=[ + Feature(name="fs1_my_feature_1", dtype=ValueType.INT64), + Feature(name="fs1_my_feature_2", dtype=ValueType.STRING), + Feature(name="fs1_my_feature_3", dtype=ValueType.STRING_LIST), + Feature(name="fs1_my_feature_4", dtype=ValueType.BYTES_LIST), + ], + entities=["fs1_my_entity_1"], + tags={"team": "matchmaking"}, + input=batch_source, + ttl=timedelta(minutes=5), + ) + + project = "project" + + # Register Feature View + test_registry.apply_feature_view(fv1, project) + + feature_views = test_registry.list_feature_views(project) + + # List Feature Views + assert ( + len(feature_views) == 1 + and feature_views[0].name == "my_feature_view_1" + and feature_views[0].features[0].name == "fs1_my_feature_1" + and feature_views[0].features[0].dtype == ValueType.INT64 + and feature_views[0].features[1].name == "fs1_my_feature_2" + and feature_views[0].features[1].dtype == ValueType.STRING + and feature_views[0].features[2].name == "fs1_my_feature_3" + and feature_views[0].features[2].dtype == ValueType.STRING_LIST + and feature_views[0].features[3].name == "fs1_my_feature_4" + and feature_views[0].features[3].dtype == ValueType.BYTES_LIST + and feature_views[0].entities[0] == "fs1_my_entity_1" + ) + + feature_view = test_registry.get_feature_view("my_feature_view_1", project) + assert ( + feature_view.name == "my_feature_view_1" + and feature_view.features[0].name == "fs1_my_feature_1" + and feature_view.features[0].dtype == ValueType.INT64 + and feature_view.features[1].name == "fs1_my_feature_2" + and feature_view.features[1].dtype == ValueType.STRING + and feature_view.features[2].name == "fs1_my_feature_3" + and feature_view.features[2].dtype == ValueType.STRING_LIST + and feature_view.features[3].name == "fs1_my_feature_4" + and feature_view.features[3].dtype == ValueType.BYTES_LIST + and feature_view.entities[0] == "fs1_my_entity_1" + ) + + test_registry.delete_feature_view("my_feature_view_1", project) + feature_views = test_registry.list_feature_views(project) + assert len(feature_views) == 0 + + +@pytest.mark.integration +@pytest.mark.parametrize( + "test_registry", [lazy_fixture("gcs_registry"), lazy_fixture("s3_registry")], +) +def test_apply_feature_view_integration(test_registry): + # Create Feature Views + batch_source = FileSource( + file_format=ParquetFormat(), + file_url="file://feast/*", + event_timestamp_column="ts_col", + created_timestamp_column="timestamp", + date_partition_column="date_partition_col", + ) + + fv1 = FeatureView( + name="my_feature_view_1", + features=[ + Feature(name="fs1_my_feature_1", dtype=ValueType.INT64), + Feature(name="fs1_my_feature_2", dtype=ValueType.STRING), + Feature(name="fs1_my_feature_3", dtype=ValueType.STRING_LIST), + Feature(name="fs1_my_feature_4", dtype=ValueType.BYTES_LIST), + ], + entities=["fs1_my_entity_1"], + tags={"team": "matchmaking"}, + input=batch_source, + ttl=timedelta(minutes=5), + ) + + project = "project" + + # Register Feature View + test_registry.apply_feature_view(fv1, project) + + feature_views = test_registry.list_feature_views(project) + + # List Feature Views + assert ( + len(feature_views) == 1 + and feature_views[0].name == "my_feature_view_1" + and feature_views[0].features[0].name == "fs1_my_feature_1" + and feature_views[0].features[0].dtype == ValueType.INT64 + and feature_views[0].features[1].name == "fs1_my_feature_2" + and feature_views[0].features[1].dtype == ValueType.STRING + and feature_views[0].features[2].name == "fs1_my_feature_3" + and feature_views[0].features[2].dtype == ValueType.STRING_LIST + and feature_views[0].features[3].name == "fs1_my_feature_4" + and feature_views[0].features[3].dtype == ValueType.BYTES_LIST + and feature_views[0].entities[0] == "fs1_my_entity_1" + ) + + feature_view = test_registry.get_feature_view("my_feature_view_1", project) + assert ( + feature_view.name == "my_feature_view_1" + and feature_view.features[0].name == "fs1_my_feature_1" + and feature_view.features[0].dtype == ValueType.INT64 + and feature_view.features[1].name == "fs1_my_feature_2" + and feature_view.features[1].dtype == ValueType.STRING + and feature_view.features[2].name == "fs1_my_feature_3" + and feature_view.features[2].dtype == ValueType.STRING_LIST + and feature_view.features[3].name == "fs1_my_feature_4" + and feature_view.features[3].dtype == ValueType.BYTES_LIST + and feature_view.entities[0] == "fs1_my_entity_1" + ) + + test_registry.delete_feature_view("my_feature_view_1", project) + feature_views = test_registry.list_feature_views(project) + assert len(feature_views) == 0 + + +def test_commit(): + fd, registry_path = mkstemp() + test_registry = Registry(registry_path, None, timedelta(600)) + + entity = Entity( + name="driver_car_id", + description="Car driver id", + value_type=ValueType.STRING, + labels={"team": "matchmaking"}, + ) + + project = "project" + + # Register Entity without commiting + test_registry.apply_entity(entity, project, commit=False) + + # Retrieving the entity should still succeed + entities = test_registry.list_entities(project, allow_cache=True) + + entity = entities[0] + assert ( + len(entities) == 1 + and entity.name == "driver_car_id" + and entity.value_type == ValueType(ValueProto.ValueType.STRING) + and entity.description == "Car driver id" + and "team" in entity.labels + and entity.labels["team"] == "matchmaking" + ) + + entity = test_registry.get_entity("driver_car_id", project, allow_cache=True) + assert ( + entity.name == "driver_car_id" + and entity.value_type == ValueType(ValueProto.ValueType.STRING) + and entity.description == "Car driver id" + and "team" in entity.labels + and entity.labels["team"] == "matchmaking" + ) + + # Create new registry that points to the same store + registry_with_same_store = Registry(registry_path, None, timedelta(600)) + + # Retrieving the entity should fail since the store is empty + entities = registry_with_same_store.list_entities(project) + assert len(entities) == 0 + + # commit from the original registry + test_registry.commit() + + # Reconstruct the new registry in order to read the newly written store + registry_with_same_store = Registry(registry_path, None, timedelta(600)) + + # Retrieving the entity should now succeed + entities = registry_with_same_store.list_entities(project) + + entity = entities[0] + assert ( + len(entities) == 1 + and entity.name == "driver_car_id" + and entity.value_type == ValueType(ValueProto.ValueType.STRING) + and entity.description == "Car driver id" + and "team" in entity.labels + and entity.labels["team"] == "matchmaking" + ) + + entity = test_registry.get_entity("driver_car_id", project) + assert ( + entity.name == "driver_car_id" + and entity.value_type == ValueType(ValueProto.ValueType.STRING) + and entity.description == "Car driver id" + and "team" in entity.labels + and entity.labels["team"] == "matchmaking" + )