use aioboto3

Author: sfc-gh-pczajka

setup.cfg

@@ -100,3 +100,4 @@ secure-local-storage =
     keyring>=23.1.0,<26.0.0
 aio =
     aiohttp
+    aioboto3

src/snowflake/connector/aio/_wif_util.py

@@ -8,17 +8,28 @@
 import json
 import logging
 import os
+from base64 import b64encode
 
 import aiohttp
 
+try:
+    import aioboto3
+    from botocore.auth import SigV4Auth
+    from botocore.awsrequest import AWSRequest
+    from botocore.utils import InstanceMetadataRegionFetcher
+except ImportError:
+    aioboto3 = None
+    SigV4Auth = None
+    AWSRequest = None
+    InstanceMetadataRegionFetcher = None
+
 from ..errorcode import ER_WIF_CREDENTIALS_NOT_FOUND
 from ..errors import ProgrammingError
 from ..wif_util import (
     DEFAULT_ENTRA_SNOWFLAKE_RESOURCE,
     SNOWFLAKE_AUDIENCE,
     AttestationProvider,
     WorkloadIdentityAttestation,
-    create_aws_attestation,
     create_oidc_attestation,
     extract_iss_and_sub_without_signature_verification,
 )
@@ -49,6 +60,91 @@ async def try_metadata_service_call(
         return None
 
 
+async def get_aws_region() -> str | None:
+    """Get the current AWS workload's region, if any."""
+    # Use sync implementation which has proper mocking support
+    from ..wif_util import get_aws_region as sync_get_aws_region
+
+    return sync_get_aws_region()
+
+
+async def get_aws_arn() -> str | None:
+    """Get the current AWS workload's ARN, if any."""
+    if aioboto3 is None:
+        logger.debug("aioboto3 not available, falling back to sync implementation")
+        from ..wif_util import get_aws_arn as sync_get_aws_arn
+
+        return sync_get_aws_arn()
+
+    try:
+        session = aioboto3.Session()
+        async with session.client("sts") as client:
+            caller_identity = await client.get_caller_identity()
+            if not caller_identity or "Arn" not in caller_identity:
+                return None
+            return caller_identity["Arn"]
+    except Exception:
+        logger.debug("Failed to get AWS ARN", exc_info=True)
+        return None
+
+
+async def create_aws_attestation() -> WorkloadIdentityAttestation | None:
+    """Tries to create a workload identity attestation for AWS.
+
+    If the application isn't running on AWS or no credentials were found, returns None.
+    """
+    if aioboto3 is None:
+        logger.debug("aioboto3 not available, falling back to sync implementation")
+        from ..wif_util import create_aws_attestation as sync_create_aws_attestation
+
+        return sync_create_aws_attestation()
+
+    try:
+        # Get credentials using aioboto3
+        session = aioboto3.Session()
+        aws_creds = await session.get_credentials()  # This IS async in aioboto3
+        if not aws_creds:
+            logger.debug("No AWS credentials were found.")
+            return None
+
+        region = await get_aws_region()
+        if not region:
+            logger.debug("No AWS region was found.")
+            return None
+
+        arn = await get_aws_arn()
+        if not arn:
+            logger.debug("No AWS caller identity was found.")
+            return None
+
+        sts_hostname = f"sts.{region}.amazonaws.com"
+        request = AWSRequest(
+            method="POST",
+            url=f"https://{sts_hostname}/?Action=GetCallerIdentity&Version=2011-06-15",
+            headers={
+                "Host": sts_hostname,
+                "X-Snowflake-Audience": SNOWFLAKE_AUDIENCE,
+            },
+        )
+
+        SigV4Auth(aws_creds, "sts", region).add_auth(request)
+
+        assertion_dict = {
+            "url": request.url,
+            "method": request.method,
+            "headers": dict(request.headers.items()),
+        }
+        credential = b64encode(json.dumps(assertion_dict).encode("utf-8")).decode(
+            "utf-8"
+        )
+        return WorkloadIdentityAttestation(
+            AttestationProvider.AWS, credential, {"arn": arn}
+        )
+    except Exception:
+        logger.debug("Failed to create AWS attestation", exc_info=True)
+        return None
+
+
 async def create_gcp_attestation() -> WorkloadIdentityAttestation | None:
     """Tries to create a workload identity attestation for GCP.
 
@@ -157,7 +253,7 @@ async def create_autodetect_attestation(
     if attestation:
         return attestation
 
-    attestation = create_aws_attestation()
+    attestation = await create_aws_attestation()
     if attestation:
         return attestation
 
@@ -188,7 +284,7 @@ async def create_attestation(
 
     attestation: WorkloadIdentityAttestation = None
     if provider == AttestationProvider.AWS:
-        attestation = create_aws_attestation()
+        attestation = await create_aws_attestation()
     elif provider == AttestationProvider.AZURE:
         attestation = await create_azure_attestation(entra_resource)
     elif provider == AttestationProvider.GCP:

test/csp_helpers.py

@@ -332,6 +332,8 @@ def sign_request(self, request: AWSRequest):
     def __enter__(self):
         # Patch the relevant functions to do what we want.
         self.patchers = []
+
+        # Patch sync boto3 calls
         self.patchers.append(
             mock.patch(
                 "boto3.session.Session.get_credentials",
@@ -354,7 +356,44 @@ def __enter__(self):
                 "snowflake.connector.wif_util.get_aws_arn", side_effect=self.get_arn
             )
         )
-        # Note: No need to patch async versions anymore since async now imports from sync
+
+        # Patch async aioboto3 calls (for when aioboto3 is used directly)
+        async def async_get_credentials():
+            return self.credentials
+
+        async def async_get_caller_identity():
+            return {"Arn": self.arn}
+
+        # Mock aioboto3.Session.get_credentials (IS async)
+        self.patchers.append(
+            mock.patch(
+                "snowflake.connector.aio._wif_util.aioboto3.Session.get_credentials",
+                side_effect=async_get_credentials,
+            )
+        )
+
+        # Mock the async STS client for direct aioboto3 usage
+        class MockStsClient:
+            async def __aenter__(self):
+                return self
+
+            async def __aexit__(self, exc_type, exc_val, exc_tb):
+                pass
+
+            async def get_caller_identity(self):
+                return await async_get_caller_identity()
+
+        def mock_session_client(service_name):
+            if service_name == "sts":
+                return MockStsClient()
+            return None
+
+        self.patchers.append(
+            mock.patch(
+                "snowflake.connector.aio._wif_util.aioboto3.Session.client",
+                side_effect=mock_session_client,
+            )
+        )
         for patcher in self.patchers:
             patcher.__enter__()
         return self