Skip to content

Commit def8633

Browse files
tmihalactmihalac
authored
fix: Refactor auth_client_manager_factory.py in function get_auth_client_m… (feast-dev#4505)
* Refactor auth_client_manager_factory.py in function get_auth_client_manager Signed-off-by: Theodor Mihalache <tmihalac@redhat.com> * Refactor auth_client_manager_factory.py in function get_auth_client_manager -Added test Signed-off-by: Theodor Mihalache <tmihalac@redhat.com> * Refactor auth_client_manager_factory.py in function get_auth_client_manager -updated test following review Signed-off-by: Theodor Mihalache <tmihalac@redhat.com> * Refactor auth_client_manager_factory.py in function get_auth_client_manager -fixed linter Signed-off-by: Theodor Mihalache <tmihalac@redhat.com> --------- Signed-off-by: Theodor Mihalache <tmihalac@redhat.com>
1 parent 0c90137 commit def8633

8 files changed

+144
-44
lines changed

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from feast.permissions.auth.auth_type import AuthType
44
from feast.permissions.auth_model import AuthConfig
5-
from feast.permissions.client.auth_client_manager_factory import get_auth_token
5+
from feast.permissions.client.client_auth_token import get_auth_token
66

77

88
class FlightBearerTokenInterceptor(fl.ClientMiddleware):

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

+41
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,49 @@
1+
import os
12
from abc import ABC, abstractmethod
23

4+
from feast.permissions.auth.auth_type import AuthType
5+
from feast.permissions.auth_model import (
6+
AuthConfig,
7+
KubernetesAuthConfig,
8+
OidcClientAuthConfig,
9+
)
10+
311

412
class AuthenticationClientManager(ABC):
513
@abstractmethod
614
def get_token(self) -> str:
715
"""Retrieves the token based on the authentication type configuration"""
816
pass
17+
18+
19+
class AuthenticationClientManagerFactory(ABC):
20+
def __init__(self, auth_config: AuthConfig):
21+
self.auth_config = auth_config
22+
23+
def get_auth_client_manager(self) -> AuthenticationClientManager:
24+
from feast.permissions.client.intra_comm_authentication_client_manager import (
25+
IntraCommAuthClientManager,
26+
)
27+
from feast.permissions.client.kubernetes_auth_client_manager import (
28+
KubernetesAuthClientManager,
29+
)
30+
from feast.permissions.client.oidc_authentication_client_manager import (
31+
OidcAuthClientManager,
32+
)
33+
34+
intra_communication_base64 = os.getenv("INTRA_COMMUNICATION_BASE64")
35+
if intra_communication_base64:
36+
return IntraCommAuthClientManager(
37+
self.auth_config, intra_communication_base64
38+
)
39+
40+
if self.auth_config.type == AuthType.OIDC.value:
41+
assert isinstance(self.auth_config, OidcClientAuthConfig)
42+
return OidcAuthClientManager(self.auth_config)
43+
elif self.auth_config.type == AuthType.KUBERNETES.value:
44+
assert isinstance(self.auth_config, KubernetesAuthConfig)
45+
return KubernetesAuthClientManager(self.auth_config)
46+
else:
47+
raise RuntimeError(
48+
f"No Auth client manager implemented for the auth type:${self.auth_config.type}"
49+
)

‎sdk/python/feast/permissions/client/auth_client_manager_factory.py

-41
This file was deleted.

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

+14
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
from feast.permissions.auth_model import (
2+
AuthConfig,
3+
)
4+
from feast.permissions.client.auth_client_manager import (
5+
AuthenticationClientManagerFactory,
6+
)
7+
8+
9+
def get_auth_token(auth_config: AuthConfig) -> str:
10+
return (
11+
AuthenticationClientManagerFactory(auth_config)
12+
.get_auth_client_manager()
13+
.get_token()
14+
)

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
from feast.errors import FeastError
66
from feast.permissions.auth_model import AuthConfig
7-
from feast.permissions.client.auth_client_manager_factory import get_auth_token
7+
from feast.permissions.client.client_auth_token import get_auth_token
88

99
logger = logging.getLogger(__name__)
1010

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from feast.permissions.auth_model import (
66
AuthConfig,
77
)
8-
from feast.permissions.client.auth_client_manager_factory import get_auth_token
8+
from feast.permissions.client.client_auth_token import get_auth_token
99

1010

1111
class AuthenticatedRequestsSession(Session):

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

+31
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
import logging
2+
3+
import jwt
4+
5+
from feast.permissions.auth.auth_type import AuthType
6+
from feast.permissions.auth_model import AuthConfig
7+
from feast.permissions.client.auth_client_manager import AuthenticationClientManager
8+
9+
logger = logging.getLogger(__name__)
10+
11+
12+
class IntraCommAuthClientManager(AuthenticationClientManager):
13+
def __init__(self, auth_config: AuthConfig, intra_communication_base64: str):
14+
self.auth_config = auth_config
15+
self.intra_communication_base64 = intra_communication_base64
16+
17+
def get_token(self):
18+
if self.auth_config.type == AuthType.OIDC.value:
19+
payload = {
20+
"preferred_username": f"{self.intra_communication_base64}", # Subject claim
21+
}
22+
elif self.auth_config.type == AuthType.KUBERNETES.value:
23+
payload = {
24+
"sub": f":::{self.intra_communication_base64}", # Subject claim
25+
}
26+
else:
27+
raise RuntimeError(
28+
f"No Auth client manager implemented for the auth type:{self.auth_config.type}"
29+
)
30+
31+
return jwt.encode(payload, "")

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

+55
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
import os
2+
from unittest import mock
3+
4+
import assertpy
5+
import jwt
6+
import pytest
7+
import yaml
8+
9+
from feast.permissions.auth.auth_type import AuthType
10+
from feast.permissions.auth_model import (
11+
AuthConfig,
12+
)
13+
from feast.permissions.client.auth_client_manager import (
14+
AuthenticationClientManagerFactory,
15+
)
16+
from feast.permissions.client.intra_comm_authentication_client_manager import (
17+
IntraCommAuthClientManager,
18+
)
19+
20+
21+
@mock.patch.dict(os.environ, {"INTRA_COMMUNICATION_BASE64": "server_intra_com_val"})
22+
def test_authentication_client_manager_factory(auth_config):
23+
raw_config = yaml.safe_load(auth_config)
24+
auth_config = AuthConfig(type=raw_config["auth"]["type"])
25+
26+
authentication_client_manager_factory = AuthenticationClientManagerFactory(
27+
auth_config
28+
)
29+
30+
authentication_client_manager = (
31+
authentication_client_manager_factory.get_auth_client_manager()
32+
)
33+
34+
if auth_config.type not in [AuthType.KUBERNETES.value, AuthType.OIDC.value]:
35+
with pytest.raises(
36+
RuntimeError,
37+
match=f"No Auth client manager implemented for the auth type:{auth_config.type}",
38+
):
39+
authentication_client_manager.get_token()
40+
else:
41+
token = authentication_client_manager.get_token()
42+
43+
decoded_token = jwt.decode(token, options={"verify_signature": False})
44+
assertpy.assert_that(authentication_client_manager).is_type_of(
45+
IntraCommAuthClientManager
46+
)
47+
48+
if AuthType.KUBERNETES.value == auth_config.type:
49+
assertpy.assert_that(decoded_token["sub"]).is_equal_to(
50+
":::server_intra_com_val"
51+
)
52+
elif AuthType.OIDC.value in auth_config.type:
53+
assertpy.assert_that(decoded_token["preferred_username"]).is_equal_to(
54+
"server_intra_com_val"
55+
)

0 commit comments

Comments
 (0)