SNOW-1572226: implement all authentication methods (#2064)

Author: sfc-gh-aling

src/snowflake/connector/aio/_connection.py

@@ -27,18 +27,20 @@
 )
 
 from .._query_context_cache import QueryContextCache
-from ..auth import AuthByIdToken
-from ..compat import quote, urlencode
+from ..compat import IS_LINUX, quote, urlencode
 from ..config_manager import CONFIG_MANAGER, _get_default_connection_params
 from ..connection import DEFAULT_CONFIGURATION
 from ..connection import SnowflakeConnection as SnowflakeConnectionSync
+from ..connection import _get_private_bytes_from_file
 from ..connection_diagnostic import ConnectionDiagnostic
 from ..constants import (
     ENV_VAR_PARTNER,
     PARAMETER_AUTOCOMMIT,
     PARAMETER_CLIENT_PREFETCH_THREADS,
+    PARAMETER_CLIENT_REQUEST_MFA_TOKEN,
     PARAMETER_CLIENT_SESSION_KEEP_ALIVE,
     PARAMETER_CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY,
+    PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL,
     PARAMETER_CLIENT_TELEMETRY_ENABLED,
     PARAMETER_CLIENT_VALIDATE_DEFAULT_PARAMETERS,
     PARAMETER_ENABLE_STAGE_S3_PRIVATELINK_FOR_US_EAST_1,
@@ -53,15 +55,34 @@
     ER_FAILED_TO_CONNECT_TO_DB,
     ER_INVALID_VALUE,
 )
-from ..network import DEFAULT_AUTHENTICATOR, REQUEST_ID, ReauthenticationRequest
+from ..network import (
+    DEFAULT_AUTHENTICATOR,
+    EXTERNAL_BROWSER_AUTHENTICATOR,
+    KEY_PAIR_AUTHENTICATOR,
+    OAUTH_AUTHENTICATOR,
+    REQUEST_ID,
+    USR_PWD_MFA_AUTHENTICATOR,
+    ReauthenticationRequest,
+)
 from ..sqlstate import SQLSTATE_CONNECTION_NOT_EXISTS, SQLSTATE_FEATURE_NOT_SUPPORTED
 from ..telemetry import TelemetryData, TelemetryField
 from ..time_util import get_time_millis
 from ..util_text import split_statements
 from ._cursor import SnowflakeCursor
 from ._network import SnowflakeRestful
 from ._time_util import HeartBeatTimer
-from .auth import Auth, AuthByDefault, AuthByPlugin
+from .auth import (
+    FIRST_PARTY_AUTHENTICATORS,
+    Auth,
+    AuthByDefault,
+    AuthByIdToken,
+    AuthByKeyPair,
+    AuthByOAuth,
+    AuthByOkta,
+    AuthByPlugin,
+    AuthByUsrPwdMfa,
+    AuthByWebBrowser,
+)
 
 logger = getLogger(__name__)
 
@@ -196,7 +217,6 @@ async def __open_connection(self):
             heartbeat_ret = await auth._rest._heartbeat()
             logger.debug(heartbeat_ret)
             if not heartbeat_ret or not heartbeat_ret.get("success"):
-                # TODO: errorhandler could be async?
                 Error.errorhandler_wrapper(
                     self,
                     None,
@@ -211,20 +231,94 @@ async def __open_connection(self):
 
         else:
             if self.auth_class is not None:
-                raise NotImplementedError(
-                    "asyncio support for auth_class is not supported"
-                )
+                if type(
+                    self.auth_class
+                ) not in FIRST_PARTY_AUTHENTICATORS and not issubclass(
+                    type(self.auth_class), AuthByKeyPair
+                ):
+                    raise TypeError("auth_class must be a child class of AuthByKeyPair")
+                self.auth_class = self.auth_class
             elif self._authenticator == DEFAULT_AUTHENTICATOR:
                 self.auth_class = AuthByDefault(
                     password=self._password,
                     timeout=self.login_timeout,
                     backoff_generator=self._backoff_generator,
                 )
+            elif self._authenticator == EXTERNAL_BROWSER_AUTHENTICATOR:
+                self._session_parameters[
+                    PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL
+                ] = (self._client_store_temporary_credential if IS_LINUX else True)
+                auth.read_temporary_credentials(
+                    self.host,
+                    self.user,
+                    self._session_parameters,
+                )
+                # Depending on whether self._rest.id_token is available we do different
+                #  auth_instance
+                if self._rest.id_token is None:
+                    self.auth_class = AuthByWebBrowser(
+                        application=self.application,
+                        protocol=self._protocol,
+                        host=self.host,
+                        port=self.port,
+                        timeout=self.login_timeout,
+                        backoff_generator=self._backoff_generator,
+                    )
+                else:
+                    self.auth_class = AuthByIdToken(
+                        id_token=self._rest.id_token,
+                        application=self.application,
+                        protocol=self._protocol,
+                        host=self.host,
+                        port=self.port,
+                        timeout=self.login_timeout,
+                        backoff_generator=self._backoff_generator,
+                    )
+
+            elif self._authenticator == KEY_PAIR_AUTHENTICATOR:
+                private_key = self._private_key
+
+                if self._private_key_file:
+                    private_key = _get_private_bytes_from_file(
+                        self._private_key_file,
+                        self._private_key_file_pwd,
+                    )
+
+                self.auth_class = AuthByKeyPair(
+                    private_key=private_key,
+                    timeout=self.login_timeout,
+                    backoff_generator=self._backoff_generator,
+                )
+            elif self._authenticator == OAUTH_AUTHENTICATOR:
+                self.auth_class = AuthByOAuth(
+                    oauth_token=self._token,
+                    timeout=self.login_timeout,
+                    backoff_generator=self._backoff_generator,
+                )
+            elif self._authenticator == USR_PWD_MFA_AUTHENTICATOR:
+                self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN] = (
+                    self._client_request_mfa_token if IS_LINUX else True
+                )
+                if self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN]:
+                    auth.read_temporary_credentials(
+                        self.host,
+                        self.user,
+                        self._session_parameters,
+                    )
+                self.auth_class = AuthByUsrPwdMfa(
+                    password=self._password,
+                    mfa_token=self.rest.mfa_token,
+                    timeout=self.login_timeout,
+                    backoff_generator=self._backoff_generator,
+                )
             else:
