From 0eebfd376f513bd360bcf91ab66f05037ca4f75a Mon Sep 17 00:00:00 2001 From: Yong Date: Sun, 23 Feb 2025 22:59:20 -0600 Subject: [PATCH 1/3] Add AWS GOV/CN partitions support --- .../aws/AwsCredentialsStorageIntegration.java | 9 +- .../aws/AwsStorageConfigurationInfo.java | 24 +++- .../AwsCredentialsStorageIntegrationTest.java | 107 ++++++++++++------ .../quarkus/entity/CatalogEntityTest.java | 6 +- 4 files changed, 97 insertions(+), 49 deletions(-) diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java index 636734831a..44a9628033 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java @@ -90,6 +90,13 @@ public EnumMap getSubscopedCreds( if (storageConfig.getRegion() != null) { credentialMap.put(PolarisCredentialProperty.CLIENT_REGION, storageConfig.getRegion()); } + + if (!storageConfig.getAwsPartition().equals("aws") + && credentialMap.get(PolarisCredentialProperty.CLIENT_REGION) == null) { + throw new IllegalArgumentException( + String.format("AWS region must be set when using partition %s", storageConfig.getAwsPartition())); + } + return credentialMap; } @@ -119,7 +126,6 @@ private IamPolicy policyString( location -> { URI uri = URI.create(location); allowGetObjectStatementBuilder.addResource( - // TODO add support for CN and GOV IamResource.create( arnPrefix + StorageUtil.concatFilePrefixes(parseS3Path(uri), "*", "/"))); final var bucket = arnPrefix + StorageUtil.getBucket(uri); @@ -155,7 +161,6 @@ private IamPolicy policyString( writeLocations.forEach( location -> { URI uri = URI.create(location); - // TODO add support for CN and GOV allowPutObjectStatementBuilder.addResource( IamResource.create( arnPrefix + StorageUtil.concatFilePrefixes(parseS3Path(uri), "*", "/"))); diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java index dc9941d9e6..88820a2735 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java @@ -36,8 +36,8 @@ public class AwsStorageConfigurationInfo extends PolarisStorageConfigurationInfo // for allowed read and write locations for subscoping creds. @JsonIgnore private static final int MAX_ALLOWED_LOCATIONS = 5; - // Technically, it should be ^arn:(aws|aws-cn|aws-us-gov):iam::(\d{12}):role/.+$, - @JsonIgnore public static final String ROLE_ARN_PATTERN = "^arn:aws:iam::(\\d{12}):role/.+$"; + // AWS ROLE ARN pattern + @JsonIgnore public static final String ROLE_ARN_PATTERN = "^arn:(aws|aws-cn|aws-us-gov):iam::(\\d{12}):role/.+$"; // AWS role to be assumed private final @Nonnull String roleARN; @@ -86,10 +86,6 @@ public static void validateArn(String arn) { if (arn == null || arn.isEmpty()) { throw new IllegalArgumentException("ARN cannot be null or empty"); } - // specifically throw errors for China and Gov - if (arn.contains("aws-cn") || arn.contains("aws-us-gov")) { - throw new IllegalArgumentException("AWS China or Gov Cloud are temporarily not supported"); - } if (!Pattern.matches(ROLE_ARN_PATTERN, arn)) { throw new IllegalArgumentException("Invalid role ARN format"); } @@ -128,7 +124,23 @@ public String getAwsAccountId() { return parseAwsAccountId(roleARN); } + @JsonIgnore + public String getAwsPartition() { + return parseAwsPartition(roleARN); + } + private static String parseAwsAccountId(String arn) { + validateArn(arn); + Pattern pattern = Pattern.compile(ROLE_ARN_PATTERN); + Matcher matcher = pattern.matcher(arn); + if (matcher.matches()) { + return matcher.group(2); + } else { + throw new IllegalArgumentException("ARN does not match the expected role ARN pattern"); + } + } + + private static String parseAwsPartition(String arn) { validateArn(arn); Pattern pattern = Pattern.compile(ROLE_ARN_PATTERN); Matcher matcher = pattern.matcher(arn); diff --git a/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java b/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java index c552494e5c..920f7910b3 100644 --- a/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java @@ -29,6 +29,7 @@ import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo; import org.apache.polaris.core.storage.aws.AwsCredentialsStorageIntegration; import org.apache.polaris.core.storage.aws.AwsStorageConfigurationInfo; +import org.assertj.core.api.Assertions; import org.assertj.core.api.InstanceOfAssertFactories; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; @@ -481,10 +482,11 @@ public void testGetSubscopedCredsInlinePolicyWithEmptyReadAndWrite() { .containsEntry(PolarisCredentialProperty.AWS_SECRET_KEY, "secretKey"); } - @Test - public void testClientRegion() { + @ParameterizedTest + @ValueSource(strings = {AWS_PARTITION, "aws-cn", "aws-us-gov"}) + public void testClientRegion(String awsPartition) { StsClient stsClient = Mockito.mock(StsClient.class); - String roleARN = "arn:aws:iam::012345678901:role/jdoe"; + String roleARN = "arn:aws:iam::012345678901:role/jdoe".replaceFirst("aws", awsPartition); String externalId = "externalId"; String bucket = "bucket"; String warehouseKeyPrefix = "path/to/warehouse"; @@ -494,28 +496,38 @@ public void testClientRegion() { invocation -> { return ASSUME_ROLE_RESPONSE; }); - EnumMap credentials = - new AwsCredentialsStorageIntegration(stsClient) - .getSubscopedCreds( - Mockito.mock(PolarisDiagnostics.class), - new AwsStorageConfigurationInfo( - PolarisStorageConfigurationInfo.StorageType.S3, - List.of(s3Path(bucket, warehouseKeyPrefix)), - roleARN, - externalId, - clientRegion), - true, /* allowList = true */ - Set.of(), - Set.of()); - assertThat(credentials) - .isNotEmpty() - .containsEntry(PolarisCredentialProperty.CLIENT_REGION, clientRegion); + switch (awsPartition) { + case AWS_PARTITION: + case "aws-cn": + case "aws-us-gov": + EnumMap credentials = + new AwsCredentialsStorageIntegration(stsClient) + .getSubscopedCreds( + Mockito.mock(PolarisDiagnostics.class), + new AwsStorageConfigurationInfo( + PolarisStorageConfigurationInfo.StorageType.S3, + List.of(s3Path(bucket, warehouseKeyPrefix)), + roleARN, + externalId, + clientRegion), + true, /* allowList = true */ + Set.of(), + Set.of()); + assertThat(credentials) + .isNotEmpty() + .containsEntry(PolarisCredentialProperty.CLIENT_REGION, clientRegion); + break; + default: + throw new IllegalArgumentException("Unknown aws partition: " + awsPartition); + + }; } - @Test - public void testNoClientRegion() { + @ParameterizedTest + @ValueSource(strings = {AWS_PARTITION, "aws-cn", "aws-us-gov"}) + public void testNoClientRegion(String awsPartition) { StsClient stsClient = Mockito.mock(StsClient.class); - String roleARN = "arn:aws:iam::012345678901:role/jdoe"; + String roleARN = "arn:aws:iam::012345678901:role/jdoe".replaceFirst("aws", awsPartition); String externalId = "externalId"; String bucket = "bucket"; String warehouseKeyPrefix = "path/to/warehouse"; @@ -524,20 +536,43 @@ public void testNoClientRegion() { invocation -> { return ASSUME_ROLE_RESPONSE; }); - EnumMap credentials = - new AwsCredentialsStorageIntegration(stsClient) - .getSubscopedCreds( - Mockito.mock(PolarisDiagnostics.class), - new AwsStorageConfigurationInfo( - PolarisStorageConfigurationInfo.StorageType.S3, - List.of(s3Path(bucket, warehouseKeyPrefix)), - roleARN, - externalId, - null), - true, /* allowList = true */ - Set.of(), - Set.of()); - assertThat(credentials).isNotEmpty().doesNotContainKey(PolarisCredentialProperty.CLIENT_REGION); + switch (awsPartition) { + case AWS_PARTITION: + EnumMap credentials = + new AwsCredentialsStorageIntegration(stsClient) + .getSubscopedCreds( + Mockito.mock(PolarisDiagnostics.class), + new AwsStorageConfigurationInfo( + PolarisStorageConfigurationInfo.StorageType.S3, + List.of(s3Path(bucket, warehouseKeyPrefix)), + roleARN, + externalId, + null), + true, /* allowList = true */ + Set.of(), + Set.of()); + assertThat(credentials).isNotEmpty().doesNotContainKey(PolarisCredentialProperty.CLIENT_REGION); + break; + case "aws-cn": + case "aws-us-gov": + Assertions.assertThatThrownBy( () -> + new AwsCredentialsStorageIntegration(stsClient) + .getSubscopedCreds( + Mockito.mock(PolarisDiagnostics.class), + new AwsStorageConfigurationInfo( + PolarisStorageConfigurationInfo.StorageType.S3, + List.of(s3Path(bucket, warehouseKeyPrefix)), + roleARN, + externalId, + null), + true, /* allowList = true */ + Set.of(), + Set.of())).isInstanceOf(IllegalArgumentException.class); + break; + default: + throw new IllegalArgumentException("Unknown aws partition: " + awsPartition); + + }; } private static @Nonnull String s3Arn(String partition, String bucket, String keyPrefix) { diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/entity/CatalogEntityTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/entity/CatalogEntityTest.java index 2d4667d106..261b779957 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/entity/CatalogEntityTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/entity/CatalogEntityTest.java @@ -185,7 +185,7 @@ public void testValidAllowedLocationPrefix() { } @ParameterizedTest - @ValueSource(strings = {"", "arn:aws:iam::0123456:role/jdoe", "aws-cn", "aws-us-gov"}) + @ValueSource(strings = {"", "arn:aws:iam::0123456:role/jdoe"}) public void testInvalidArn(String roleArn) { String basedLocation = "s3://externally-owned-bucket"; AwsStorageConfigInfo awsStorageConfigModel = @@ -209,10 +209,6 @@ public void testInvalidArn(String roleArn) { case "": expectedMessage = "ARN cannot be null or empty"; break; - case "aws-cn": - case "aws-us-gov": - expectedMessage = "AWS China or Gov Cloud are temporarily not supported"; - break; default: expectedMessage = "Invalid role ARN format"; } From fd2455653a83e3fd126673b03ca0a92e6a2a0faf Mon Sep 17 00:00:00 2001 From: Yong Zheng Date: Sun, 23 Feb 2025 23:55:32 -0600 Subject: [PATCH 2/3] Fix style --- .../aws/AwsCredentialsStorageIntegration.java | 5 +- .../aws/AwsStorageConfigurationInfo.java | 4 +- .../AwsCredentialsStorageIntegrationTest.java | 126 +++++++++--------- 3 files changed, 71 insertions(+), 64 deletions(-) diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java index 44a9628033..42a048ef1f 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java @@ -92,9 +92,10 @@ public EnumMap getSubscopedCreds( } if (!storageConfig.getAwsPartition().equals("aws") - && credentialMap.get(PolarisCredentialProperty.CLIENT_REGION) == null) { + && credentialMap.get(PolarisCredentialProperty.CLIENT_REGION) == null) { throw new IllegalArgumentException( - String.format("AWS region must be set when using partition %s", storageConfig.getAwsPartition())); + String.format( + "AWS region must be set when using partition %s", storageConfig.getAwsPartition())); } return credentialMap; diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java index 88820a2735..16caa7cb44 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java @@ -37,7 +37,9 @@ public class AwsStorageConfigurationInfo extends PolarisStorageConfigurationInfo @JsonIgnore private static final int MAX_ALLOWED_LOCATIONS = 5; // AWS ROLE ARN pattern - @JsonIgnore public static final String ROLE_ARN_PATTERN = "^arn:(aws|aws-cn|aws-us-gov):iam::(\\d{12}):role/.+$"; + @JsonIgnore + public static final String ROLE_ARN_PATTERN = + "^arn:(aws|aws-cn|aws-us-gov):iam::(\\d{12}):role/.+$"; // AWS role to be assumed private final @Nonnull String roleARN; diff --git a/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java b/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java index 920f7910b3..0cde1b7944 100644 --- a/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java @@ -497,30 +497,30 @@ public void testClientRegion(String awsPartition) { return ASSUME_ROLE_RESPONSE; }); switch (awsPartition) { - case AWS_PARTITION: - case "aws-cn": - case "aws-us-gov": - EnumMap credentials = - new AwsCredentialsStorageIntegration(stsClient) - .getSubscopedCreds( - Mockito.mock(PolarisDiagnostics.class), - new AwsStorageConfigurationInfo( - PolarisStorageConfigurationInfo.StorageType.S3, - List.of(s3Path(bucket, warehouseKeyPrefix)), - roleARN, - externalId, - clientRegion), - true, /* allowList = true */ - Set.of(), - Set.of()); - assertThat(credentials) - .isNotEmpty() - .containsEntry(PolarisCredentialProperty.CLIENT_REGION, clientRegion); - break; - default: - throw new IllegalArgumentException("Unknown aws partition: " + awsPartition); - - }; + case AWS_PARTITION: + case "aws-cn": + case "aws-us-gov": + EnumMap credentials = + new AwsCredentialsStorageIntegration(stsClient) + .getSubscopedCreds( + Mockito.mock(PolarisDiagnostics.class), + new AwsStorageConfigurationInfo( + PolarisStorageConfigurationInfo.StorageType.S3, + List.of(s3Path(bucket, warehouseKeyPrefix)), + roleARN, + externalId, + clientRegion), + true, /* allowList = true */ + Set.of(), + Set.of()); + assertThat(credentials) + .isNotEmpty() + .containsEntry(PolarisCredentialProperty.CLIENT_REGION, clientRegion); + break; + default: + throw new IllegalArgumentException("Unknown aws partition: " + awsPartition); + } + ; } @ParameterizedTest @@ -536,43 +536,47 @@ public void testNoClientRegion(String awsPartition) { invocation -> { return ASSUME_ROLE_RESPONSE; }); - switch (awsPartition) { - case AWS_PARTITION: - EnumMap credentials = - new AwsCredentialsStorageIntegration(stsClient) - .getSubscopedCreds( - Mockito.mock(PolarisDiagnostics.class), - new AwsStorageConfigurationInfo( - PolarisStorageConfigurationInfo.StorageType.S3, - List.of(s3Path(bucket, warehouseKeyPrefix)), - roleARN, - externalId, - null), - true, /* allowList = true */ - Set.of(), - Set.of()); - assertThat(credentials).isNotEmpty().doesNotContainKey(PolarisCredentialProperty.CLIENT_REGION); - break; - case "aws-cn": - case "aws-us-gov": - Assertions.assertThatThrownBy( () -> - new AwsCredentialsStorageIntegration(stsClient) - .getSubscopedCreds( - Mockito.mock(PolarisDiagnostics.class), - new AwsStorageConfigurationInfo( - PolarisStorageConfigurationInfo.StorageType.S3, - List.of(s3Path(bucket, warehouseKeyPrefix)), - roleARN, - externalId, - null), - true, /* allowList = true */ - Set.of(), - Set.of())).isInstanceOf(IllegalArgumentException.class); - break; - default: - throw new IllegalArgumentException("Unknown aws partition: " + awsPartition); - - }; + switch (awsPartition) { + case AWS_PARTITION: + EnumMap credentials = + new AwsCredentialsStorageIntegration(stsClient) + .getSubscopedCreds( + Mockito.mock(PolarisDiagnostics.class), + new AwsStorageConfigurationInfo( + PolarisStorageConfigurationInfo.StorageType.S3, + List.of(s3Path(bucket, warehouseKeyPrefix)), + roleARN, + externalId, + null), + true, /* allowList = true */ + Set.of(), + Set.of()); + assertThat(credentials) + .isNotEmpty() + .doesNotContainKey(PolarisCredentialProperty.CLIENT_REGION); + break; + case "aws-cn": + case "aws-us-gov": + Assertions.assertThatThrownBy( + () -> + new AwsCredentialsStorageIntegration(stsClient) + .getSubscopedCreds( + Mockito.mock(PolarisDiagnostics.class), + new AwsStorageConfigurationInfo( + PolarisStorageConfigurationInfo.StorageType.S3, + List.of(s3Path(bucket, warehouseKeyPrefix)), + roleARN, + externalId, + null), + true, /* allowList = true */ + Set.of(), + Set.of())) + .isInstanceOf(IllegalArgumentException.class); + break; + default: + throw new IllegalArgumentException("Unknown aws partition: " + awsPartition); + } + ; } private static @Nonnull String s3Arn(String partition, String bucket, String keyPrefix) { From 8ace28a700d558ad24a4021df5d63bbadca19ae0 Mon Sep 17 00:00:00 2001 From: Yong Date: Mon, 24 Feb 2025 23:53:42 -0600 Subject: [PATCH 3/3] Disable CN partition support --- .../aws/AwsCredentialsStorageIntegration.java | 7 +- .../aws/AwsStorageConfigurationInfo.java | 9 +- .../AwsCredentialsStorageIntegrationTest.java | 204 +++++++++++------- .../quarkus/entity/CatalogEntityTest.java | 5 +- 4 files changed, 141 insertions(+), 84 deletions(-) diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java index 44a9628033..f522b77b57 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java @@ -91,10 +91,11 @@ public EnumMap getSubscopedCreds( credentialMap.put(PolarisCredentialProperty.CLIENT_REGION, storageConfig.getRegion()); } - if (!storageConfig.getAwsPartition().equals("aws") - && credentialMap.get(PolarisCredentialProperty.CLIENT_REGION) == null) { + if (storageConfig.getAwsPartition().equals("aws-us-gov") + && credentialMap.get(PolarisCredentialProperty.CLIENT_REGION) == null) { throw new IllegalArgumentException( - String.format("AWS region must be set when using partition %s", storageConfig.getAwsPartition())); + String.format( + "AWS region must be set when using partition %s", storageConfig.getAwsPartition())); } return credentialMap; diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java index 88820a2735..ebd2454035 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java @@ -36,8 +36,9 @@ public class AwsStorageConfigurationInfo extends PolarisStorageConfigurationInfo // for allowed read and write locations for subscoping creds. @JsonIgnore private static final int MAX_ALLOWED_LOCATIONS = 5; - // AWS ROLE ARN pattern - @JsonIgnore public static final String ROLE_ARN_PATTERN = "^arn:(aws|aws-cn|aws-us-gov):iam::(\\d{12}):role/.+$"; + // Technically, it should be ^arn:(aws|aws-cn|aws-us-gov):iam::(\d{12}):role/.+$, + @JsonIgnore + public static final String ROLE_ARN_PATTERN = "^arn:(aws|aws-us-gov):iam::(\\d{12}):role/.+$"; // AWS role to be assumed private final @Nonnull String roleARN; @@ -86,6 +87,10 @@ public static void validateArn(String arn) { if (arn == null || arn.isEmpty()) { throw new IllegalArgumentException("ARN cannot be null or empty"); } + // specifically throw errors for China + if (arn.contains("aws-cn")) { + throw new IllegalArgumentException("AWS China is temporarily not supported"); + } if (!Pattern.matches(ROLE_ARN_PATTERN, arn)) { throw new IllegalArgumentException("Invalid role ARN format"); } diff --git a/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java b/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java index 920f7910b3..5ef145d65e 100644 --- a/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java @@ -102,15 +102,19 @@ public void testGetSubscopedCredsInlinePolicy(String awsPartition) { PolarisStorageConfigurationInfo.StorageType storageType = PolarisStorageConfigurationInfo.StorageType.S3; String roleARN; + String region; switch (awsPartition) { case AWS_PARTITION: roleARN = "arn:aws:iam::012345678901:role/jdoe"; + region = "us-east-1"; break; case "aws-cn": roleARN = "arn:aws-cn:iam::012345678901:role/jdoe"; + region = "Beijing"; break; case "aws-us-gov": roleARN = "arn:aws-us-gov:iam::012345678901:role/jdoe"; + region = "us-gov-west-1"; break; default: throw new IllegalArgumentException("Unknown aws partition: " + awsPartition); @@ -214,24 +218,48 @@ public void testGetSubscopedCredsInlinePolicy(String awsPartition) { }); return ASSUME_ROLE_RESPONSE; }); - EnumMap credentials = - new AwsCredentialsStorageIntegration(stsClient) - .getSubscopedCreds( - Mockito.mock(PolarisDiagnostics.class), - new AwsStorageConfigurationInfo( - storageType, - List.of(s3Path(bucket, warehouseKeyPrefix)), - roleARN, - externalId, - "us-east-2"), - true, - Set.of(s3Path(bucket, firstPath), s3Path(bucket, secondPath)), - Set.of(s3Path(bucket, firstPath))); - assertThat(credentials) - .isNotEmpty() - .containsEntry(PolarisCredentialProperty.AWS_TOKEN, "sess") - .containsEntry(PolarisCredentialProperty.AWS_KEY_ID, "accessKey") - .containsEntry(PolarisCredentialProperty.AWS_SECRET_KEY, "secretKey"); + switch (awsPartition) { + case "aws-cn": + Assertions.assertThatThrownBy( + () -> + new AwsCredentialsStorageIntegration(stsClient) + .getSubscopedCreds( + Mockito.mock(PolarisDiagnostics.class), + new AwsStorageConfigurationInfo( + storageType, + List.of(s3Path(bucket, warehouseKeyPrefix)), + roleARN, + externalId, + region), + true, + Set.of(s3Path(bucket, firstPath), s3Path(bucket, secondPath)), + Set.of(s3Path(bucket, firstPath)))) + .isInstanceOf(IllegalArgumentException.class); + break; + case AWS_PARTITION: + case "aws-us-gov": + EnumMap credentials = + new AwsCredentialsStorageIntegration(stsClient) + .getSubscopedCreds( + Mockito.mock(PolarisDiagnostics.class), + new AwsStorageConfigurationInfo( + storageType, + List.of(s3Path(bucket, warehouseKeyPrefix)), + roleARN, + externalId, + region), + true, + Set.of(s3Path(bucket, firstPath), s3Path(bucket, secondPath)), + Set.of(s3Path(bucket, firstPath))); + assertThat(credentials) + .isNotEmpty() + .containsEntry(PolarisCredentialProperty.AWS_TOKEN, "sess") + .containsEntry(PolarisCredentialProperty.AWS_KEY_ID, "accessKey") + .containsEntry(PolarisCredentialProperty.AWS_SECRET_KEY, "secretKey"); + break; + default: + throw new IllegalArgumentException("Unknown aws partition: " + awsPartition); + } } @Test @@ -497,30 +525,46 @@ public void testClientRegion(String awsPartition) { return ASSUME_ROLE_RESPONSE; }); switch (awsPartition) { - case AWS_PARTITION: - case "aws-cn": - case "aws-us-gov": - EnumMap credentials = + case "aws-cn": + Assertions.assertThatThrownBy( + () -> new AwsCredentialsStorageIntegration(stsClient) - .getSubscopedCreds( - Mockito.mock(PolarisDiagnostics.class), - new AwsStorageConfigurationInfo( - PolarisStorageConfigurationInfo.StorageType.S3, - List.of(s3Path(bucket, warehouseKeyPrefix)), - roleARN, - externalId, - clientRegion), - true, /* allowList = true */ - Set.of(), - Set.of()); - assertThat(credentials) - .isNotEmpty() - .containsEntry(PolarisCredentialProperty.CLIENT_REGION, clientRegion); - break; - default: - throw new IllegalArgumentException("Unknown aws partition: " + awsPartition); - - }; + .getSubscopedCreds( + Mockito.mock(PolarisDiagnostics.class), + new AwsStorageConfigurationInfo( + PolarisStorageConfigurationInfo.StorageType.S3, + List.of(s3Path(bucket, warehouseKeyPrefix)), + roleARN, + externalId, + clientRegion), + true, /* allowList = true */ + Set.of(), + Set.of())) + .isInstanceOf(IllegalArgumentException.class); + break; + case AWS_PARTITION: + case "aws-us-gov": + EnumMap credentials = + new AwsCredentialsStorageIntegration(stsClient) + .getSubscopedCreds( + Mockito.mock(PolarisDiagnostics.class), + new AwsStorageConfigurationInfo( + PolarisStorageConfigurationInfo.StorageType.S3, + List.of(s3Path(bucket, warehouseKeyPrefix)), + roleARN, + externalId, + clientRegion), + true, /* allowList = true */ + Set.of(), + Set.of()); + assertThat(credentials) + .isNotEmpty() + .containsEntry(PolarisCredentialProperty.CLIENT_REGION, clientRegion); + break; + default: + throw new IllegalArgumentException("Unknown aws partition: " + awsPartition); + } + ; } @ParameterizedTest @@ -536,43 +580,47 @@ public void testNoClientRegion(String awsPartition) { invocation -> { return ASSUME_ROLE_RESPONSE; }); - switch (awsPartition) { - case AWS_PARTITION: - EnumMap credentials = - new AwsCredentialsStorageIntegration(stsClient) - .getSubscopedCreds( - Mockito.mock(PolarisDiagnostics.class), - new AwsStorageConfigurationInfo( - PolarisStorageConfigurationInfo.StorageType.S3, - List.of(s3Path(bucket, warehouseKeyPrefix)), - roleARN, - externalId, - null), - true, /* allowList = true */ - Set.of(), - Set.of()); - assertThat(credentials).isNotEmpty().doesNotContainKey(PolarisCredentialProperty.CLIENT_REGION); - break; - case "aws-cn": - case "aws-us-gov": - Assertions.assertThatThrownBy( () -> - new AwsCredentialsStorageIntegration(stsClient) - .getSubscopedCreds( - Mockito.mock(PolarisDiagnostics.class), - new AwsStorageConfigurationInfo( - PolarisStorageConfigurationInfo.StorageType.S3, - List.of(s3Path(bucket, warehouseKeyPrefix)), - roleARN, - externalId, - null), - true, /* allowList = true */ - Set.of(), - Set.of())).isInstanceOf(IllegalArgumentException.class); - break; - default: - throw new IllegalArgumentException("Unknown aws partition: " + awsPartition); - - }; + switch (awsPartition) { + case AWS_PARTITION: + EnumMap credentials = + new AwsCredentialsStorageIntegration(stsClient) + .getSubscopedCreds( + Mockito.mock(PolarisDiagnostics.class), + new AwsStorageConfigurationInfo( + PolarisStorageConfigurationInfo.StorageType.S3, + List.of(s3Path(bucket, warehouseKeyPrefix)), + roleARN, + externalId, + null), + true, /* allowList = true */ + Set.of(), + Set.of()); + assertThat(credentials) + .isNotEmpty() + .doesNotContainKey(PolarisCredentialProperty.CLIENT_REGION); + break; + case "aws-cn": + case "aws-us-gov": + Assertions.assertThatThrownBy( + () -> + new AwsCredentialsStorageIntegration(stsClient) + .getSubscopedCreds( + Mockito.mock(PolarisDiagnostics.class), + new AwsStorageConfigurationInfo( + PolarisStorageConfigurationInfo.StorageType.S3, + List.of(s3Path(bucket, warehouseKeyPrefix)), + roleARN, + externalId, + null), + true, /* allowList = true */ + Set.of(), + Set.of())) + .isInstanceOf(IllegalArgumentException.class); + break; + default: + throw new IllegalArgumentException("Unknown aws partition: " + awsPartition); + } + ; } private static @Nonnull String s3Arn(String partition, String bucket, String keyPrefix) { diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/entity/CatalogEntityTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/entity/CatalogEntityTest.java index 261b779957..4a4859251a 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/entity/CatalogEntityTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/entity/CatalogEntityTest.java @@ -185,7 +185,7 @@ public void testValidAllowedLocationPrefix() { } @ParameterizedTest - @ValueSource(strings = {"", "arn:aws:iam::0123456:role/jdoe"}) + @ValueSource(strings = {"", "arn:aws:iam::0123456:role/jdoe", "aws-cn"}) public void testInvalidArn(String roleArn) { String basedLocation = "s3://externally-owned-bucket"; AwsStorageConfigInfo awsStorageConfigModel = @@ -209,6 +209,9 @@ public void testInvalidArn(String roleArn) { case "": expectedMessage = "ARN cannot be null or empty"; break; + case "aws-cn": + expectedMessage = "AWS China is temporarily not supported"; + break; default: expectedMessage = "Invalid role ARN format"; }