[async] WIF impersonation for AWS (#2517)

Author: sfc-gh-pczajka

src/snowflake/connector/aio/_connection.py

@@ -420,14 +420,18 @@ async def __open_connection(self):
                     )
                 if (
                     self._workload_identity_impersonation_path
-                    and self._workload_identity_provider != AttestationProvider.GCP
+                    and self._workload_identity_provider
+                    not in (
+                        AttestationProvider.GCP,
+                        AttestationProvider.AWS,
+                    )
                 ):
                     Error.errorhandler_wrapper(
                         self,
                         None,
                         ProgrammingError,
                         {
-                            "msg": "workload_identity_impersonation_path is currently only supported for GCP.",
+                            "msg": "workload_identity_impersonation_path is currently only supported for GCP and AWS.",
                             "errno": ER_INVALID_WIF_SETTINGS,
                         },
                     )

src/snowflake/connector/aio/_wif_util.py

@@ -45,12 +45,37 @@ async def get_aws_region() -> str:
     return region
 
 
-async def create_aws_attestation() -> WorkloadIdentityAttestation:
+async def get_aws_session(impersonation_path: list[str] | None = None):
+    """Creates an aioboto3 session with the appropriate credentials.
+
+    If impersonation_path is provided, this uses the role at the end of the path. Otherwise, this uses the role attached to the current workload.
+    """
+    session = aioboto3.Session()
+
+    impersonation_path = impersonation_path or []
+    for arn in impersonation_path:
+        async with session.client("sts") as sts_client:
+            response = await sts_client.assume_role(
+                RoleArn=arn, RoleSessionName="identity-federation-session"
+            )
+        creds = response["Credentials"]
+        session = aioboto3.Session(
+            aws_access_key_id=creds["AccessKeyId"],
+            aws_secret_access_key=creds["SecretAccessKey"],
+            aws_session_token=creds["SessionToken"],
+        )
+    return session
+
+
+async def create_aws_attestation(
+    impersonation_path: list[str] | None = None,
+) -> WorkloadIdentityAttestation:
     """Tries to create a workload identity attestation for AWS.
 
     If the application isn't running on AWS or no credentials were found, raises an error.
     """
-    session = aioboto3.Session()
+    session = await get_aws_session(impersonation_path)
+
     aws_creds = await session.get_credentials()
     if not aws_creds:
         raise ProgrammingError(
@@ -274,7 +299,7 @@ async def create_attestation(
     )
 
     if provider == AttestationProvider.AWS:
-        return await create_aws_attestation()
+        return await create_aws_attestation(impersonation_path)
     elif provider == AttestationProvider.AZURE:
         return await create_azure_attestation(entra_resource, session_manager)
     elif provider == AttestationProvider.GCP:

test/unit/aio/csp_helpers_async.py

@@ -202,6 +202,8 @@ async def async_get_arn():
         )
 
         # Mock the async STS client for direct aioboto3 usage
+        fake_aws_self = self
+
         class MockStsClient:
             async def __aenter__(self):
                 return self
@@ -212,6 +214,9 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
             async def get_caller_identity(self):
                 return await async_get_caller_identity()
 
+            async def assume_role(self, **kwargs):
+                return fake_aws_self.assume_role(**kwargs)
+
         def mock_session_client(service_name):
             if service_name == "sts":
                 return MockStsClient()

test/unit/aio/test_auth_workload_identity_async.py

@@ -253,6 +253,22 @@ async def test_explicit_aws_generates_unique_assertion_content(
     )
 
 
+async def test_aws_impersonation_calls_correct_apis_for_each_role_in_impersonation_path(
+    fake_aws_environment: FakeAwsEnvironmentAsync,
+):
+    impersonation_path = [
+        "arn:aws:iam::123456789:role/role2",
+        "arn:aws:iam::123456789:role/role3",
+    ]
+    fake_aws_environment.assumption_path = impersonation_path
+    auth_class = AuthByWorkloadIdentity(
+        provider=AttestationProvider.AWS, impersonation_path=impersonation_path
+    )
+    await auth_class.prepare(conn=None)
+
+    assert fake_aws_environment.assume_role_call_count == 2
+
+
 # -- GCP Tests --
 
 

test/unit/aio/test_connection_async_unit.py

@@ -659,16 +659,14 @@ async def test_workload_identity_provider_is_required_for_wif_authenticator(
     "provider_param",
     [
         # Strongly-typed values.
-        AttestationProvider.AWS,
         AttestationProvider.AZURE,
         AttestationProvider.OIDC,
         # String values.
-        "AWS",
         "AZURE",
         "OIDC",
     ],
 )
-async def test_workload_identity_impersonation_path_unsupported_for_non_gcp_providers(
+async def test_workload_identity_impersonation_path_errors_for_unsupported_providers(
     monkeypatch, provider_param
 ):
     async def mock_authenticate(*_):
@@ -690,20 +688,22 @@ async def mock_authenticate(*_):
                 ],
             )
         assert (
-            "workload_identity_impersonation_path is currently only supported for GCP."
+            "workload_identity_impersonation_path is currently only supported for GCP and AWS."
             in str(excinfo.value)
         )
 
 
 @pytest.mark.parametrize(
-    "provider_param",
+    "provider_param,impersonation_path",
     [
-        AttestationProvider.GCP,
-        "GCP",
+        (AttestationProvider.GCP, ["sa2@project.iam.gserviceaccount.com"]),
+        (AttestationProvider.AWS, ["arn:aws:iam::1234567890:role/role2"]),
+        ("GCP", ["sa2@project.iam.gserviceaccount.com"]),
+        ("AWS", ["arn:aws:iam::1234567890:role/role2"]),
     ],
 )
-async def test_workload_identity_impersonation_path_supported_for_gcp_provider(
-    monkeypatch, provider_param
+async def test_workload_identity_impersonation_path_populates_auth_class_for_supported_provider(
+    monkeypatch, provider_param, impersonation_path
 ):
     async def mock_authenticate(*_):
         pass
@@ -718,14 +718,9 @@ async def mock_authenticate(*_):
             account="account",
             authenticator="WORKLOAD_IDENTITY",
             workload_identity_provider=provider_param,
-            workload_identity_impersonation_path=[
-                "sa2@project.iam.gserviceaccount.com"
-            ],
+            workload_identity_impersonation_path=impersonation_path,
         )
-        assert conn.auth_class.provider == AttestationProvider.GCP
-        assert conn.auth_class.impersonation_path == [
-            "sa2@project.iam.gserviceaccount.com"
-        ]
+        assert conn.auth_class.impersonation_path == impersonation_path
 
 
 @pytest.mark.parametrize(