Added the arrow flight interceptor to inject the auth header. (feast-dev#68)

lokeshrangineni · tmihalac · commit 44c339f53a39 · 2024-08-21T13:39:12.000-04:00
* * 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 &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Fixing the failing integration test cases by setting the header in binary format.

Signed-off-by: Lokesh Rangineni &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* 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 &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

* Incorproating code review comments from Daniel.

Signed-off-by: Lokesh Rangineni &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;

---------

Signed-off-by: Lokesh Rangineni &lt;19699092+lokeshrangineni@users.noreply.github.com&gt;
Signed-off-by: Abdul Hameed &lt;ahameed@redhat.com&gt;
diff --git a/sdk/python/feast/infra/offline_stores/remote.py b/sdk/python/feast/infra/offline_stores/remote.py
@@ -354,7 +354,7 @@ def _call_get(
     client: fl.FlightClient,
     command_descriptor: fl.FlightDescriptor,
 ):
-    flight = client.get_flight_info(command_descriptor, options)
+    flight = client.get_flight_info(command_descriptor)
     ticket = flight.endpoints[0].ticket
     reader = client.do_get(ticket)
     return reader.read_all()
diff --git a/sdk/python/feast/permissions/client/auth_client_manager.py b/sdk/python/feast/permissions/client/auth_client_manager.py
@@ -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}"
-        )
diff --git a/sdk/python/feast/permissions/client/utils.py b/sdk/python/feast/permissions/client/utils.py
diff --git a/sdk/python/tests/unit/infra/offline_stores/test_offline_store.py b/sdk/python/tests/unit/infra/offline_stores/test_offline_store.py
@@ -3,7 +3,6 @@
 
 import pandas as pd
 import pyarrow
-import pyarrow.flight as fl
 import pytest
 
 from feast.infra.offline_stores.contrib.athena_offline_store.athena import (
@@ -216,7 +215,6 @@ def retrieval_job(request, environment):
 
         return RemoteRetrievalJob(
             client=MagicMock(),
-            options=fl.FlightCallOptions(),
             api_parameters={
                 "str": "str",
             },
