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

lsaiken · web-flow · commit 74c75f1f4c91 · 2022-08-10T10:18:53.000-07:00
Signed-off-by: Lisa Aiken &lt;laiken@manifold.ai&gt;

Signed-off-by: Lisa Aiken &lt;laiken@manifold.ai&gt;
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

Original file line number	Diff line number	Diff line change
`@@ -83,7 +83,9 @@ def get_snowflake_conn(config, autocommit=True) -> SnowflakeConnection:`
`83`	`83`
`84`	`84`	`try:`
`85`	`85`	`conn = snowflake.connector.connect(`
`86`		`- application="feast", autocommit=autocommit, **kwargs`
	`86`	`+ application="feast",`
	`87`	`+ autocommit=autocommit,`
	`88`	`+ **kwargs,`
`87`	`89`	`)`
`88`	`90`
`89`	`91`	`return conn`