fix: Add X-Trino-Extra-Credential header and remove user override (f…

…east-dev#3246) * do changes with signing Signed-off-by: ammarar <[email protected]> * lint fix Signed-off-by: ammarar <[email protected]> Signed-off-by: ammarar <[email protected]>
Affirm · Sep 26, 2022 · f38506c · f38506c
1 parent f242e24
commit f38506c
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/contrib/trino_offline_store/trino_queries.py b/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/trino_queries.py
@@ -36,13 +36,17 @@ def __init__(
         catalog: Optional[str] = None,
         auth: Optional[Any] = None,
         http_scheme: Optional[str] = None,
+        source: Optional[str] = None,
+        extra_credential: Optional[str] = None,
     ):
         self.host = host or os.getenv("TRINO_HOST")
         self.port = port or os.getenv("TRINO_PORT")
         self.user = user or os.getenv("TRINO_USER")
         self.catalog = catalog or os.getenv("TRINO_CATALOG")
         self.auth = auth or os.getenv("TRINO_AUTH")
         self.http_scheme = http_scheme or os.getenv("TRINO_HTTP_SCHEME")
+        self.source = source or os.getenv("TRINO_SOURCE")
+        self.extra_credential = extra_credential or os.getenv("TRINO_EXTRA_CREDENTIAL")
         self._cursor: Optional[Cursor] = None
 
         if self.host is None:
@@ -56,13 +60,20 @@ def __init__(
 
     def _get_cursor(self) -> Cursor:
         if self._cursor is None:
+            headers = (
+                {trino.constants.HEADER_EXTRA_CREDENTIAL: self.extra_credential}
+                if self.extra_credential
+                else {}
+            )
             self._cursor = trino.dbapi.connect(
                 host=self.host,
                 port=self.port,
                 user=self.user,
                 catalog=self.catalog,
                 auth=self.auth,
                 http_scheme=self.http_scheme,
+                source=self.source,
+                http_headers=headers,
             ).cursor()
 
         return self._cursor

diff --git a/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/trino_source.py b/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/trino_source.py
@@ -228,7 +228,6 @@ def get_table_column_names_and_types(
         self, config: RepoConfig
     ) -> Iterable[Tuple[str, str]]:
         client = Trino(
-            user="user",
             catalog=config.offline_store.catalog,
             host=config.offline_store.host,
             port=config.offline_store.port,