fix: Refactor auth_client_manager_factory.py in function get_auth_cli…

…ent_m… (feast-dev#4505)

* Refactor auth_client_manager_factory.py in function get_auth_client_manager

Signed-off-by: Theodor Mihalache <[email protected]>

* Refactor auth_client_manager_factory.py in function get_auth_client_manager
-Added test

Signed-off-by: Theodor Mihalache <[email protected]>

* Refactor auth_client_manager_factory.py in function get_auth_client_manager
-updated test following review

Signed-off-by: Theodor Mihalache <[email protected]>

* Refactor auth_client_manager_factory.py in function get_auth_client_manager
-fixed linter

Signed-off-by: Theodor Mihalache <[email protected]>

---------

Signed-off-by: Theodor Mihalache <[email protected]>

sdk/python/feast/permissions/client/arrow_flight_auth_interceptor.py

@@ -2,7 +2,7 @@
 
 from feast.permissions.auth.auth_type import AuthType
 from feast.permissions.auth_model import AuthConfig
-from feast.permissions.client.auth_client_manager_factory import get_auth_token
+from feast.permissions.client.client_auth_token import get_auth_token
 
 
 class FlightBearerTokenInterceptor(fl.ClientMiddleware):

sdk/python/feast/permissions/client/auth_client_manager.py

@@ -1,8 +1,49 @@
+import os
 from abc import ABC, abstractmethod
 
+from feast.permissions.auth.auth_type import AuthType
+from feast.permissions.auth_model import (
+    AuthConfig,
+    KubernetesAuthConfig,
+    OidcClientAuthConfig,
+)
+
 
 class AuthenticationClientManager(ABC):
     @abstractmethod
     def get_token(self) -> str:
         """Retrieves the token based on the authentication type configuration"""
         pass
+
+
+class AuthenticationClientManagerFactory(ABC):
+    def __init__(self, auth_config: AuthConfig):
+        self.auth_config = auth_config
+
+    def get_auth_client_manager(self) -> AuthenticationClientManager:
+        from feast.permissions.client.intra_comm_authentication_client_manager import (
+            IntraCommAuthClientManager,
+        )
+        from feast.permissions.client.kubernetes_auth_client_manager import (
+            KubernetesAuthClientManager,
+        )
+        from feast.permissions.client.oidc_authentication_client_manager import (
+            OidcAuthClientManager,
+        )
+
+        intra_communication_base64 = os.getenv("INTRA_COMMUNICATION_BASE64")
+        if intra_communication_base64:
+            return IntraCommAuthClientManager(
+                self.auth_config, intra_communication_base64
+            )
+
+        if self.auth_config.type == AuthType.OIDC.value:
+            assert isinstance(self.auth_config, OidcClientAuthConfig)
+            return OidcAuthClientManager(self.auth_config)
+        elif self.auth_config.type == AuthType.KUBERNETES.value:
+            assert isinstance(self.auth_config, KubernetesAuthConfig)
+            return KubernetesAuthClientManager(self.auth_config)
+        else:
+            raise RuntimeError(
+                f"No Auth client manager implemented for the auth type:${self.auth_config.type}"
+            )

sdk/python/feast/permissions/client/client_auth_token.py

@@ -0,0 +1,14 @@
+from feast.permissions.auth_model import (
+    AuthConfig,
+)
+from feast.permissions.client.auth_client_manager import (
+    AuthenticationClientManagerFactory,
+)
+
+
+def get_auth_token(auth_config: AuthConfig) -> str:
+    return (
+        AuthenticationClientManagerFactory(auth_config)
+        .get_auth_client_manager()
+        .get_token()
+    )

sdk/python/feast/permissions/client/grpc_client_auth_interceptor.py

@@ -4,7 +4,7 @@
 
 from feast.errors import FeastError
 from feast.permissions.auth_model import AuthConfig
-from feast.permissions.client.auth_client_manager_factory import get_auth_token
+from feast.permissions.client.client_auth_token import get_auth_token
 
 logger = logging.getLogger(__name__)
 

sdk/python/feast/permissions/client/http_auth_requests_wrapper.py

@@ -5,7 +5,7 @@
 from feast.permissions.auth_model import (
     AuthConfig,
 )
-from feast.permissions.client.auth_client_manager_factory import get_auth_token
+from feast.permissions.client.client_auth_token import get_auth_token
 
 
 class AuthenticatedRequestsSession(Session):

sdk/python/feast/permissions/client/intra_comm_authentication_client_manager.py

@@ -0,0 +1,31 @@
+import logging
+
+import jwt
+
+from feast.permissions.auth.auth_type import AuthType
+from feast.permissions.auth_model import AuthConfig
+from feast.permissions.client.auth_client_manager import AuthenticationClientManager
+
+logger = logging.getLogger(__name__)
+
+
+class IntraCommAuthClientManager(AuthenticationClientManager):
+    def __init__(self, auth_config: AuthConfig, intra_communication_base64: str):
+        self.auth_config = auth_config
+        self.intra_communication_base64 = intra_communication_base64
+
+    def get_token(self):
+        if self.auth_config.type == AuthType.OIDC.value:
+            payload = {
+                "preferred_username": f"{self.intra_communication_base64}",  # Subject claim
+            }
+        elif self.auth_config.type == AuthType.KUBERNETES.value:
+            payload = {
+                "sub": f":::{self.intra_communication_base64}",  # Subject claim
+            }
+        else:
+            raise RuntimeError(
+                f"No Auth client manager implemented for the auth type:{self.auth_config.type}"
+            )
+
+        return jwt.encode(payload, "")

sdk/python/tests/unit/permissions/auth/client/test_authentication_client_manager_factory.py

@@ -0,0 +1,55 @@
+import os
+from unittest import mock
+
+import assertpy
+import jwt
+import pytest
+import yaml
+
+from feast.permissions.auth.auth_type import AuthType
+from feast.permissions.auth_model import (
+    AuthConfig,
+)
+from feast.permissions.client.auth_client_manager import (
+    AuthenticationClientManagerFactory,
+)
+from feast.permissions.client.intra_comm_authentication_client_manager import (
+    IntraCommAuthClientManager,
+)
+
+
+@mock.patch.dict(os.environ, {"INTRA_COMMUNICATION_BASE64": "server_intra_com_val"})
+def test_authentication_client_manager_factory(auth_config):
+    raw_config = yaml.safe_load(auth_config)
+    auth_config = AuthConfig(type=raw_config["auth"]["type"])
+
+    authentication_client_manager_factory = AuthenticationClientManagerFactory(
+        auth_config
+    )
+
+    authentication_client_manager = (
+        authentication_client_manager_factory.get_auth_client_manager()
+    )
+
+    if auth_config.type not in [AuthType.KUBERNETES.value, AuthType.OIDC.value]:
+        with pytest.raises(
+            RuntimeError,
+            match=f"No Auth client manager implemented for the auth type:{auth_config.type}",
+        ):
+            authentication_client_manager.get_token()
+    else:
+        token = authentication_client_manager.get_token()
+
+        decoded_token = jwt.decode(token, options={"verify_signature": False})
+        assertpy.assert_that(authentication_client_manager).is_type_of(
+            IntraCommAuthClientManager
+        )
+
+        if AuthType.KUBERNETES.value == auth_config.type:
+            assertpy.assert_that(decoded_token["sub"]).is_equal_to(
+                ":::server_intra_com_val"
+            )
+        elif AuthType.OIDC.value in auth_config.type:
+            assertpy.assert_that(decoded_token["preferred_username"]).is_equal_to(
+                "server_intra_com_val"
+            )