[Async] Apply #2203 to async code

Author: sfc-gh-pczajka

src/snowflake/connector/aio/_connection.py

@@ -35,6 +35,7 @@
 from ..connection import _get_private_bytes_from_file
 from ..constants import (
     _CONNECTIVITY_ERR_MSG,
+    ENV_VAR_EXPERIMENTAL_AUTHENTICATION,
     ENV_VAR_PARTNER,
     PARAMETER_AUTOCOMMIT,
     PARAMETER_CLIENT_PREFETCH_THREADS,
@@ -55,6 +56,7 @@
     ER_CONNECTION_IS_CLOSED,
     ER_FAILED_TO_CONNECT_TO_DB,
     ER_INVALID_VALUE,
+    ER_INVALID_WIF_SETTINGS,
 )
 from ..network import (
     DEFAULT_AUTHENTICATOR,
@@ -64,12 +66,14 @@
     PROGRAMMATIC_ACCESS_TOKEN,
     REQUEST_ID,
     USR_PWD_MFA_AUTHENTICATOR,
+    WORKLOAD_IDENTITY_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 ..wif_util import AttestationProvider
 from ._cursor import SnowflakeCursor
 from ._description import CLIENT_NAME
 from ._network import SnowflakeRestful
@@ -87,6 +91,7 @@
     AuthByPlugin,
     AuthByUsrPwdMfa,
     AuthByWebBrowser,
+    AuthByWorkloadIdentity,
 )
 
 logger = getLogger(__name__)
@@ -320,6 +325,29 @@ async def __open_connection(self):
                     timeout=self.login_timeout,
                     backoff_generator=self._backoff_generator,
                 )
+            elif self._authenticator == WORKLOAD_IDENTITY_AUTHENTICATOR:
+                if ENV_VAR_EXPERIMENTAL_AUTHENTICATION not in os.environ:
+                    Error.errorhandler_wrapper(
+                        self,
+                        None,
+                        ProgrammingError,
+                        {
+                            "msg": f"Please set the '{ENV_VAR_EXPERIMENTAL_AUTHENTICATION}' environment variable to use the '{WORKLOAD_IDENTITY_AUTHENTICATOR}' authenticator.",
+                            "errno": ER_INVALID_WIF_SETTINGS,
+                        },
+                    )
+                # Standardize the provider enum.
+                if self._workload_identity_provider and isinstance(
+                    self._workload_identity_provider, str
+                ):
+                    self._workload_identity_provider = AttestationProvider.from_string(
+                        self._workload_identity_provider
+                    )
+                self.auth_class = AuthByWorkloadIdentity(
+                    provider=self._workload_identity_provider,
+                    token=self._token,
+                    entra_resource=self._workload_identity_entra_resource,
+                )
             else:
                 # okta URL, e.g., https://<account>.okta.com/
                 self.auth_class = AuthByOkta(

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

@@ -16,6 +16,7 @@
 from ._pat import AuthByPAT
 from ._usrpwdmfa import AuthByUsrPwdMfa
 from ._webbrowser import AuthByWebBrowser
+from ._workload_identity import AuthByWorkloadIdentity
 
 FIRST_PARTY_AUTHENTICATORS = frozenset(
     (
@@ -27,6 +28,7 @@
         AuthByWebBrowser,
         AuthByIdToken,
         AuthByPAT,
+        AuthByWorkloadIdentity,
         AuthNoAuth,
     )
 )
@@ -40,6 +42,7 @@
     "AuthByOkta",
     "AuthByUsrPwdMfa",
     "AuthByWebBrowser",
+    "AuthByWorkloadIdentity",
     "AuthNoAuth",
     "Auth",
     "AuthType",

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

@@ -0,0 +1,43 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+from typing import Any
+
+from ...auth.workload_identity import (
+    AuthByWorkloadIdentity as AuthByWorkloadIdentitySync,
+)
+from ._by_plugin import AuthByPlugin as AuthByPluginAsync
+
+
+class AuthByWorkloadIdentity(AuthByPluginAsync, AuthByWorkloadIdentitySync):
+    def __init__(
+        self,
+        *,
+        provider=None,
+        token: str | None = None,
+        entra_resource: str | None = None,
+        **kwargs,
+    ) -> None:
+        """Initializes an instance with workload identity authentication."""
+        AuthByWorkloadIdentitySync.__init__(
+            self,
+            provider=provider,
+            token=token,
+            entra_resource=entra_resource,
+            **kwargs,
+        )
+
+    async def reset_secrets(self) -> None:
+        AuthByWorkloadIdentitySync.reset_secrets(self)
+
+    async def prepare(self, **kwargs: Any) -> None:
+        AuthByWorkloadIdentitySync.prepare(self, **kwargs)
+
+    async def reauthenticate(self, **kwargs: Any) -> dict[str, bool]:
+        return AuthByWorkloadIdentitySync.reauthenticate(self, **kwargs)
+
+    async def update_body(self, body: dict[Any, Any]) -> None:
+        AuthByWorkloadIdentitySync.update_body(self, body)