Skip to content
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

Jump to bottom

Store AWS region in AwsStorageConfigurationInfo #455

Open
wants to merge 10 commits into
base: main
Choose a base branch
from
+82 −12
Open

Store AWS region in AwsStorageConfigurationInfo #455

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
e02e8d0
debugging
eric-maynard Nov 17, 2024
b2f3cc1
working
eric-maynard Nov 17, 2024
799f151
fix
eric-maynard Nov 17, 2024
134c77b
fix more invocations
eric-maynard Nov 17, 2024
75fa1f0
autolint
eric-maynard Nov 17, 2024
1a89f6b
patch one thing
eric-maynard Nov 18, 2024
09b6c43
Merge branch 'main' into issue-342
eric-maynard Nov 20, 2024
13f7218
add test
eric-maynard Nov 25, 2024
fc76b73
autolint
eric-maynard Nov 25, 2024
679bb54
Merge branch 'issue-342' of git.luolix.top-oss:eric-maynard/polaris into …
eric-maynard Nov 25, 2024
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
4 changes: 3 additions & 1 deletion polaris-core/src/main/java/org/apache/polaris/core/entity/CatalogEntity.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,7 @@ private StorageConfigInfo getStorageInfo(Map<String, String> internalProperties)
.setUserArn(awsConfig.getUserARN())
.setStorageType(StorageConfigInfo.StorageTypeEnum.S3)
.setAllowedLocations(awsConfig.getAllowedLocations())
.setRegion(awsConfig.getRegion())
.build();
}
if (configInfo instanceof AzureStorageConfigurationInfo) {
Expand Down Expand Up @@ -244,7 +245,8 @@ public Builder setStorageConfigurationInfo(
PolarisStorageConfigurationInfo.StorageType.S3,
new ArrayList<>(allowedLocations),
awsConfigModel.getRoleArn(),
awsConfigModel.getExternalId());
awsConfigModel.getExternalId(),
awsConfigModel.getRegion());
awsConfig.validateArn(awsConfigModel.getRoleArn());
config = awsConfig;
break;
Expand Down
2 changes: 2 additions & 0 deletions polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisCredentialProperty.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@ public enum PolarisCredentialProperty {
AWS_KEY_ID(String.class, "s3.access-key-id", "the aws access key id"),
AWS_SECRET_KEY(String.class, "s3.secret-access-key", "the aws access key secret"),
AWS_TOKEN(String.class, "s3.session-token", "the aws scoped access token"),
CLIENT_REGION(
String.class, "client.region", "region to configure client for making requests to AWS"),

GCS_ACCESS_TOKEN(String.class, "gcs.oauth2.token", "the gcs scoped access token"),
GCS_ACCESS_TOKEN_EXPIRES_AT(
Expand Down
1 change: 1 addition & 0 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 @@ -76,6 +76,7 @@ public EnumMap<PolarisCredentialProperty, String> getSubscopedCreds(
credentialMap.put(
PolarisCredentialProperty.AWS_SECRET_KEY, response.credentials().secretAccessKey());
credentialMap.put(PolarisCredentialProperty.AWS_TOKEN, response.credentials().sessionToken());
credentialMap.put(PolarisCredentialProperty.CLIENT_REGION, storageConfig.getRegion());
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: should we add client.region only if region is set?

return credentialMap;
}

Expand Down
22 changes: 19 additions & 3 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 @@ -49,23 +49,30 @@ public class AwsStorageConfigurationInfo extends PolarisStorageConfigurationInfo
@JsonProperty(value = "userARN")
private @Nullable String userARN = null;

/** User ARN for the service principal */
@JsonProperty(value = "region")
private @Nullable String region = null;

@JsonCreator
public AwsStorageConfigurationInfo(
@JsonProperty(value = "storageType", required = true) @NotNull StorageType storageType,
@JsonProperty(value = "allowedLocations", required = true) @NotNull
List<String> allowedLocations,
@JsonProperty(value = "roleARN", required = true) @NotNull String roleARN) {
this(storageType, allowedLocations, roleARN, null);
@JsonProperty(value = "roleARN", required = true) @NotNull String roleARN,
@JsonProperty(value = "region", required = false) @NotNull String region) {
this(storageType, allowedLocations, roleARN, null, region);
}

public AwsStorageConfigurationInfo(
@NotNull StorageType storageType,
@NotNull List<String> allowedLocations,
@NotNull String roleARN,
@Nullable String externalId) {
@Nullable String externalId,
@Nullable String region) {
super(storageType, allowedLocations);
this.roleARN = roleARN;
this.externalId = externalId;
this.region = region;
validateMaxAllowedLocations(MAX_ALLOWED_LOCATIONS);
}

Expand Down Expand Up @@ -107,6 +114,14 @@ public void setUserARN(@Nullable String userARN) {
this.userARN = userARN;
}

public @Nullable String getRegion() {
return region;
}

public void setRegion(@Nullable String region) {
this.region = region;
}

@Override
public String toString() {
return MoreObjects.toStringHelper(this)
Expand All @@ -116,6 +131,7 @@ public String toString() {
.add("userARN", userARN)
.add("externalId", externalId)
.add("allowedLocation", getAllowedLocations())
.add("region", region)
.toString();
}
}
9 changes: 6 additions & 3 deletions ...is-core/src/test/java/org/apache/polaris/core/storage/InMemoryStorageIntegrationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,8 @@ public void testValidateAccessToLocations() {
"s3://bucket/path/to/warehouse",
"s3://bucket/anotherpath/to/warehouse",
"s3://bucket2/warehouse/"),
"arn:aws:iam::012345678901:role/jdoe"),
"arn:aws:iam::012345678901:role/jdoe",
"us-east-2"),
Set.of(PolarisStorageActions.READ),
Set.of(
"s3://bucket/path/to/warehouse/namespace/table",
Expand Down Expand Up @@ -136,7 +137,8 @@ public void testValidateAccessToLocationsNoAllowedLocations() {
new AwsStorageConfigurationInfo(
PolarisStorageConfigurationInfo.StorageType.S3,
List.of(),
"arn:aws:iam::012345678901:role/jdoe"),
"arn:aws:iam::012345678901:role/jdoe",
"us-east-2"),
Set.of(PolarisStorageActions.READ),
Set.of(
"s3://bucket/path/to/warehouse/namespace/table",
Expand Down Expand Up @@ -169,7 +171,8 @@ public void testValidateAccessToLocationsWithPrefixOfAllowedLocation() {
new AwsStorageConfigurationInfo(
PolarisStorageConfigurationInfo.StorageType.S3,
List.of("s3://bucket/path/to/warehouse"),
"arn:aws:iam::012345678901:role/jdoe"),
"arn:aws:iam::012345678901:role/jdoe",
"us-east-2"),
Set.of(PolarisStorageActions.READ),
// trying to read a prefix under the allowed location
Set.of("s3://bucket/path/to"));
Expand Down
52 changes: 47 additions & 5 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 @@ -83,7 +83,8 @@ public void testGetSubscopedCreds() {
PolarisStorageConfigurationInfo.StorageType.S3,
List.of(warehouseDir),
roleARN,
externalId),
externalId,
null),
true,
Set.of(warehouseDir + "/namespace/table"),
Set.of(warehouseDir + "/namespace/table"));
Expand Down Expand Up @@ -217,7 +218,11 @@ public void testGetSubscopedCredsInlinePolicy(String awsPartition) {
.getSubscopedCreds(
Mockito.mock(PolarisDiagnostics.class),
new AwsStorageConfigurationInfo(
storageType, List.of(s3Path(bucket, warehouseKeyPrefix)), roleARN, externalId),
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)));
Expand Down Expand Up @@ -310,7 +315,8 @@ public void testGetSubscopedCredsInlinePolicyWithoutList() {
PolarisStorageConfigurationInfo.StorageType.S3,
List.of(s3Path(bucket, warehouseKeyPrefix)),
roleARN,
externalId),
externalId,
"us-east-2"),
false, /* allowList = false*/
Set.of(s3Path(bucket, firstPath), s3Path(bucket, secondPath)),
Set.of(s3Path(bucket, firstPath)));
Expand Down Expand Up @@ -398,7 +404,11 @@ public void testGetSubscopedCredsInlinePolicyWithoutWrites() {
.getSubscopedCreds(
Mockito.mock(PolarisDiagnostics.class),
new AwsStorageConfigurationInfo(
storageType, List.of(s3Path(bucket, warehouseKeyPrefix)), roleARN, externalId),
storageType,
List.of(s3Path(bucket, warehouseKeyPrefix)),
roleARN,
externalId,
"us-east-2"),
true, /* allowList = true */
Set.of(s3Path(bucket, firstPath), s3Path(bucket, secondPath)),
Set.of());
Expand Down Expand Up @@ -459,7 +469,8 @@ public void testGetSubscopedCredsInlinePolicyWithEmptyReadAndWrite() {
PolarisStorageConfigurationInfo.StorageType.S3,
List.of(s3Path(bucket, warehouseKeyPrefix)),
roleARN,
externalId),
externalId,
"us-east-2"),
true, /* allowList = true */
Set.of(),
Set.of());
Expand All @@ -470,6 +481,37 @@ public void testGetSubscopedCredsInlinePolicyWithEmptyReadAndWrite() {
.containsEntry(PolarisCredentialProperty.AWS_SECRET_KEY, "secretKey");
}

@Test
public void testClientRegion() {
StsClient stsClient = Mockito.mock(StsClient.class);
String roleARN = "arn:aws:iam::012345678901:role/jdoe";
String externalId = "externalId";
String bucket = "bucket";
String warehouseKeyPrefix = "path/to/warehouse";
String clientRegion = "test-region";
Mockito.when(stsClient.assumeRole(Mockito.isA(AssumeRoleRequest.class)))
.thenAnswer(
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);
}

private static @NotNull String s3Arn(String partition, String bucket, String keyPrefix) {
String bucketArn = "arn:" + partition + ":s3:::" + bucket;
if (keyPrefix == null) {
Expand Down
4 changes: 4 additions & 0 deletions spec/polaris-management-service.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -898,6 +898,10 @@ components:
type: string
description: the aws user arn used to assume the aws role
example: "arn:aws:iam::123456789001:user/abc1-b-self1234"
region:
type: string
description: the aws region where data is stored
example: "us-east-2"
required:
- roleArn

Expand Down