Fix bug in AWS sovereign partition support (#2459)

Author: sfc-gh-pmansour

src/snowflake/connector/wif_util.py

@@ -92,41 +92,6 @@ def get_aws_region() -> str:
     return region
 
 
-def get_aws_arn() -> str:
-    """Get the current AWS workload's ARN."""
-    caller_identity = boto3.client("sts").get_caller_identity()
-    if not caller_identity or "Arn" not in caller_identity:
-        raise ProgrammingError(
-            msg="No AWS identity was found. Ensure the application is running on AWS with an IAM role attached.",
-            errno=ER_WIF_CREDENTIALS_NOT_FOUND,
-        )
-    return caller_identity["Arn"]
-
-
-def get_aws_partition(arn: str) -> str:
-    """Get the current AWS partition from ARN.
-
-    Args:
-        arn (str): The Amazon Resource Name (ARN) string.
-
-    Returns:
-        str: The AWS partition (e.g., 'aws', 'aws-cn', 'aws-us-gov').
-
-    Raises:
-        ProgrammingError: If the ARN is invalid or does not contain a valid partition.
-
-    Reference: https://docs.aws.amazon.com/IAM/latest/UserGuide/reference-arns.html.
-    """
-    parts = arn.split(":")
-    if len(parts) > 1 and parts[0] == "arn" and parts[1]:
-        return parts[1]
-
-    raise ProgrammingError(
-        msg=f"Invalid AWS ARN: '{arn}'.",
-        errno=ER_WIF_CREDENTIALS_NOT_FOUND,
-    )
-
-
 def get_aws_sts_hostname(region: str, partition: str) -> str:
     """Constructs the AWS STS hostname for a given region and partition.
 
@@ -169,15 +134,15 @@ def create_aws_attestation() -> WorkloadIdentityAttestation:
 
     If the application isn't running on AWS or no credentials were found, raises an error.
     """
-    aws_creds = boto3.session.Session().get_credentials()
+    session = boto3.session.Session()
+    aws_creds = session.get_credentials()
     if not aws_creds:
         raise ProgrammingError(
             msg="No AWS credentials were found. Ensure the application is running on AWS with an IAM role attached.",
             errno=ER_WIF_CREDENTIALS_NOT_FOUND,
         )
     region = get_aws_region()
-    arn = get_aws_arn()
-    partition = get_aws_partition(arn)
+    partition = session.get_partition_for_region(region)
     sts_hostname = get_aws_sts_hostname(region, partition)
     request = AWSRequest(
         method="POST",
@@ -196,8 +161,10 @@ def create_aws_attestation() -> WorkloadIdentityAttestation:
         "headers": dict(request.headers.items()),
     }
     credential = b64encode(json.dumps(assertion_dict).encode("utf-8")).decode("utf-8")
+    # Unlike other providers, for AWS, we only include general identifiers (region and partition)
+    # rather than specific user identifiers, since we don't actually execute a GetCallerIdentity call.
     return WorkloadIdentityAttestation(
-        AttestationProvider.AWS, credential, {"arn": arn}
+        AttestationProvider.AWS, credential, {"region": region, "partition": partition}
     )
 
 

test/csp_helpers.py

@@ -358,9 +358,6 @@ def __init__(self):
     def get_region(self):
         return self.region
 
-    def get_arn(self):
-        return self.arn
-
     def get_credentials(self):
         return self.credentials
 
@@ -403,11 +400,6 @@ def __enter__(self):
                 side_effect=self.get_region,
             )
         )