-                raise NotImplementedError(
-                    f"asyncio support for authenticator is not supported {self._authenticator}"
+                # okta URL, e.g., https://<account>.okta.com/
+                self.auth_class = AuthByOkta(
+                    application=self.application,
+                    timeout=self.login_timeout,
+                    backoff_generator=self._backoff_generator,
                 )
-            # TODO: asyncio support for other authenticators
+
             await self.authenticate_with_retry(self.auth_class)
 
             self._password = None  # ensure password won't persist

src/snowflake/connector/aio/auth/__init__.py

@@ -4,12 +4,38 @@
 
 from __future__ import annotations
 
+from ...auth.by_plugin import AuthType
 from ._auth import Auth
 from ._by_plugin import AuthByPlugin
 from ._default import AuthByDefault
+from ._idtoken import AuthByIdToken
+from ._keypair import AuthByKeyPair
+from ._oauth import AuthByOAuth
+from ._okta import AuthByOkta
+from ._usrpwdmfa import AuthByUsrPwdMfa
+from ._webbrowser import AuthByWebBrowser
+
+FIRST_PARTY_AUTHENTICATORS = frozenset(
+    (
+        AuthByDefault,
+        AuthByKeyPair,
+        AuthByOAuth,
+        AuthByOkta,
+        AuthByUsrPwdMfa,
+        AuthByWebBrowser,
+        AuthByIdToken,
+    )
+)
 
 __all__ = [
-    AuthByDefault,
-    Auth,
-    AuthByPlugin,
+    "AuthByPlugin",
+    "AuthByDefault",
+    "AuthByKeyPair",
+    "AuthByOAuth",
+    "AuthByOkta",
+    "AuthByUsrPwdMfa",
+    "AuthByWebBrowser",
+    "Auth",
+    "AuthType",
+    "FIRST_PARTY_AUTHENTICATORS",
 ]

src/snowflake/connector/aio/auth/_auth.py

@@ -4,14 +4,16 @@
 
 from __future__ import annotations
 
+import asyncio
 import copy
 import json
 import logging
 import uuid
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Callable
 
 from ...auth import Auth as AuthSync
-from ...auth._auth import ID_TOKEN, delete_temporary_credential
+from ...auth._auth import ID_TOKEN, MFA_TOKEN, delete_temporary_credential
 from ...compat import urlencode
 from ...constants import (
     HTTP_HEADER_ACCEPT,
@@ -62,9 +64,10 @@ async def authenticate(
         timeout: int | None = None,
     ) -> dict[str, str | int | bool]:
         if mfa_callback or password_callback:
-            # TODO: what's the usage of callback here and whether callback should be async?
+            # check SNOW-1707210 for mfa_callback and password_callback support
             raise NotImplementedError(
-                "mfa_callback or password_callback not supported for asyncio"
+                "mfa_callback or password_callback is not supported in asyncio connector, please open a feature"
+                " request issue in github: https://github.com/snowflakedb/snowflake-connector-python/issues/new/choose"
             )
         logger.debug("authenticate")
 
@@ -148,7 +151,6 @@ async def authenticate(
                 json.dumps(body),
                 socket_timeout=auth_instance._socket_timeout,
             )
-        # TODO: encapsulate error handling logic to be shared between sync and async
         except ForbiddenError as err:
             # HTTP 403
             raise err.__class__(
@@ -181,7 +183,65 @@ async def authenticate(
             "EXT_AUTHN_DUO_ALL",
             "EXT_AUTHN_DUO_PUSH_N_PASSCODE",
         ):
-            raise NotImplementedError("asyncio MFA not supported")
+            body["inFlightCtx"] = ret["data"].get("inFlightCtx")
+            body["data"]["EXT_AUTHN_DUO_METHOD"] = "push"
+            self.ret = {"message": "Timeout", "data": {}}
+
+            async def post_request_wrapper(self, url, headers, body) -> None:
+                # get the MFA response
+                self.ret = await self._rest._post_request(
+                    url,
+                    headers,
+                    body,
+                    socket_timeout=auth_instance._socket_timeout,
+                )
+
+            # send new request to wait until MFA is approved
+            try:
+                await asyncio.wait_for(
+                    post_request_wrapper(self, url, headers, json.dumps(body)),
+                    timeout=timeout,
+                )
+            except asyncio.TimeoutError:
+                logger.debug("get the MFA response timed out")
+
+            ret = self.ret
+            if (
+                ret
+                and ret["data"]
+                and ret["data"].get("nextAction") == "EXT_AUTHN_SUCCESS"
+            ):
+                body = copy.deepcopy(body_template)
+                body["inFlightCtx"] = ret["data"].get("inFlightCtx")
+                # final request to get tokens
+                ret = await self._rest._post_request(
+                    url,
+                    headers,
+                    json.dumps(body),
+                    socket_timeout=auth_instance._socket_timeout,
+                )
+            elif not ret or not ret["data"] or not ret["data"].get("token"):
+                # not token is returned.
+                Error.errorhandler_wrapper(
+                    self._rest._connection,
+                    None,
+                    DatabaseError,
+                    {
+                        "msg": (
+                            "Failed to connect to DB. MFA "
+                            "authentication failed: {"
+                            "host}:{port}. {message}"
+                        ).format(
+                            host=self._rest._host,
+                            port=self._rest._port,
+                            message=ret["message"],
+                        ),
+                        "errno": ER_FAILED_TO_CONNECT_TO_DB,
+                        "sqlstate": SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
+                    },
+                )
+                return session_parameters  # required for unit test
+
         elif ret["data"] and ret["data"].get("nextAction") == "PWD_CHANGE":
             if callable(password_callback):
                 body = copy.deepcopy(body_template)
@@ -216,23 +276,20 @@ async def authenticate(
                         sqlstate=SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
                     )
                 )
-            # TODO: error handling for AuthByKeyPairAsync and AuthByUsrPwdMfaAsync
-            # from . import AuthByKeyPair
-            #
-            # if isinstance(auth_instance, AuthByKeyPair):
-            #     logger.debug(
-            #         "JWT Token authentication failed. "
-            #         "Token expires at: %s. "
-            #         "Current Time: %s",
-            #         str(auth_instance._jwt_token_exp),
-            #         str(datetime.now(timezone.utc).replace(tzinfo=None)),
-            #     )
-            # from . import AuthByUsrPwdMfa
-            #
-            # if isinstance(auth_instance, AuthByUsrPwdMfa):
-            #     delete_temporary_credential(self._rest._host, user, MFA_TOKEN)
-            # TODO: can errorhandler of a connection be async? should we support both sync and async
-            #  users could perform async ops in the error handling
+            from . import AuthByKeyPair
+
+            if isinstance(auth_instance, AuthByKeyPair):
+                logger.debug(
+                    "JWT Token authentication failed. "
+                    "Token expires at: %s. "
+                    "Current Time: %s",
+                    str(auth_instance._jwt_token_exp),
+                    str(datetime.now(timezone.utc).replace(tzinfo=None)),
+                )
+            from . import AuthByUsrPwdMfa
+
+            if isinstance(auth_instance, AuthByUsrPwdMfa):
+                delete_temporary_credential(self._rest._host, user, MFA_TOKEN)
             Error.errorhandler_wrapper(
                 self._rest._connection,
                 None,

src/snowflake/connector/aio/auth/_by_plugin.py

@@ -7,17 +7,28 @@
 import asyncio
 import logging
 from abc import abstractmethod
-from typing import Any
+from typing import TYPE_CHECKING, Any, Iterator
 
-from ... import DatabaseError, Error, OperationalError, SnowflakeConnection
+from ... import DatabaseError, Error, OperationalError
 from ...auth import AuthByPlugin as AuthByPluginSync
 from ...errorcode import ER_FAILED_TO_CONNECT_TO_DB
 from ...sqlstate import SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED
 
+if TYPE_CHECKING:
+    from .. import SnowflakeConnection
+
 logger = logging.getLogger(__name__)
 
 
 class AuthByPlugin(AuthByPluginSync):
+    def __init__(
+        self,
+        timeout: int | None = None,
+        backoff_generator: Iterator | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(timeout, backoff_generator, **kwargs)
+
     @abstractmethod
     async def prepare(
         self,