Pczajka fix aio wif (#2678)

Author: sfc-gh-pczajka

.github/workflows/check_installation.yml

@@ -16,8 +16,8 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
 
 jobs:
-  test-installation:
-    name: Test Boto Dependency
+  test-installation-boto:
+    name: Test boto dependency
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
@@ -54,3 +54,44 @@ jobs:
           # Deactivate and clean up
           deactivate
           rm -rf test_no_boto_env
+
+  test-installation-aioboto:
+    name: Test aioboto dependency
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.12
+
+      - name: Test installation with [aio] (should include aioboto)
+        shell: bash
+        run: |
+          python -m venv test_no_boto_env
+          source test_no_boto_env/bin/activate
+
+          python -m pip install '.[aio]'
+
+          # Check that aioboto3 and aiobotocore are installed
+          pip freeze | grep aioboto || exit 1  # aioboto3 and aiobotocore should be installed
+
+          # Deactivate and clean up
+          deactivate
+          rm -rf test_no_boto_env
+
+      - name: Test [aio] installation with SNOWFLAKE_NO_BOTO=1 (should exclude aioboto)
+        shell: bash
+        run: |
+          python -m venv test_no_boto_env
+          source test_no_boto_env/bin/activate
+
+          SNOWFLAKE_NO_BOTO=1 python -m pip install .
+
+          # Check that boto3, botocore, aioboto, aiobotocore are NOT installed
+          pip freeze | grep boto && exit 1
+
+          # Deactivate and clean up
+          deactivate
+          rm -rf test_no_boto_env

setup.py

@@ -186,12 +186,15 @@ class SetDefaultInstallationExtras(egg_info):
 
     def finalize_options(self):
         super().finalize_options()
-
         # if not explicitly excluded, add boto dependencies to install_requires
         if not SNOWFLAKE_NO_BOTO:
             boto_extras = self.distribution.extras_require.get("boto", [])
             self.distribution.install_requires += boto_extras
 
+            if "aio" in self.distribution.extras_require:
+                aioboto_extras = self.distribution.extras_require.get("aioboto", [])
+                self.distribution.extras_require["aio"] += aioboto_extras
+
 
 # Update command classes
 cmd_class["egg_info"] = SetDefaultInstallationExtras

src/snowflake/connector/aio/_wif_util.py

@@ -80,10 +80,7 @@ async def create_aws_attestation(
     If the application isn't running on AWS or no credentials were found, raises an error.
     """
     if not installed_aioboto:
-        raise MissingDependencyError(
-            msg="AWS Workload Identity Federation can't be used because aioboto3 or aiobotocore optional dependency is not installed. Try installing missing dependencies.",
-            errno=ER_WIF_CREDENTIALS_NOT_FOUND,
-        )
+        raise MissingDependencyError("aioboto3 or aiobotocore")
 
     session = await get_aws_session(impersonation_path)
     aws_creds = await session.get_credentials()
@@ -105,7 +102,9 @@ async def create_aws_attestation(
         },
     )
 
-    botocore.auth.SigV4Auth(aws_creds, "sts", region).add_auth(request)
+    # Freeze aiobotocore credentials for use with synchronous botocore signing
+    frozen_creds = await aws_creds.get_frozen_credentials()
+    botocore.auth.SigV4Auth(frozen_creds, "sts", region).add_auth(request)
 
     assertion_dict = {
         "url": request.url,

test/unit/aio/csp_helpers_async.py

@@ -154,6 +154,21 @@ class FakeGceMetadataServiceAsync(FakeMetadataServiceAsync, FakeGceMetadataServi
     pass
 
 
+class AsyncCredentialsWrapper:
+    """Wrapper around boto credentials to make get_frozen_credentials async for testing."""
+
+    def __init__(self, credentials):
+        self._credentials = credentials
+
+    async def get_frozen_credentials(self):
+        """Async version of get_frozen_credentials that returns the wrapped credentials."""
+        return self._credentials
+
+    def __getattr__(self, name):
+        """Delegate all other attributes to the wrapped credentials."""
+        return getattr(self._credentials, name)
+
+
 class FakeAwsEnvironmentAsync(FakeAwsEnvironment):
     """Emulates the AWS environment-specific functions used in async wif_util.py.
 
@@ -166,15 +181,19 @@ async def get_region(self):
         return self.region
 
     async def get_credentials(self):
-        return self.credentials
+        if self.credentials is None:
+            return None
+        return AsyncCredentialsWrapper(self.credentials)
 
     def __enter__(self):
         # First call the parent's __enter__ to get base functionality
         super().__enter__()
 
         # Then add async-specific patches
         async def async_get_credentials():
-            return self.credentials
+            if self.credentials is None:
+                return None
+            return AsyncCredentialsWrapper(self.credentials)
 
         async def async_get_caller_identity():
             return {"Arn": self.arn}