-        self.patchers.append(
-            mock.patch(
-                "snowflake.connector.wif_util.get_aws_arn", side_effect=self.get_arn
-            )
-        )
         self.patchers.append(
             mock.patch(
                 "snowflake.connector.platform_detection.IMDSFetcher._get_request",

test/unit/test_auth_workload_identity.py

@@ -15,11 +15,7 @@
     HTTPError,
     Timeout,
 )
-from snowflake.connector.wif_util import (
-    AttestationProvider,
-    get_aws_partition,
-    get_aws_sts_hostname,
-)
+from snowflake.connector.wif_util import AttestationProvider, get_aws_sts_hostname
 
 from ..csp_helpers import FakeAwsEnvironment, FakeGceMetadataService, gen_dummy_id_token
 
@@ -129,8 +125,19 @@ def test_explicit_aws_encodes_audience_host_signature_to_api(
     verify_aws_token(data["TOKEN"], fake_aws_environment.region)
 
 
-def test_explicit_aws_uses_regional_hostname(fake_aws_environment: FakeAwsEnvironment):
-    fake_aws_environment.region = "antarctica-northeast-3"
+@pytest.mark.parametrize(
+    "region,expected_hostname",
+    [
+        ("us-east-1", "sts.us-east-1.amazonaws.com"),
+        ("af-south-1", "sts.af-south-1.amazonaws.com"),
+        ("us-gov-west-1", "sts.us-gov-west-1.amazonaws.com"),
+        ("cn-north-1", "sts.cn-north-1.amazonaws.com.cn"),
+    ],
+)
+def test_explicit_aws_uses_regional_hostnames(
+    fake_aws_environment: FakeAwsEnvironment, region: str, expected_hostname: str
+):
+    fake_aws_environment.region = region
 
     auth_class = AuthByWorkloadIdentity(provider=AttestationProvider.AWS)
     auth_class.prepare()
@@ -140,59 +147,23 @@ def test_explicit_aws_uses_regional_hostname(fake_aws_environment: FakeAwsEnviro
     hostname_from_url = urlparse(decoded_token["url"]).hostname
     hostname_from_header = decoded_token["headers"]["Host"]
 
-    expected_hostname = "sts.antarctica-northeast-3.amazonaws.com"
     assert expected_hostname == hostname_from_url
     assert expected_hostname == hostname_from_header
 
 
 def test_explicit_aws_generates_unique_assertion_content(
     fake_aws_environment: FakeAwsEnvironment,
 ):
-    fake_aws_environment.arn = (
-        "arn:aws:sts::123456789:assumed-role/A-Different-Role/i-34afe100cad287fab"
-    )
+    fake_aws_environment.region = "us-east-1"
     auth_class = AuthByWorkloadIdentity(provider=AttestationProvider.AWS)
     auth_class.prepare()
 
     assert (
-        '{"_provider":"AWS","arn":"arn:aws:sts::123456789:assumed-role/A-Different-Role/i-34afe100cad287fab"}'
+        '{"_provider":"AWS","partition":"aws","region":"us-east-1"}'
         == auth_class.assertion_content
     )
 
 
-@pytest.mark.parametrize(
-    "arn, expected_partition",
-    [
-        ("arn:aws:iam::123456789012:role/MyTestRole", "aws"),
-        (
-            "arn:aws-cn:ec2:cn-north-1:987654321098:instance/i-1234567890abcdef0",
-            "aws-cn",
-        ),
-        ("arn:aws-us-gov:s3:::my-gov-bucket", "aws-us-gov"),
-        ("arn:aws:s3:::my-bucket/my/key", "aws"),
-        ("arn:aws:lambda:us-east-1:123456789012:function:my-function", "aws"),
-        ("arn:aws:sns:eu-west-1:111122223333:my-topic", "aws"),
-        ("arn:aws:iam:", "aws"),  # Incomplete ARN, but partition is present
-    ],
-)
-def test_get_aws_partition_valid_arns(arn, expected_partition):
-    assert get_aws_partition(arn) == expected_partition
-
-
-@pytest.mark.parametrize(
-    "arn",
-    [
-        "invalid-arn",
-        "arn::service:region:account:resource",  # Missing partition
-        "",  # Empty string
-    ],
-)
-def test_get_aws_partition_invalid_arns(arn):
-    with pytest.raises(ProgrammingError) as excinfo:
-        get_aws_partition(arn)
-    assert "Invalid AWS ARN" in str(excinfo.value)
-
-
 @pytest.mark.parametrize(
     "region, partition, expected_hostname",
     [