This is an automated email from the ASF dual-hosted git repository.
yzheng pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/polaris.git
The following commit(s) were added to refs/heads/main by this push:
new f0c5f05c Add AWS GOV partitions support (#1056)
f0c5f05c is described below
commit f0c5f05cadeee56c0f8a9dede0c8f311acb2e67d
Author: MonkeyCanCode <[email protected]>
AuthorDate: Tue Feb 25 19:11:27 2025 -0600
Add AWS GOV partitions support (#1056)
* Add AWS GOV/CN partitions support
* Fix style
* Disable CN partition support
---
.../aws/AwsCredentialsStorageIntegration.java | 10 +-
.../storage/aws/AwsStorageConfigurationInfo.java | 25 ++-
.../aws/AwsCredentialsStorageIntegrationTest.java | 191 +++++++++++++++------
.../service/quarkus/entity/CatalogEntityTest.java | 5 +-
4 files changed, 168 insertions(+), 63 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 63673483..f522b77b 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,14 @@ public class AwsCredentialsStorageIntegration
if (storageConfig.getRegion() != null) {
credentialMap.put(PolarisCredentialProperty.CLIENT_REGION,
storageConfig.getRegion());
}
+
+ 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()));
+ }
+
return credentialMap;
}
@@ -119,7 +127,6 @@ public class AwsCredentialsStorageIntegration
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 +162,6 @@ public class AwsCredentialsStorageIntegration
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 dc9941d9..ebd24540 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,8 @@ public class AwsStorageConfigurationInfo extends
PolarisStorageConfigurationInfo
@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/.+$";
+ @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,9 +87,9 @@ public class AwsStorageConfigurationInfo extends
PolarisStorageConfigurationInfo
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");
+ // 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");
@@ -128,7 +129,23 @@ public class AwsStorageConfigurationInfo extends
PolarisStorageConfigurationInfo
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 c552494e..5ef145d6 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.PolarisCredentialProperty;
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;
@@ -101,15 +102,19 @@ class AwsCredentialsStorageIntegrationTest {
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);
@@ -213,24 +218,48 @@ class AwsCredentialsStorageIntegrationTest {
});
return ASSUME_ROLE_RESPONSE;
});
- EnumMap<PolarisCredentialProperty, String> 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<PolarisCredentialProperty, String> 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
@@ -481,10 +510,11 @@ class AwsCredentialsStorageIntegrationTest {
.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 +524,54 @@ class AwsCredentialsStorageIntegrationTest {
invocation -> {
return ASSUME_ROLE_RESPONSE;
});
- EnumMap<PolarisCredentialProperty, String> 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-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()))
+ .isInstanceOf(IllegalArgumentException.class);
+ break;
+ case AWS_PARTITION:
+ case "aws-us-gov":
+ EnumMap<PolarisCredentialProperty, String> 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 +580,47 @@ class AwsCredentialsStorageIntegrationTest {
invocation -> {
return ASSUME_ROLE_RESPONSE;
});
- EnumMap<PolarisCredentialProperty, String> 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<PolarisCredentialProperty, String> 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 2d4667d1..4a485925 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 class CatalogEntityTest {
}
@ParameterizedTest
- @ValueSource(strings = {"", "arn:aws:iam::0123456:role/jdoe", "aws-cn",
"aws-us-gov"})
+ @ValueSource(strings = {"", "arn:aws:iam::0123456:role/jdoe", "aws-cn"})
public void testInvalidArn(String roleArn) {
String basedLocation = "s3://externally-owned-bucket";
AwsStorageConfigInfo awsStorageConfigModel =
@@ -210,8 +210,7 @@ public class CatalogEntityTest {
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";
+ expectedMessage = "AWS China is temporarily not supported";
break;
default:
expectedMessage = "Invalid role ARN format";
