Skip to content

Commit

Permalink
code review fixes including the unit test and integration test as sug…
Browse files Browse the repository at this point in the history
…gested

Signed-off-by: Abdul Hameed <ahameed@redhat.com>
  • Loading branch information
redhatHameed committed Aug 14, 2024
1 parent 2d25dc6 commit 7fd3331
Show file tree
Hide file tree
Showing 5 changed files with 21 additions and 27 deletions.
3 changes: 2 additions & 1 deletion sdk/python/feast/feature_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
ServerType,
init_auth_manager,
init_security_manager,
str_to_auth_manager_type,
)

# Define prometheus metrics
Expand Down Expand Up @@ -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,
Expand Down
24 changes: 14 additions & 10 deletions sdk/python/feast/offline_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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}")
Expand Down Expand Up @@ -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}")
Expand Down
2 changes: 1 addition & 1 deletion sdk/python/feast/permissions/server/arrow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
14 changes: 0 additions & 14 deletions sdk/python/feast/permissions/server/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@

import enum
import logging
import os

import feast
from feast.permissions.auth.auth_manager import (
Expand Down Expand Up @@ -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.
Expand Down
5 changes: 4 additions & 1 deletion sdk/python/tests/unit/test_offline_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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)


Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 7fd3331

Please sign in to comment.