From 5aa3d0026a38b954dace8ca36951484a34ad113c Mon Sep 17 00:00:00 2001 From: Lisa Aiken Date: Wed, 10 Aug 2022 10:18:53 -0700 Subject: [PATCH] feat: Add authentication option for snowflake connector (#3039) Signed-off-by: Lisa Aiken Signed-off-by: Lisa Aiken Signed-off-by: Francisco Javier Arceo --- sdk/python/feast/infra/offline_stores/snowflake.py | 3 +++ sdk/python/feast/infra/offline_stores/snowflake_source.py | 6 +++++- sdk/python/feast/infra/online_stores/snowflake.py | 3 +++ sdk/python/feast/infra/utils/snowflake_utils.py | 4 +++- 4 files changed, 14 insertions(+), 2 deletions(-) diff --git a/sdk/python/feast/infra/offline_stores/snowflake.py b/sdk/python/feast/infra/offline_stores/snowflake.py index 5936151d71..7ad764f34d 100644 --- a/sdk/python/feast/infra/offline_stores/snowflake.py +++ b/sdk/python/feast/infra/offline_stores/snowflake.py @@ -91,6 +91,9 @@ class SnowflakeOfflineStoreConfig(FeastConfigBaseModel): schema_: Optional[str] = Field(None, alias="schema") """ Snowflake schema name """ + authenticator: Optional[str] = None + """ Snowflake authenticator name """ + storage_integration_name: Optional[str] = None """ Storage integration name in snowflake """ diff --git a/sdk/python/feast/infra/offline_stores/snowflake_source.py b/sdk/python/feast/infra/offline_stores/snowflake_source.py index 63ccfb9dae..df0aef2ade 100644 --- a/sdk/python/feast/infra/offline_stores/snowflake_source.py +++ b/sdk/python/feast/infra/offline_stores/snowflake_source.py @@ -298,7 +298,11 @@ class SavedDatasetSnowflakeStorage(SavedDatasetStorage): def __init__(self, table_ref: str): self.snowflake_options = SnowflakeOptions( - database=None, schema=None, table=table_ref, query=None, warehouse=None + database=None, + schema=None, + table=table_ref, + query=None, + warehouse=None, ) @staticmethod diff --git a/sdk/python/feast/infra/online_stores/snowflake.py b/sdk/python/feast/infra/online_stores/snowflake.py index 73c68e4bc0..88b141981f 100644 --- a/sdk/python/feast/infra/online_stores/snowflake.py +++ b/sdk/python/feast/infra/online_stores/snowflake.py @@ -53,6 +53,9 @@ class SnowflakeOnlineStoreConfig(FeastConfigBaseModel): schema_: Optional[str] = Field("PUBLIC", alias="schema") """ Snowflake schema name """ + authenticator: Optional[str] = None + """ Snowflake authenticator name """ + class Config: allow_population_by_field_name = True diff --git a/sdk/python/feast/infra/utils/snowflake_utils.py b/sdk/python/feast/infra/utils/snowflake_utils.py index f54288e45d..6ba9e54489 100644 --- a/sdk/python/feast/infra/utils/snowflake_utils.py +++ b/sdk/python/feast/infra/utils/snowflake_utils.py @@ -83,7 +83,9 @@ def get_snowflake_conn(config, autocommit=True) -> SnowflakeConnection: try: conn = snowflake.connector.connect( - application="feast", autocommit=autocommit, **kwargs + application="feast", + autocommit=autocommit, + **kwargs, ) return conn