Skip to content

Add AWS GOV partitions support #1056

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Feb 26, 2025
Merged

Add AWS GOV partitions support #1056

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
0eebfd3
Add AWS GOV/CN partitions support
MonkeyCanCode Feb 24, 2025
fd24556
Fix style
MonkeyCanCode Feb 24, 2025
8ace28a
Disable CN partition support
MonkeyCanCode Feb 25, 2025
fbf09d1
Merge branch 'main' into gov_support
MonkeyCanCode Feb 25, 2025
935c615
Disable CN partition support
MonkeyCanCode Feb 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions ...e/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,14 @@ public EnumMap<PolarisCredentialProperty, String> getSubscopedCreds(
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;
}

Expand Down Expand Up @@ -119,7 +127,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);
Expand Down Expand Up @@ -155,7 +162,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), "*", "/")));
Expand Down
25 changes: 21 additions & 4 deletions ...s-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -86,9 +87,9 @@ 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");
// 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");
Expand Down Expand Up @@ -128,7 +129,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);
Expand Down
191 changes: 137 additions & 54 deletions ...est/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -101,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);
Expand Down Expand Up @@ -213,24 +218,48 @@ public void testGetSubscopedCredsInlinePolicy(String awsPartition) {
});
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
Expand Down Expand Up @@ -481,10 +510,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";
Expand All @@ -494,28 +524,54 @@ public void testClientRegion() {
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";
Expand All @@ -524,20 +580,47 @@ public void testNoClientRegion() {
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) {
Expand Down
5 changes: 2 additions & 3 deletions ...us/service/src/test/java/org/apache/polaris/service/quarkus/entity/CatalogEntityTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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", "aws-cn"})
public void testInvalidArn(String roleArn) {
String basedLocation = "s3://externally-owned-bucket";
AwsStorageConfigInfo awsStorageConfigModel =
Expand All @@ -210,8 +210,7 @@ public void testInvalidArn(String roleArn) {
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";
Expand Down