diff --git a/sdk/python/feast/feature_server.py b/sdk/python/feast/feature_server.py index 1f4918fe7a..e22ff43b9c 100644 --- a/sdk/python/feast/feature_server.py +++ b/sdk/python/feast/feature_server.py @@ -76,6 +76,38 @@ class GetOnlineFeaturesRequest(BaseModel): full_feature_names: bool = False +def _get_features(request: GetOnlineFeaturesRequest, store: "feast.FeatureStore"): + if request.feature_service: + feature_service = store.get_feature_service( + request.feature_service, allow_cache=True + ) + assert_permissions( + resource=feature_service, actions=[AuthzedAction.READ_ONLINE] + ) + features = feature_service # type: ignore + else: + all_feature_views, all_on_demand_feature_views = ( + utils._get_feature_views_to_use( + store.registry, + store.project, + request.features, + allow_cache=True, + hide_dummy_entity=False, + ) + ) + for feature_view in all_feature_views: + assert_permissions( + resource=feature_view, actions=[AuthzedAction.READ_ONLINE] + ) + for od_feature_view in all_on_demand_feature_views: + assert_permissions( + resource=od_feature_view, actions=[AuthzedAction.READ_ONLINE] + ) + features = request.features # type: ignore + + return features + + def get_app( store: "feast.FeatureStore", registry_ttl_sec: int = DEFAULT_FEATURE_SERVER_REGISTRY_TTL, @@ -121,33 +153,7 @@ async def lifespan(app: FastAPI): ) async def get_online_features(request: GetOnlineFeaturesRequest) -> Dict[str, Any]: # Initialize parameters for FeatureStore.get_online_features(...) call - if request.feature_service: - feature_service = store.get_feature_service( - request.feature_service, allow_cache=True - ) - assert_permissions( - resource=feature_service, actions=[AuthzedAction.READ_ONLINE] - ) - features = feature_service # type: ignore - else: - all_feature_views, all_on_demand_feature_views = ( - utils._get_feature_views_to_use( - store.registry, - store.project, - request.features, - allow_cache=True, - hide_dummy_entity=False, - ) - ) - for feature_view in all_feature_views: - assert_permissions( - resource=feature_view, actions=[AuthzedAction.READ_ONLINE] - ) - for od_feature_view in all_on_demand_feature_views: - assert_permissions( - resource=od_feature_view, actions=[AuthzedAction.READ_ONLINE] - ) - features = request.features # type: ignore + features = await run_in_threadpool(_get_features, request, store) read_params = dict( features=features, @@ -163,9 +169,13 @@ async def get_online_features(request: GetOnlineFeaturesRequest) -> Dict[str, An ) # Convert the Protobuf object to JSON and return it - return MessageToDict( - response.proto, preserving_proto_field_name=True, float_precision=18 + response_dict = await run_in_threadpool( + MessageToDict, + response.proto, + preserving_proto_field_name=True, + float_precision=18, ) + return response_dict @app.post("/push", dependencies=[Depends(inject_user_details)]) async def push(request: PushFeaturesRequest) -> None: