Link sync implementation of Oauth to async code

Author: sfc-gh-pczajka

src/snowflake/connector/aio/_connection.py

@@ -32,6 +32,7 @@
 from ..connection import _get_private_bytes_from_file
 from ..constants import (
     _CONNECTIVITY_ERR_MSG,
+    _OAUTH_DEFAULT_SCOPE,
     ENV_VAR_EXPERIMENTAL_AUTHENTICATION,
     ENV_VAR_PARTNER,
     PARAMETER_AUTOCOMMIT,
@@ -51,15 +52,19 @@
 from ..description import PLATFORM, PYTHON_VERSION, SNOWFLAKE_CONNECTOR_VERSION
 from ..errorcode import (
     ER_CONNECTION_IS_CLOSED,
+    ER_EXPERIMENTAL_AUTHENTICATION_NOT_SUPPORTED,
     ER_FAILED_TO_CONNECT_TO_DB,
     ER_INVALID_VALUE,
     ER_INVALID_WIF_SETTINGS,
+    ER_NO_CLIENT_ID,
 )
 from ..network import (
     DEFAULT_AUTHENTICATOR,
     EXTERNAL_BROWSER_AUTHENTICATOR,
     KEY_PAIR_AUTHENTICATOR,
     OAUTH_AUTHENTICATOR,
+    OAUTH_AUTHORIZATION_CODE,
+    OAUTH_CLIENT_CREDENTIALS,
     PROGRAMMATIC_ACCESS_TOKEN,
     REQUEST_ID,
     USR_PWD_MFA_AUTHENTICATOR,
@@ -84,6 +89,8 @@
     AuthByIdToken,
     AuthByKeyPair,
     AuthByOAuth,
+    AuthByOauthCode,
+    AuthByOauthCredentials,
     AuthByOkta,
     AuthByPAT,
     AuthByPlugin,
@@ -307,6 +314,56 @@ async def __open_connection(self):
                     timeout=self.login_timeout,
                     backoff_generator=self._backoff_generator,
                 )
+            elif self._authenticator == OAUTH_AUTHORIZATION_CODE:
+                self._check_experimental_authentication_flag()
+                self._check_oauth_required_parameters()
+                features = self.oauth_security_features
+                if self._role and (self._oauth_scope == ""):
+                    # if role is known then let's inject it into scope
+                    self._oauth_scope = _OAUTH_DEFAULT_SCOPE.format(role=self._role)
+                self.auth_class = AuthByOauthCode(
+                    application=self.application,
+                    client_id=self._oauth_client_id,
+                    client_secret=self._oauth_client_secret,
+                    authentication_url=self._oauth_authorization_url.format(
+                        host=self.host, port=self.port
+                    ),
+                    token_request_url=self._oauth_token_request_url.format(
+                        host=self.host, port=self.port
+                    ),
+                    redirect_uri=self._oauth_redirect_uri,
+                    scope=self._oauth_scope,
+                    pkce_enabled=features.pkce_enabled,
+                    token_cache=(
+                        auth.get_token_cache()
+                        if self._client_store_temporary_credential
+                        else None
+                    ),
+                    refresh_token_enabled=features.refresh_token_enabled,
+                    external_browser_timeout=self._external_browser_timeout,
+                )
+            elif self._authenticator == OAUTH_CLIENT_CREDENTIALS:
+                self._check_experimental_authentication_flag()
+                self._check_oauth_required_parameters()
+                features = self.oauth_security_features
+                if self._role and (self._oauth_scope == ""):
+                    # if role is known then let's inject it into scope
+                    self._oauth_scope = _OAUTH_DEFAULT_SCOPE.format(role=self._role)
+                self.auth_class = AuthByOauthCredentials(
+                    application=self.application,
+                    client_id=self._oauth_client_id,
+                    client_secret=self._oauth_client_secret,
+                    token_request_url=self._oauth_token_request_url.format(
+                        host=self.host, port=self.port
+                    ),
+                    scope=self._oauth_scope,
+                    token_cache=(
+                        auth.get_token_cache()
+                        if self._client_store_temporary_credential
+                        else None
+                    ),
+                    refresh_token_enabled=features.refresh_token_enabled,
+                )
             elif self._authenticator == PROGRAMMATIC_ACCESS_TOKEN:
                 self.auth_class = AuthByPAT(self._token)
             elif self._authenticator == USR_PWD_MFA_AUTHENTICATOR:
@@ -1052,3 +1109,37 @@ async def is_valid(self) -> bool:
         except Exception as e:
             logger.debug("session could not be validated due to exception: %s", e)
             return False
+
+    def _check_experimental_authentication_flag(self) -> None:
+        if os.getenv(ENV_VAR_EXPERIMENTAL_AUTHENTICATION, "false").lower() != "true":
+            Error.errorhandler_wrapper(
+                self,
+                None,
+                ProgrammingError,
+                {
+                    "msg": f"Please set the '{ENV_VAR_EXPERIMENTAL_AUTHENTICATION}' environment variable true to use the '{self._authenticator}' authenticator.",
+                    "errno": ER_EXPERIMENTAL_AUTHENTICATION_NOT_SUPPORTED,
+                },
+            )
+
+    def _check_oauth_required_parameters(self) -> None:
+        if self._oauth_client_id is None:
+            Error.errorhandler_wrapper(
+                self,
+                None,
+                ProgrammingError,
+                {
+                    "msg": "Oauth code flow requirement 'client_id' is empty",
+                    "errno": ER_NO_CLIENT_ID,
+                },
+            )
+        if self._oauth_client_secret is None:
+            Error.errorhandler_wrapper(
+                self,
+                None,
+                ProgrammingError,
+                {
+                    "msg": "Oauth code flow requirement 'client_secret' is empty",
+                    "errno": ER_NO_CLIENT_ID,
+                },
+            )

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

@@ -8,6 +8,8 @@
 from ._keypair import AuthByKeyPair
 from ._no_auth import AuthNoAuth
 from ._oauth import AuthByOAuth
+from ._oauth_code import AuthByOauthCode
+from ._oauth_credentials import AuthByOauthCredentials
 from ._okta import AuthByOkta
 from ._pat import AuthByPAT
 from ._usrpwdmfa import AuthByUsrPwdMfa
@@ -19,6 +21,8 @@
         AuthByDefault,
         AuthByKeyPair,
         AuthByOAuth,
+        AuthByOauthCode,
+        AuthByOauthCredentials,
         AuthByOkta,
         AuthByUsrPwdMfa,
         AuthByWebBrowser,
@@ -35,6 +39,8 @@
     "AuthByKeyPair",
     "AuthByPAT",
     "AuthByOAuth",
+    "AuthByOauthCode",
+    "AuthByOauthCredentials",
     "AuthByOkta",
     "AuthByUsrPwdMfa",
     "AuthByWebBrowser",

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

@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from ...auth.oauth_code import AuthByOauthCode as AuthByOauthCodeSync
+from ...token_cache import TokenCache
+from ._by_plugin import AuthByPlugin as AuthByPluginAsync
+
+logger = logging.getLogger(__name__)
+
+
+class AuthByOauthCode(AuthByPluginAsync, AuthByOauthCodeSync):
+    """Async version of OAuth authorization code authenticator."""
+
+    def __init__(
+        self,
+        application: str,
+        client_id: str,
+        client_secret: str,
+        authentication_url: str,
+        token_request_url: str,
+        redirect_uri: str,
+        scope: str,
+        pkce_enabled: bool = True,
+        token_cache: TokenCache | None = None,
+        refresh_token_enabled: bool = False,
+        external_browser_timeout: int | None = None,
+        **kwargs,
+    ) -> None:
+        """Initializes an instance with OAuth authorization code parameters."""
+        logger.debug(
+            "OAuth authentication is not supported in async version - falling back to sync implementation"
+        )
+        AuthByOauthCodeSync.__init__(
+            self,
+            application=application,
+            client_id=client_id,
+            client_secret=client_secret,
+            authentication_url=authentication_url,
+            token_request_url=token_request_url,
+            redirect_uri=redirect_uri,
+            scope=scope,
+            pkce_enabled=pkce_enabled,
+            token_cache=token_cache,
+            refresh_token_enabled=refresh_token_enabled,
+            external_browser_timeout=external_browser_timeout,
+            **kwargs,
+        )
+
+    async def reset_secrets(self) -> None:
+        AuthByOauthCodeSync.reset_secrets(self)
+
+    async def prepare(self, **kwargs: Any) -> None:
+        AuthByOauthCodeSync.prepare(self, **kwargs)
+
+    async def reauthenticate(self, **kwargs: Any) -> dict[str, bool]:
+        return AuthByOauthCodeSync.reauthenticate(self, **kwargs)
+
+    async def update_body(self, body: dict[Any, Any]) -> None:
+        AuthByOauthCodeSync.update_body(self, body)

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

@@ -0,0 +1,57 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from ...auth.oauth_credentials import (
+    AuthByOauthCredentials as AuthByOauthCredentialsSync,
+)
+from ...token_cache import TokenCache
+from ._by_plugin import AuthByPlugin as AuthByPluginAsync
+
+logger = logging.getLogger(__name__)
+
+
+class AuthByOauthCredentials(AuthByPluginAsync, AuthByOauthCredentialsSync):
+    """Async version of OAuth client credentials authenticator."""
+
+    def __init__(
+        self,
+        application: str,
+        client_id: str,
+        client_secret: str,
+        token_request_url: str,
+        scope: str,
+        token_cache: TokenCache | None = None,
+        refresh_token_enabled: bool = False,
+        **kwargs,
+    ) -> None:
+        """Initializes an instance with OAuth client credentials parameters."""
+        logger.debug(
+            "OAuth authentication is not supported in async version - falling back to sync implementation"
+        )
+        AuthByOauthCredentialsSync.__init__(
+            self,
+            application=application,
+            client_id=client_id,
+            client_secret=client_secret,
+            token_request_url=token_request_url,
+            scope=scope,
+            token_cache=token_cache,
+            refresh_token_enabled=refresh_token_enabled,
+            **kwargs,
+        )
+
+    async def reset_secrets(self) -> None:
+        AuthByOauthCredentialsSync.reset_secrets(self)
+
+    async def prepare(self, **kwargs: Any) -> None:
+        AuthByOauthCredentialsSync.prepare(self, **kwargs)
+
+    async def reauthenticate(self, **kwargs: Any) -> dict[str, bool]:
+        return AuthByOauthCredentialsSync.reauthenticate(self, **kwargs)
+
+    async def update_body(self, body: dict[Any, Any]) -> None:
+        AuthByOauthCredentialsSync.update_body(self, body)