remove duplication in workflow identity

Author: sfc-gh-pczajka

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

@@ -4,39 +4,16 @@
 
 from __future__ import annotations
 
-from enum import Enum, unique
 from typing import Any
 
-from ...auth.by_plugin import AuthType
-from ...network import WORKLOAD_IDENTITY_AUTHENTICATOR
+from ...auth.workload_identity import (
+    AuthByWorkloadIdentity as AuthByWorkloadIdentitySync,
+)
 from .._wif_util import AttestationProvider, create_attestation
 from ._by_plugin import AuthByPlugin as AuthByPluginAsync
 
 
-@unique
-class ApiFederatedAuthenticationType(Enum):
-    """An API-specific enum of the WIF authentication type."""
-
-    AWS = "AWS"
-    AZURE = "AZURE"
-    GCP = "GCP"
-    OIDC = "OIDC"
-
-    @staticmethod
-    def from_attestation(attestation) -> ApiFederatedAuthenticationType:
-        """Maps the internal / driver-specific attestation providers to API authenticator types."""
-        if attestation.provider == AttestationProvider.AWS:
-            return ApiFederatedAuthenticationType.AWS
-        if attestation.provider == AttestationProvider.AZURE:
-            return ApiFederatedAuthenticationType.AZURE
-        if attestation.provider == AttestationProvider.GCP:
-            return ApiFederatedAuthenticationType.GCP
-        if attestation.provider == AttestationProvider.OIDC:
-            return ApiFederatedAuthenticationType.OIDC
-        raise ValueError(f"Unknown attestation provider '{attestation.provider}'")
-
-
-class AuthByWorkloadIdentity(AuthByPluginAsync):
+class AuthByWorkloadIdentity(AuthByWorkloadIdentitySync, AuthByPluginAsync):
     """Plugin to authenticate via workload identity."""
 
     def __init__(
@@ -48,17 +25,16 @@ def __init__(
         **kwargs,
     ) -> None:
         """Initializes an instance with workload identity authentication."""
-        super().__init__(**kwargs)
-        self.provider = provider
-        self.token = token
-        self.entra_resource = entra_resource
-        self.attestation = None
-
-    def type_(self) -> AuthType:
-        return AuthType.WORKLOAD_IDENTITY
+        AuthByWorkloadIdentitySync.__init__(
+            self,
+            provider=provider,
+            token=token,
+            entra_resource=entra_resource,
+            **kwargs,
+        )
 
     async def reset_secrets(self) -> None:
-        self.attestation = None
+        AuthByWorkloadIdentitySync.reset_secrets(self)
 
     async def prepare(self, **kwargs: Any) -> None:
         """Fetch the token using async wif_util."""
@@ -71,19 +47,4 @@ async def reauthenticate(self, **kwargs: Any) -> dict[str, bool]:
         return {"success": False}
 
     async def update_body(self, body: dict[Any, Any]) -> None:
-        body["data"]["AUTHENTICATOR"] = WORKLOAD_IDENTITY_AUTHENTICATOR
-        body["data"]["PROVIDER"] = ApiFederatedAuthenticationType.from_attestation(
-            self.attestation
-        ).value
-        body["data"]["TOKEN"] = self.attestation.credential
-
-    @property
-    def assertion_content(self) -> str:
-        """Returns the CSP provider name and an identifier. Used for logging purposes."""
-        if not self.attestation:
-            return ""
-        properties = self.attestation.user_identifier_components
-        properties["_provider"] = self.attestation.provider.value
-        import json
-
-        return json.dumps(properties, sort_keys=True, separators=(",", ":"))
+        AuthByWorkloadIdentitySync.update_body(self, body)

test/csp_helpers.py

@@ -354,19 +354,7 @@ def __enter__(self):
                 "snowflake.connector.wif_util.get_aws_arn", side_effect=self.get_arn
             )
         )
-        # Also patch the async versions
-        self.patchers.append(
-            mock.patch(
-                "snowflake.connector.aio._wif_util.get_aws_region",
-                side_effect=self.get_region,
-            )
-        )
-        self.patchers.append(
-            mock.patch(
-                "snowflake.connector.aio._wif_util.get_aws_arn",
-                side_effect=self.get_arn,
-            )
-        )
+        # Note: No need to patch async versions anymore since async now imports from sync
         for patcher in self.patchers:
             patcher.__enter__()
         return self