Store AWS region in AwsStorageConfigurationInfo (apache#455)

* debugging * working * fix * fix more invocations * autolint * patch one thing * add test * autolint * null check * autolint * remove notnull
cgpoh · Dec 12, 2024 · f9a24d4 · f9a24d4
1 parent 85f0beb
commit f9a24d4
Show file tree

Hide file tree

Showing 7 changed files with 112 additions and 12 deletions.
diff --git a/polaris-core/src/main/java/org/apache/polaris/core/entity/CatalogEntity.java b/polaris-core/src/main/java/org/apache/polaris/core/entity/CatalogEntity.java
@@ -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) {
@@ -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;

diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisCredentialProperty.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisCredentialProperty.java
@@ -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(

diff --git a/...e/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java b/...e/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java
@@ -76,6 +76,9 @@ public EnumMap<PolarisCredentialProperty, String> getSubscopedCreds(
     credentialMap.put(
         PolarisCredentialProperty.AWS_SECRET_KEY, response.credentials().secretAccessKey());
     credentialMap.put(PolarisCredentialProperty.AWS_TOKEN, response.credentials().sessionToken());
+    if (storageConfig.getRegion() != null) {
+      credentialMap.put(PolarisCredentialProperty.CLIENT_REGION, storageConfig.getRegion());
+    }
     return credentialMap;
   }
 

diff --git a/...s-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java b/...s-core/src/main/java/org/apache/polaris/core/storage/aws/AwsStorageConfigurationInfo.java
@@ -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) @Nonnull StorageType storageType,
       @JsonProperty(value = "allowedLocations", required = true) @Nonnull
           List<String> allowedLocations,
-      @JsonProperty(value = "roleARN", required = true) @Nonnull String roleARN) {
-    this(storageType, allowedLocations, roleARN, null);
+      @JsonProperty(value = "roleARN", required = true) @Nonnull String roleARN,
+      @JsonProperty(value = "region", required = false) @Nullable String region) {
+    this(storageType, allowedLocations, roleARN, null, region);
   }
 
   public AwsStorageConfigurationInfo(
       @Nonnull StorageType storageType,
       @Nonnull List<String> allowedLocations,
       @Nonnull 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);
   }
 
@@ -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)
@@ -116,6 +131,7 @@ public String toString() {
         .add("userARN", userARN)
         .add("externalId", externalId)
         .add("allowedLocation", getAllowedLocations())
+        .add("region", region)
         .toString();
   }
 }
diff --git a/...is-core/src/test/java/org/apache/polaris/core/storage/InMemoryStorageIntegrationTest.java b/...is-core/src/test/java/org/apache/polaris/core/storage/InMemoryStorageIntegrationTest.java
@@ -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",
@@ -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",
@@ -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"));

diff --git a/...est/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java b/...est/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java
@@ -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"));
@@ -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)));
@@ -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)));
@@ -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());
@@ -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());
@@ -470,6 +481,65 @@ 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);
+  }
+
+  @Test
+  public void testNoClientRegion() {
+    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";
+    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,
+                    null),
+                true, /* allowList = true */
+                Set.of(),
+                Set.of());
+    assertThat(credentials).isNotEmpty().doesNotContainKey(PolarisCredentialProperty.CLIENT_REGION);
+  }
+
   private static @Nonnull String s3Arn(String partition, String bucket, String keyPrefix) {
     String bucketArn = "arn:" + partition + ":s3:::" + bucket;
     if (keyPrefix == null) {

diff --git a/spec/polaris-management-service.yml b/spec/polaris-management-service.yml
@@ -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