From 4da806e6594f55ed373a740668e6dbdbed4191b1 Mon Sep 17 00:00:00 2001 From: Felix Wang Date: Tue, 25 Jan 2022 13:56:33 -0800 Subject: [PATCH] Switch from hardcoded names to enum Signed-off-by: Felix Wang --- sdk/python/feast/diff/FcoDiff.py | 80 ++++++++++++++++++++--------- sdk/python/feast/registry.py | 16 +++--- sdk/python/feast/repo_operations.py | 17 +++--- 3 files changed, 74 insertions(+), 39 deletions(-) diff --git a/sdk/python/feast/diff/FcoDiff.py b/sdk/python/feast/diff/FcoDiff.py index f1cbf55c41..e2aac16bf5 100644 --- a/sdk/python/feast/diff/FcoDiff.py +++ b/sdk/python/feast/diff/FcoDiff.py @@ -16,9 +16,19 @@ from feast.protos.feast.core.RequestFeatureView_pb2 import ( RequestFeatureView as RequestFeatureViewProto, ) -from feast.registry import REGISTRY_OBJECT_TYPE_TO_STR, REGISTRY_OBJECT_TYPES, Registry +from feast.registry import FeastObjectType, Registry from feast.repo_contents import RepoContents +FEAST_OBJECT_TYPE_TO_STR = { + FeastObjectType.ENTITY: "entity", + FeastObjectType.FEATURE_VIEW: "feature view", + FeastObjectType.ON_DEMAND_FEATURE_VIEW: "on demand feature view", + FeastObjectType.REQUEST_FEATURE_VIEW: "request feature view", + FeastObjectType.FEATURE_SERVICE: "feature service", +} + +FEAST_OBJECT_TYPES = FEAST_OBJECT_TYPE_TO_STR.keys() + Fco = TypeVar("Fco", Entity, BaseFeatureView, FeatureService) @@ -105,14 +115,22 @@ def diff_registry_objects(current: Fco, new: Fco, object_type: str) -> FcoDiff: ) ) return FcoDiff( - new_proto.spec.name, object_type, current, new, property_diffs, transition, + name=new_proto.spec.name, + fco_type=object_type, + current_fco=current, + new_fco=new, + fco_property_diffs=property_diffs, + transition_type=transition, ) def extract_objects_for_keep_delete_update_add( registry: Registry, current_project: str, desired_repo_contents: RepoContents, ) -> Tuple[ - Dict[str, Set[Fco]], Dict[str, Set[Fco]], Dict[str, Set[Fco]], Dict[str, Set[Fco]] + Dict[FeastObjectType, Set[Fco]], + Dict[FeastObjectType, Set[Fco]], + Dict[FeastObjectType, Set[Fco]], + Dict[FeastObjectType, Set[Fco]], ]: """ Returns the objects in the registry that must be modified to achieve the desired repo state. @@ -127,20 +145,32 @@ def extract_objects_for_keep_delete_update_add( objs_to_update = {} objs_to_add = {} - registry_object_type_to_objects: Dict[str, List[Any]] + registry_object_type_to_objects: Dict[FeastObjectType, List[Any]] registry_object_type_to_objects = { - "entities": registry.list_entities(project=current_project), - "feature_views": registry.list_feature_views(project=current_project), - "on_demand_feature_views": registry.list_on_demand_feature_views( + FeastObjectType.ENTITY: registry.list_entities(project=current_project), + FeastObjectType.FEATURE_VIEW: registry.list_feature_views( + project=current_project + ), + FeastObjectType.ON_DEMAND_FEATURE_VIEW: registry.list_on_demand_feature_views( project=current_project ), - "request_feature_views": registry.list_request_feature_views( + FeastObjectType.REQUEST_FEATURE_VIEW: registry.list_request_feature_views( project=current_project ), - "feature_services": registry.list_feature_services(project=current_project), + FeastObjectType.FEATURE_SERVICE: registry.list_feature_services( + project=current_project + ), + } + registry_object_type_to_repo_contents: Dict[FeastObjectType, Set[Any]] + registry_object_type_to_repo_contents = { + FeastObjectType.ENTITY: desired_repo_contents.entities, + FeastObjectType.FEATURE_VIEW: desired_repo_contents.feature_views, + FeastObjectType.ON_DEMAND_FEATURE_VIEW: desired_repo_contents.on_demand_feature_views, + FeastObjectType.REQUEST_FEATURE_VIEW: desired_repo_contents.request_feature_views, + FeastObjectType.FEATURE_SERVICE: desired_repo_contents.feature_services, } - for object_type in REGISTRY_OBJECT_TYPES: + for object_type in FEAST_OBJECT_TYPES: ( to_keep, to_delete, @@ -148,7 +178,7 @@ def extract_objects_for_keep_delete_update_add( to_add, ) = tag_objects_for_keep_delete_update_add( registry_object_type_to_objects[object_type], - getattr(desired_repo_contents, object_type), + registry_object_type_to_repo_contents[object_type], ) objs_to_keep[object_type] = to_keep @@ -181,7 +211,7 @@ def diff_between( registry, current_project, desired_repo_contents ) - for object_type in REGISTRY_OBJECT_TYPES: + for object_type in FEAST_OBJECT_TYPES: objects_to_keep = objs_to_keep[object_type] objects_to_delete = objs_to_delete[object_type] objects_to_update = objs_to_update[object_type] @@ -190,30 +220,30 @@ def diff_between( for e in objects_to_add: diff.add_fco_diff( FcoDiff( - e.name, - REGISTRY_OBJECT_TYPE_TO_STR[object_type], - None, - e, - [], - TransitionType.CREATE, + name=e.name, + fco_type=FEAST_OBJECT_TYPE_TO_STR[object_type], + current_fco=None, + new_fco=e, + fco_property_diffs=[], + transition_type=TransitionType.CREATE, ) ) for e in objects_to_delete: diff.add_fco_diff( FcoDiff( - e.name, - REGISTRY_OBJECT_TYPE_TO_STR[object_type], - e, - None, - [], - TransitionType.DELETE, + name=e.name, + fco_type=FEAST_OBJECT_TYPE_TO_STR[object_type], + current_fco=e, + new_fco=None, + fco_property_diffs=[], + transition_type=TransitionType.DELETE, ) ) for e in objects_to_update: current_obj = [_e for _e in objects_to_keep if _e.name == e.name][0] diff.add_fco_diff( diff_registry_objects( - current_obj, e, REGISTRY_OBJECT_TYPE_TO_STR[object_type] + current_obj, e, FEAST_OBJECT_TYPE_TO_STR[object_type] ) ) diff --git a/sdk/python/feast/registry.py b/sdk/python/feast/registry.py index 3254990adc..f05abc6d9a 100644 --- a/sdk/python/feast/registry.py +++ b/sdk/python/feast/registry.py @@ -14,6 +14,7 @@ import logging from collections import defaultdict from datetime import datetime, timedelta +from enum import Enum from pathlib import Path from threading import Lock from typing import Any, Dict, List, Optional @@ -58,15 +59,14 @@ "": "LocalRegistryStore", } -REGISTRY_OBJECT_TYPE_TO_STR = { - "entities": "entity", - "feature_views": "feature view", - "on_demand_feature_views": "on demand feature view", - "request_feature_views": "request feature view", - "feature_services": "feature service", -} -REGISTRY_OBJECT_TYPES = REGISTRY_OBJECT_TYPE_TO_STR.keys() +class FeastObjectType(Enum): + ENTITY = 0 + FEATURE_VIEW = 1 + ON_DEMAND_FEATURE_VIEW = 2 + REQUEST_FEATURE_VIEW = 3 + FEATURE_SERVICE = 4 + logger = logging.getLogger(__name__) diff --git a/sdk/python/feast/repo_operations.py b/sdk/python/feast/repo_operations.py index 99004eac36..17d0530d4e 100644 --- a/sdk/python/feast/repo_operations.py +++ b/sdk/python/feast/repo_operations.py @@ -12,16 +12,17 @@ from click.exceptions import BadParameter from feast.diff.FcoDiff import ( - TransitionType, + FEAST_OBJECT_TYPES, extract_objects_for_keep_delete_update_add, ) +from feast.diff.property_diff import TransitionType from feast.entity import Entity from feast.feature_service import FeatureService from feast.feature_store import FeatureStore from feast.feature_view import DUMMY_ENTITY, DUMMY_ENTITY_NAME, FeatureView from feast.names import adjectives, animals from feast.on_demand_feature_view import OnDemandFeatureView -from feast.registry import REGISTRY_OBJECT_TYPES, Registry +from feast.registry import FeastObjectType, Registry from feast.repo_config import RepoConfig from feast.repo_contents import RepoContents from feast.request_feature_view import RequestFeatureView @@ -176,7 +177,7 @@ def extract_objects_for_apply_delete(project, registry, repo): Entity, FeatureView, RequestFeatureView, OnDemandFeatureView, FeatureService ] ] = [] - for object_type in REGISTRY_OBJECT_TYPES: + for object_type in FEAST_OBJECT_TYPES: to_apply = set(objs_to_add[object_type]).union(objs_to_update[object_type]) all_to_apply.extend(to_apply) @@ -185,14 +186,18 @@ def extract_objects_for_apply_delete(project, registry, repo): Entity, FeatureView, RequestFeatureView, OnDemandFeatureView, FeatureService ] ] = [] - for object_type in REGISTRY_OBJECT_TYPES: + for object_type in FEAST_OBJECT_TYPES: all_to_delete.extend(objs_to_delete[object_type]) return ( all_to_apply, all_to_delete, - set(objs_to_add["feature_views"].union(objs_to_update["feature_views"])), - objs_to_delete["feature_views"], + set( + objs_to_add[FeastObjectType.FEATURE_VIEW].union( + objs_to_update[FeastObjectType.FEATURE_VIEW] + ) + ), + objs_to_delete[FeastObjectType.FEATURE_VIEW], )