feat: Add authentication option for snowflake connector (#3039)

Signed-off-by: Lisa Aiken <laiken@manifold.ai> Signed-off-by: Lisa Aiken <laiken@manifold.ai>
feast-dev · Aug 10, 2022 · 74c75f1 · 74c75f1
1 parent 8527457
commit 74c75f1
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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