diff --git a/.github/workflows/check_installation.yml b/.github/workflows/check_installation.yml index 4c8d31b80..0695d39ae 100644 --- a/.github/workflows/check_installation.yml +++ b/.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 diff --git a/setup.py b/setup.py index 0b7ab60f4..bc0c06596 100644 --- a/setup.py +++ b/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 diff --git a/src/snowflake/connector/aio/_wif_util.py b/src/snowflake/connector/aio/_wif_util.py index 9474ab744..b1d96fa56 100644 --- a/src/snowflake/connector/aio/_wif_util.py +++ b/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, diff --git a/test/unit/aio/csp_helpers_async.py b/test/unit/aio/csp_helpers_async.py index fab005be6..9095542e8 100644 --- a/test/unit/aio/csp_helpers_async.py +++ b/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,7 +181,9 @@ 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 @@ -174,7 +191,9 @@ def __enter__(self): # 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}