From 74c75f1f4c91f0097f9a1085a4e68a07c524037d Mon Sep 17 00:00:00 2001
From: Lisa Aiken <lisa.aiken@gmail.com>
Date: Wed, 10 Aug 2022 10:18:53 -0700
Subject: [PATCH] feat: Add authentication option for snowflake connector
 (#3039)

Signed-off-by: Lisa Aiken <laiken@manifold.ai>

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