fix aws impersonation

Author: sfc-gh-eqin

src/snowflake/connector/wif_util.py

@@ -145,6 +145,24 @@ def get_aws_sts_hostname(region: str, partition: str) -> str:
         )
 
 
+def get_aws_session(impersonation_path: list[str] | None = None):
+    """Creates a boto3 session with the appropriate credentials. If impersonation_path is provided, uses role at end of the path."""
+    session = boto3.session.Session()
+
+    impersonation_path = impersonation_path or []
+    for arn in impersonation_path:
+        response = session.client("sts").assume_role(
+            RoleArn=arn, RoleSessionName="identity-federation-session"
+        )
+        creds = response["Credentials"]
+        session = boto3.Session(
+            aws_access_key_id=creds["AccessKeyId"],
+            aws_secret_access_key=creds["SecretAccessKey"],
+            aws_session_token=creds["SessionToken"],
+        )
+    return session
+
+
 def create_aws_attestation(
     impersonation_path: list[str] | None = None,
 ) -> WorkloadIdentityAttestation:
@@ -153,26 +171,7 @@ def create_aws_attestation(
     If the application isn't running on AWS or no credentials were found, raises an error.
     """
     # TODO: SNOW-2223669 Investigate if our adapters - containing settings of http traffic - should be passed here as boto urllib3session. Those requests go to local servers, so they do not need Proxy setup or Headers customization in theory. But we may want to have all the traffic going through one class (e.g. Adapter or mixin).
-    session = boto3.session.Session()
-    if impersonation_path:
-        sts_client = boto3.client("sts")
-        for arn in impersonation_path:
-            # Assume target role
-            response = sts_client.assume_role(
-                RoleArn=arn, RoleSessionName="identity-federation-session"
-            )
-
-        # Use the credentials from the last assumed role
-        creds = response["Credentials"]
-        access_key = creds["AccessKeyId"]
-        secret_key = creds["SecretAccessKey"]
-        session_token = creds["SessionToken"]
-
-        session = boto3.Session(
-            aws_access_key_id=access_key,
-            aws_secret_access_key=secret_key,
-            aws_session_token=session_token,
-        )
+    session = get_aws_session(impersonation_path)
 
     aws_creds = session.get_credentials()
     if not aws_creds: