Implementation of oidc client authentication. (feast-dev#40)

* Adding initial draft code to manage the oidc client authentication. Signed-off-by: Lokesh Rangineni <lokeshforjava@gmail.com> * Adding initial draft code to manage the oidc client authentication. Signed-off-by: Lokesh Rangineni <lokeshforjava@gmail.com> * Incorporating code review comments. Signed-off-by: Lokesh Rangineni <lokeshforjava@gmail.com> --------- Signed-off-by: Lokesh Rangineni <lokeshforjava@gmail.com> Signed-off-by: Abdul Hameed <ahameed@redhat.com> Signed-off-by: Lokesh Rangineni <19699092+lokeshrangineni@users.noreply.github.com>
lokeshrangineni · Aug 7, 2024 · ee71332 · ee71332
1 parent 338b443
commit ee71332
Show file tree

Hide file tree

Showing 8 changed files with 137 additions and 9 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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)