diff --git a/sdk/python/feast/feature_server.py b/sdk/python/feast/feature_server.py index f901d0e442..df62a8e013 100644 --- a/sdk/python/feast/feature_server.py +++ b/sdk/python/feast/feature_server.py @@ -27,6 +27,7 @@ ServerType, init_auth_manager, init_security_manager, + str_to_auth_manager_type, ) # Define prometheus metrics @@ -345,7 +346,7 @@ def start_server( monitoring_thread.start() # TODO RBAC remove and use the auth section of the feature store config instead - auth_manager_type = store.config.auth_config.type + auth_manager_type = str_to_auth_manager_type(store.config.auth_config.type) init_security_manager(auth_manager_type=auth_manager_type, fs=store) init_auth_manager( auth_manager_type=auth_manager_type, diff --git a/sdk/python/feast/offline_server.py b/sdk/python/feast/offline_server.py index cc6aa0795a..748237f339 100644 --- a/sdk/python/feast/offline_server.py +++ b/sdk/python/feast/offline_server.py @@ -16,13 +16,13 @@ from feast.permissions.security_manager import assert_permissions from feast.permissions.server.arrow import ( arrowflight_middleware, - inject_user_details, inject_user_details_decorator, ) from feast.permissions.server.utils import ( ServerType, init_auth_manager, init_security_manager, + str_to_auth_manager_type, ) from feast.saved_dataset import SavedDatasetStorage @@ -34,7 +34,9 @@ class OfflineServer(fl.FlightServerBase): def __init__(self, store: FeatureStore, location: str, **kwargs): super(OfflineServer, self).__init__( location, - middleware=arrowflight_middleware(store.config.auth_config.type), + middleware=arrowflight_middleware( + str_to_auth_manager_type(store.config.auth_config.type) + ), **kwargs, ) self._location = location @@ -178,8 +180,6 @@ def _validate_do_get_parameters(self, command: dict): # and returns the stream of data @inject_user_details_decorator def do_get(self, context: fl.ServerCallContext, ticket: fl.Ticket): - inject_user_details(context) - key = ast.literal_eval(ticket.ticket.decode()) if key not in self.flights: logger.error(f"Unknown key {key}") @@ -448,18 +448,22 @@ def remove_dummies(fv: FeatureView) -> FeatureView: return fv -def start_server( - store: FeatureStore, - host: str, - port: int, -): - auth_manager_type = store.config.auth_config.type +def _init_auth_manager(store: FeatureStore): + auth_manager_type = str_to_auth_manager_type(store.config.auth_config.type) init_security_manager(auth_manager_type=auth_manager_type, fs=store) init_auth_manager( auth_manager_type=auth_manager_type, server_type=ServerType.ARROW, ) + +def start_server( + store: FeatureStore, + host: str, + port: int, +): + _init_auth_manager(store) + location = "grpc+tcp://{}:{}".format(host, port) server = OfflineServer(store, location) logger.info(f"Offline store server serving on {location}") diff --git a/sdk/python/feast/permissions/server/arrow.py b/sdk/python/feast/permissions/server/arrow.py index 3eb5419354..47594bdaf9 100644 --- a/sdk/python/feast/permissions/server/arrow.py +++ b/sdk/python/feast/permissions/server/arrow.py @@ -24,7 +24,7 @@ def arrowflight_middleware( - auth_manager_type: str, + auth_manager_type: AuthManagerType, ) -> Optional[dict[str, fl.ServerMiddlewareFactory]]: """ A dictionary with the configured middlewares to support extracting the user details when the authorization manager is defined. diff --git a/sdk/python/feast/permissions/server/utils.py b/sdk/python/feast/permissions/server/utils.py index ddf6ffa5dc..b4c0340426 100644 --- a/sdk/python/feast/permissions/server/utils.py +++ b/sdk/python/feast/permissions/server/utils.py @@ -4,7 +4,6 @@ import enum import logging -import os import feast from feast.permissions.auth.auth_manager import ( @@ -62,19 +61,6 @@ def str_to_auth_manager_type(value: str) -> AuthManagerType: return AuthManagerType.NONE -# TODO RBAC: will remove once we can manage the auth configuration -def auth_manager_type_from_env() -> AuthManagerType: - type = os.getenv("AUTH_MANAGER_TYPE", "none") - print(f"Configuring authentication manager for {type}") - - if type.lower() == AuthManagerType.OIDC.value: - return AuthManagerType.OIDC - if type.lower() == AuthManagerType.KUBERNETES.value: - return AuthManagerType.KUBERNETES - - return AuthManagerType.NONE - - def init_security_manager(auth_manager_type: AuthManagerType, fs: "feast.FeatureStore"): """ Initialize the global security manager. diff --git a/sdk/python/tests/unit/test_offline_server.py b/sdk/python/tests/unit/test_offline_server.py index 5991e7450d..237e2ecad4 100644 --- a/sdk/python/tests/unit/test_offline_server.py +++ b/sdk/python/tests/unit/test_offline_server.py @@ -14,7 +14,7 @@ RemoteOfflineStore, RemoteOfflineStoreConfig, ) -from feast.offline_server import OfflineServer +from feast.offline_server import OfflineServer, _init_auth_manager from feast.repo_config import RepoConfig from tests.utils.cli_repo_creator import CliRunner @@ -26,6 +26,7 @@ def empty_offline_server(environment): store = environment.feature_store location = "grpc+tcp://localhost:0" + _init_auth_manager(store=store) return OfflineServer(store=store, location=location) @@ -102,6 +103,8 @@ def test_remote_offline_store_apis(): with tempfile.TemporaryDirectory() as temp_dir: store = default_store(str(temp_dir)) location = "grpc+tcp://localhost:0" + + _init_auth_manager(store=store) server = OfflineServer(store=store, location=location) assertpy.assert_that(server).is_not_none