From 3293c85e83084333c26e68c935d80b6d2b9394ba Mon Sep 17 00:00:00 2001 From: ammarar Date: Fri, 23 Sep 2022 18:10:11 -0700 Subject: [PATCH 1/2] do changes with signing Signed-off-by: ammarar --- .../contrib/trino_offline_store/trino_queries.py | 7 +++++++ .../contrib/trino_offline_store/trino_source.py | 1 - 2 files changed, 7 insertions(+), 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 index 1d4b5881240..95efd881e09 100644 --- 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,6 +36,8 @@ 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") @@ -43,6 +45,8 @@ def __init__( 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,6 +60,7 @@ 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, @@ -63,6 +68,8 @@ def _get_cursor(self) -> Cursor: 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 index 6e989bd40ce..f09b79069ca 100644 --- 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, From dae2b2f2f1a9aaa9dd9425159ff79bfc5583b2bd Mon Sep 17 00:00:00 2001 From: ammarar Date: Fri, 23 Sep 2022 18:39:40 -0700 Subject: [PATCH 2/2] lint fix Signed-off-by: ammarar --- .../contrib/trino_offline_store/trino_queries.py | 6 +++++- 1 file changed, 5 insertions(+), 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 index 95efd881e09..97c61f78a60 100644 --- 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 @@ -60,7 +60,11 @@ 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 {} + 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,