diff --git a/sdk/python/feast/infra/online_stores/remote.py b/sdk/python/feast/infra/online_stores/remote.py index 19e1b7d5159..93fbcaf7716 100644 --- a/sdk/python/feast/infra/online_stores/remote.py +++ b/sdk/python/feast/infra/online_stores/remote.py @@ -16,11 +16,13 @@ from datetime import datetime from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple -import requests from pydantic import StrictStr from feast import Entity, FeatureView, RepoConfig from feast.infra.online_stores.online_store import OnlineStore +from feast.permissions.client.http_auth_requests_wrapper import ( + get_http_auth_requests_session, +) from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto from feast.protos.feast.types.Value_pb2 import Value as ValueProto from feast.repo_config import FeastConfigBaseModel @@ -70,7 +72,7 @@ def online_read( req_body = self._construct_online_read_api_json_request( entity_keys, table, requested_features ) - response = requests.post( + response = get_http_auth_requests_session(config.auth_config).post( f"{config.online_store.path}/get-online-features", data=req_body ) if response.status_code == 200: diff --git a/sdk/python/feast/permissions/auth_model.py b/sdk/python/feast/permissions/auth_model.py index 860b4d1f891..afb0a22bc9e 100644 --- a/sdk/python/feast/permissions/auth_model.py +++ b/sdk/python/feast/permissions/auth_model.py @@ -1,4 +1,4 @@ -from typing import Literal +from typing import Literal, Optional from feast.repo_config import FeastConfigBaseModel @@ -8,9 +8,10 @@ class AuthConfig(FeastConfigBaseModel): class OidcAuthConfig(AuthConfig): - auth_server_url: str + auth_server_url: Optional[str] = None + auth_discovery_url: str client_id: str - client_secret: str + client_secret: Optional[str] = None username: str password: str realm: str = "master" @@ -20,5 +21,5 @@ class NoAuthConfig(AuthConfig): pass -class K8AuthConfig(AuthConfig): +class KubernetesAuthConfig(AuthConfig): pass diff --git a/sdk/python/feast/permissions/client/auth_client_manager.py b/sdk/python/feast/permissions/client/auth_client_manager.py new file mode 100644 index 00000000000..82f9b7433e8 --- /dev/null +++ b/sdk/python/feast/permissions/client/auth_client_manager.py @@ -0,0 +1,8 @@ +from abc import ABC, abstractmethod + + +class AuthenticationClientManager(ABC): + @abstractmethod + def get_token(self) -> str: + """Retrieves the token based on the authentication type configuration""" + pass diff --git a/sdk/python/feast/permissions/client/http_auth_requests_wrapper.py b/sdk/python/feast/permissions/client/http_auth_requests_wrapper.py new file mode 100644 index 00000000000..25b52ff5a54 --- /dev/null +++ b/sdk/python/feast/permissions/client/http_auth_requests_wrapper.py @@ -0,0 +1,43 @@ +import requests +from requests import Session + +from feast.permissions.auth_model import ( + AuthConfig, + KubernetesAuthConfig, + OidcAuthConfig, +) +from feast.permissions.client.kubernetes_auth_client_manager import ( + KubernetesAuthClientManager, +) +from feast.permissions.client.oidc_authentication_client_manager import ( + OidcAuthClientManager, +) + + +class AuthenticatedRequestsSession(Session): + def __init__(self, auth_token: str): + super().__init__() + self.auth_token = auth_token + self.headers.update({"Authorization": f"Bearer {self.auth_token}"}) + + +def get_auth_client_manager(auth_config: AuthConfig): + if auth_config.type == "oidc": + assert isinstance(auth_config, OidcAuthConfig) + return OidcAuthClientManager(auth_config) + elif auth_config.type == "kubernetes": + assert isinstance(auth_config, KubernetesAuthConfig) + return KubernetesAuthClientManager(auth_config) + else: + raise RuntimeError( + f"No Auth client manager implemented for the auth type:${auth_config.type}" + ) + + +def get_http_auth_requests_session(auth_config: AuthConfig) -> Session: + if auth_config.type == "no_auth": + request_session = requests.session() + else: + auth_client_manager = get_auth_client_manager(auth_config) + request_session = AuthenticatedRequestsSession(auth_client_manager.get_token()) + return request_session diff --git a/sdk/python/feast/permissions/client/kubernetes_auth_client_manager.py b/sdk/python/feast/permissions/client/kubernetes_auth_client_manager.py new file mode 100644 index 00000000000..ef74fba3a99 --- /dev/null +++ b/sdk/python/feast/permissions/client/kubernetes_auth_client_manager.py @@ -0,0 +1,11 @@ +from feast.permissions.auth_model import KubernetesAuthConfig +from feast.permissions.client.auth_client_manager import AuthenticationClientManager + + +class KubernetesAuthClientManager(AuthenticationClientManager): + def __init__(self, auth_config: KubernetesAuthConfig): + self.auth_config = auth_config + + # TODO: needs to implement this for k8 auth. + def get_token(self): + return "" diff --git a/sdk/python/feast/permissions/client/oidc_authentication_client_manager.py b/sdk/python/feast/permissions/client/oidc_authentication_client_manager.py new file mode 100644 index 00000000000..a16edd6ebde --- /dev/null +++ b/sdk/python/feast/permissions/client/oidc_authentication_client_manager.py @@ -0,0 +1,52 @@ +import logging + +import requests + +from feast.permissions.auth_model import OidcAuthConfig +from feast.permissions.client.auth_client_manager import AuthenticationClientManager + +logger = logging.getLogger(__name__) + + +class OidcAuthClientManager(AuthenticationClientManager): + def __init__(self, auth_config: OidcAuthConfig): + self.auth_config = auth_config + + def _get_token_endpoint(self): + response = requests.get(self.auth_config.auth_discovery_url) + if response.status_code == 200: + oidc_config = response.json() + if not oidc_config["token_endpoint"]: + raise RuntimeError( + " OIDC token_endpoint is not available from discovery url response." + ) + return oidc_config["token_endpoint"] + else: + raise RuntimeError( + f"Error fetching OIDC token endpoint configuration: {response.status_code} - {response.text}" + ) + + def get_token(self): + # Fetch the token endpoint from the discovery URL + token_endpoint = self._get_token_endpoint() + + token_request_body = { + "grant_type": "password", + "client_id": self.auth_config.client_id, + "username": self.auth_config.username, + "password": self.auth_config.password, + } + + token_response = requests.post(token_endpoint, data=token_request_body) + if token_response.status_code == 200: + access_token = token_response.json()["access_token"] + if not access_token: + logger.debug( + f"access_token is empty for the client_id=${self.auth_config.client_id}" + ) + raise RuntimeError("access token is empty") + return access_token + else: + raise RuntimeError( + "Failed to obtain access token: {token_response.status_code} - {token_response.text}" + ) diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py index 4f09b971693..c5f6c93c6c5 100644 --- a/sdk/python/feast/repo_config.py +++ b/sdk/python/feast/repo_config.py @@ -89,7 +89,7 @@ AUTH_CONFIGS_CLASS_FOR_TYPE = { "no_auth": "feast.permissions.auth_model.NoAuthConfig", - "kubernetes": "feast.permissions.auth_model.K8AuthConfig", + "kubernetes": "feast.permissions.auth_model.KubernetesAuthConfig", "oidc": "feast.permissions.auth_model.OidcAuthConfig", } diff --git a/sdk/python/tests/unit/infra/scaffolding/test_repo_config.py b/sdk/python/tests/unit/infra/scaffolding/test_repo_config.py index cc4a77662f3..76ad85ca7c6 100644 --- a/sdk/python/tests/unit/infra/scaffolding/test_repo_config.py +++ b/sdk/python/tests/unit/infra/scaffolding/test_repo_config.py @@ -4,7 +4,11 @@ from typing import Optional from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig -from feast.permissions.auth_model import K8AuthConfig, NoAuthConfig, OidcAuthConfig +from feast.permissions.auth_model import ( + KubernetesAuthConfig, + NoAuthConfig, + OidcAuthConfig, +) from feast.repo_config import FeastConfigError, load_repo_config @@ -210,6 +214,7 @@ def test_auth_config(): password: test_password realm: master auth_server_url: http://localhost:8712 + auth_discovery_url: http://localhost:8080/realms/master/.well-known/openid-configuration registry: "registry.db" provider: local online_store: @@ -232,6 +237,7 @@ def test_auth_config(): password: test_password realm: master auth_server_url: http://localhost:8712 + auth_discovery_url: http://localhost:8080/realms/master/.well-known/openid-configuration registry: "registry.db" provider: local online_store: @@ -254,6 +260,7 @@ def test_auth_config(): password: test_password realm: master auth_server_url: http://localhost:8712 + auth_discovery_url: http://localhost:8080/realms/master/.well-known/openid-configuration registry: "registry.db" provider: local online_store: @@ -271,6 +278,10 @@ def test_auth_config(): assert oidc_repo_config.auth_config.password == "test_password" assert oidc_repo_config.auth_config.realm == "master" assert oidc_repo_config.auth_config.auth_server_url == "http://localhost:8712" + assert ( + oidc_repo_config.auth_config.auth_discovery_url + == "http://localhost:8080/realms/master/.well-known/openid-configuration" + ) no_auth_repo_config = _test_config( dedent( @@ -304,4 +315,4 @@ def test_auth_config(): expect_error=None, ) assert k8_repo_config.auth.get("type") == "kubernetes" - assert isinstance(k8_repo_config.auth_config, K8AuthConfig) + assert isinstance(k8_repo_config.auth_config, KubernetesAuthConfig)