forked from feast-dev/feast
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Added the arrow flight interceptor to inject the auth header. (feast-…
…dev#68) * * Added the arrow flight interceptor to inject the auth header. * Injecting grpc interceptor if it is needed when auth type is not NO_AUTH. Signed-off-by: Lokesh Rangineni <19699092+lokeshrangineni@users.noreply.github.com> * Fixing the failing integration test cases by setting the header in binary format. Signed-off-by: Lokesh Rangineni <19699092+lokeshrangineni@users.noreply.github.com> * Refactored method and moved to factory class to incorporate code review comment. Fixed lint error by removing the type of port. and other minor changes. Signed-off-by: Lokesh Rangineni <19699092+lokeshrangineni@users.noreply.github.com> * Incorproating code review comments from Daniel. Signed-off-by: Lokesh Rangineni <19699092+lokeshrangineni@users.noreply.github.com> --------- Signed-off-by: Lokesh Rangineni <19699092+lokeshrangineni@users.noreply.github.com>
- Loading branch information
1 parent
c13f229
commit 0aad7a8
Showing
9 changed files
with
106 additions
and
113 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
38 changes: 38 additions & 0 deletions
38
sdk/python/feast/permissions/client/arrow_flight_auth_interceptor.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,38 @@ | ||
| import pyarrow.flight as fl | ||
|
|
||
| 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 | ||
|
|
||
|
|
||
| class FlightBearerTokenInterceptor(fl.ClientMiddleware): | ||
| def __init__(self, auth_config: AuthConfig): | ||
| super().__init__() | ||
| self.auth_config = auth_config | ||
|
|
||
| def call_completed(self, exception): | ||
| pass | ||
|
|
||
| def received_headers(self, headers): | ||
| pass | ||
|
|
||
| def sending_headers(self): | ||
| access_token = get_auth_token(self.auth_config) | ||
| return {b"authorization": b"Bearer " + access_token.encode("utf-8")} | ||
|
|
||
|
|
||
| class FlightAuthInterceptorFactory(fl.ClientMiddlewareFactory): | ||
| def __init__(self, auth_config: AuthConfig): | ||
| super().__init__() | ||
| self.auth_config = auth_config | ||
|
|
||
| def start_call(self, info): | ||
| return FlightBearerTokenInterceptor(self.auth_config) | ||
|
|
||
|
|
||
| def build_arrow_flight_client(host: str, port, auth_config: AuthConfig): | ||
| if auth_config.type != AuthType.NONE.value: | ||
| middleware_factory = FlightAuthInterceptorFactory(auth_config) | ||
| return fl.FlightClient(f"grpc://{host}:{port}", middleware=[middleware_factory]) | ||
| else: | ||
| return fl.FlightClient(f"grpc://{host}:{port}") |
30 changes: 0 additions & 30 deletions
30
sdk/python/feast/permissions/client/auth_client_manager.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,38 +1,8 @@ | ||
| from abc import ABC, abstractmethod | ||
|
|
||
| from feast.permissions.auth.auth_type import AuthType | ||
| from feast.permissions.auth_model import ( | ||
| AuthConfig, | ||
| KubernetesAuthConfig, | ||
| OidcAuthConfig, | ||
| ) | ||
|
|
||
|
|
||
| class AuthenticationClientManager(ABC): | ||
| @abstractmethod | ||
| def get_token(self) -> str: | ||
| """Retrieves the token based on the authentication type configuration""" | ||
| pass | ||
|
|
||
|
|
||
| def get_auth_client_manager(auth_config: AuthConfig) -> AuthenticationClientManager: | ||
| if auth_config.type == AuthType.OIDC.value: | ||
| assert isinstance(auth_config, OidcAuthConfig) | ||
|
|
||
| from feast.permissions.client.oidc_authentication_client_manager import ( | ||
| OidcAuthClientManager, | ||
| ) | ||
|
|
||
| return OidcAuthClientManager(auth_config) | ||
| elif auth_config.type == AuthType.KUBERNETES.value: | ||
| assert isinstance(auth_config, KubernetesAuthConfig) | ||
|
|
||
| from feast.permissions.client.kubernetes_auth_client_manager import ( | ||
| KubernetesAuthClientManager, | ||
| ) | ||
|
|
||
| return KubernetesAuthClientManager(auth_config) | ||
| else: | ||
| raise RuntimeError( | ||
| f"No Auth client manager implemented for the auth type:${auth_config.type}" | ||
| ) |
30 changes: 30 additions & 0 deletions
30
sdk/python/feast/permissions/client/auth_client_manager_factory.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,30 @@ | ||
| from feast.permissions.auth.auth_type import AuthType | ||
| from feast.permissions.auth_model import ( | ||
| AuthConfig, | ||
| KubernetesAuthConfig, | ||
| OidcAuthConfig, | ||
| ) | ||
| from feast.permissions.client.auth_client_manager import AuthenticationClientManager | ||
| from feast.permissions.client.kubernetes_auth_client_manager import ( | ||
| KubernetesAuthClientManager, | ||
| ) | ||
| from feast.permissions.client.oidc_authentication_client_manager import ( | ||
| OidcAuthClientManager, | ||
| ) | ||
|
|
||
|
|
||
| def get_auth_client_manager(auth_config: AuthConfig) -> AuthenticationClientManager: | ||
| if auth_config.type == AuthType.OIDC.value: | ||
| assert isinstance(auth_config, OidcAuthConfig) | ||
| return OidcAuthClientManager(auth_config) | ||
| elif auth_config.type == AuthType.KUBERNETES.value: | ||
| 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_auth_token(auth_config: AuthConfig) -> str: | ||
| return get_auth_client_manager(auth_config).get_token() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.