Client module-grpc

- Made changes following code review Signed-off-by: Theodor Mihalache <tmihalac@redhat.com>
dmartinol · Jul 12, 2024 · e76a97b · e76a97b
1 parent bc4d8d3
commit e76a97b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/sdk/python/feast/infra/offline_stores/remote.py b/sdk/python/feast/infra/offline_stores/remote.py
@@ -27,7 +27,7 @@
     RetrievalMetadata,
 )
 from feast.infra.registry.base_registry import BaseRegistry
-from feast.permissions.client.auth_client_manager import create_flight_call_options
+from feast.permissions.client.utils import create_flight_call_options
 from feast.repo_config import FeastConfigBaseModel, RepoConfig
 from feast.saved_dataset import SavedDatasetStorage
 

diff --git a/sdk/python/feast/permissions/client/utils.py b/sdk/python/feast/permissions/client/utils.py
@@ -1,3 +1,6 @@
+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 import get_auth_client_manager
 
@@ -9,3 +12,10 @@ def create_auth_header(
     token = auth_client_manager.get_token()
 
     return (("authorization", "Bearer " + token),)
+
+
+def create_flight_call_options(auth_config: AuthConfig) -> fl.FlightCallOptions:
+    if auth_config.type != AuthType.NONE.value:
+        headers = create_auth_header(auth_config)
+        return fl.FlightCallOptions(headers=headers)
+    return fl.FlightCallOptions()