Allow authenticator parameter when connecting to snowflake

Author: pappacena

landingai/predict/snowflake.py

@@ -28,6 +28,8 @@ def __init__(
         # TODO: Remove this once we remove the API key auth from snowflake
         api_key: Optional[str] = None,
         check_server_ready: bool = True,
+        *,
+        snowflake_authenticator: Optional[str] = None,
     ) -> None:
         assert (
             snowflake_password is not None or snowflake_private_key is not None
@@ -40,6 +42,7 @@ def __init__(
         self.snowflake_user = snowflake_user
         self.snowflake_password = snowflake_password
         self.snowflake_private_key = snowflake_private_key
+        self.snowflake_authenticator = snowflake_authenticator
 
         self._auth_token = None
         self._last_auth_token_fetch: Optional[datetime.datetime] = None
@@ -86,6 +89,8 @@ def _get_auth_token(self) -> str:
                 format=serialization.PrivateFormat.PKCS8,
                 encryption_algorithm=serialization.NoEncryption(),
             )
+        if self.snowflake_authenticator is not None:
+            connect_params["authenticator"] = self.snowflake_authenticator
 
         ctx = snowflake.connector.connect(**connect_params)
         ctx._all_async_queries_finished = lambda: False  # type: ignore