diff --git a/sdk/python/feast/feature_server.py b/sdk/python/feast/feature_server.py index cd36a90943..c5abafc23c 100644 --- a/sdk/python/feast/feature_server.py +++ b/sdk/python/feast/feature_server.py @@ -3,7 +3,7 @@ import time import traceback from contextlib import asynccontextmanager -from typing import List, Optional +from typing import Any, Dict, List, Mapping, Optional, Sequence, Union import pandas as pd import psutil @@ -33,6 +33,7 @@ init_security_manager, str_to_auth_manager_type, ) +from feast.protos.feast.types.Value_pb2 import RepeatedValue, Value # Define prometheus metrics cpu_usage_gauge = Gauge( @@ -69,8 +70,11 @@ class MaterializeIncrementalRequest(BaseModel): class GetOnlineFeaturesRequest(BaseModel): - entities: dict[str, list] - feature_service: str = None + entities: Union[ + List[Dict[str, Any]], + Mapping[str, Union[Sequence[Any], Sequence[Value], RepeatedValue]], + ] + feature_service: Optional[str] = None features: List[str] = Field(default_factory=list) full_feature_names: bool = False @@ -131,6 +135,7 @@ async def get_online_features(request: GetOnlineFeaturesRequest): ) features = feature_service else: + features = request.features all_feature_views, all_on_demand_feature_views = ( utils._get_feature_views_to_use( store.registry, @@ -250,7 +255,7 @@ async def health(): @app.post("/materialize", dependencies=[Depends(inject_user_details)]) def materialize(request: MaterializeRequest): - for feature_view in request.feature_views: + for feature_view in request.feature_views or []: assert_permissions( resource=feature_view, actions=[AuthzedAction.WRITE_ONLINE] ) @@ -262,7 +267,7 @@ def materialize(request: MaterializeRequest): @app.post("/materialize-incremental", dependencies=[Depends(inject_user_details)]) def materialize_incremental(request: MaterializeIncrementalRequest): - for feature_view in request.feature_views: + for feature_view in request.feature_views or []: assert_permissions( resource=feature_view, actions=[AuthzedAction.WRITE_ONLINE] )