adnanhemani commented on code in PR #3224:
URL: https://github.com/apache/polaris/pull/3224#discussion_r2600652919
##########
polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java:
##########
@@ -97,6 +124,54 @@ public void testGetSubscopedCreds(String scheme) {
true,
Set.of(warehouseDir + "/namespace/table"),
Set.of(warehouseDir + "/namespace/table"),
+ POLARIS_PRINCIPAL,
+ Optional.of("/namespace/table/credentials"));
+ assertThat(storageAccessConfig.credentials())
+ .isNotEmpty()
+ .containsEntry(StorageAccessProperty.AWS_TOKEN.getPropertyName(),
"sess")
+ .containsEntry(StorageAccessProperty.AWS_KEY_ID.getPropertyName(),
"accessKey")
+ .containsEntry(StorageAccessProperty.AWS_SECRET_KEY.getPropertyName(),
"secretKey")
+ .containsEntry(
+
StorageAccessProperty.AWS_SESSION_TOKEN_EXPIRES_AT_MS.getPropertyName(),
+ String.valueOf(EXPIRE_TIME.toEpochMilli()));
+ assertThat(storageAccessConfig.extraProperties())
+ .containsEntry(
+
StorageAccessProperty.AWS_REFRESH_CREDENTIALS_ENDPOINT.getPropertyName(),
+ "/namespace/table/credentials");
+ }
+
+ @Test
+ public void testGetSubscopedCredsWithNameInclude() {
+ StsClient stsClient = Mockito.mock(StsClient.class);
+ String roleARN = "arn:aws:iam::012345678901:role/jdoe";
+ String externalId = "externalId";
+
+ Mockito.when(stsClient.assumeRole(Mockito.isA(AssumeRoleRequest.class)))
+ .thenAnswer(
+ invocation -> {
+ assertThat(invocation.getArguments()[0])
+ .isInstanceOf(AssumeRoleRequest.class)
+
.asInstanceOf(InstanceOfAssertFactories.type(AssumeRoleRequest.class))
+ .returns(externalId, AssumeRoleRequest::externalId)
+ .returns(roleARN, AssumeRoleRequest::roleArn)
+ .returns("polaris-test-principal",
AssumeRoleRequest::roleSessionName);
+ return ASSUME_ROLE_RESPONSE;
+ });
+ String warehouseDir = "s3://bucket/path/to/warehouse";
+ StorageAccessConfig storageAccessConfig =
+ new AwsCredentialsStorageIntegration(
+ AwsStorageConfigurationInfo.builder()
+ .addAllowedLocation(warehouseDir)
+ .roleARN(roleARN)
+ .externalId(externalId)
+ .build(),
+ stsClient)
+ .getSubscopedCreds(
+ PRINCIPAL_INCLUDER_REALM_CONFIG,
Review Comment:
nit: Similar comment as above. I also had to play "spot-the-difference"
between this test and the one above it so a quick comment on this line would
make that very obvious to a reader that we are overriding behavior here since
it is not clear in the test name.
##########
polaris-core/src/test/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheTest.java:
##########
@@ -135,10 +144,133 @@ public void testCacheHit() {
true,
Set.of("s3://bucket1/path", "s3://bucket2/path"),
Set.of("s3://bucket3/path", "s3://bucket4/path"),
+ polarisPrincipal,
+ Optional.empty());
+
Assertions.assertThat(storageCredentialCache.getEstimatedSize()).isEqualTo(1);
+
+ Optional<PolarisPrincipal> emptyPrincipal = Optional.empty();
+
+ storageCredentialCache.getOrGenerateSubScopeCreds(
+ storageCredentialsVendor,
+ polarisEntity,
+ true,
+ Set.of("s3://bucket1/path", "s3://bucket2/path"),
+ Set.of("s3://bucket3/path", "s3://bucket4/path"),
+ polarisPrincipal,
Optional.empty());
Assertions.assertThat(storageCredentialCache.getEstimatedSize()).isEqualTo(1);
}
+ @Test
+ public void testCacheMissForAnotherPrincipal() {
+
+ List<ScopedCredentialsResult> mockedScopedCreds =
+ getFakeScopedCreds(3, /* expireSoon= */ false);
+ Mockito.when(
+ storageCredentialsVendor.getSubscopedCredsForEntity(
+ Mockito.any(),
+ Mockito.anyBoolean(),
+ Mockito.anySet(),
+ Mockito.anySet(),
+ Mockito.any(),
+ Mockito.any()))
+ .thenReturn(mockedScopedCreds.get(0))
+ .thenReturn(mockedScopedCreds.get(1))
+ .thenReturn(mockedScopedCreds.get(1));
+ PolarisBaseEntity baseEntity =
+ new PolarisBaseEntity(
+ 1, 2, PolarisEntityType.CATALOG,
PolarisEntitySubType.ICEBERG_TABLE, 0, "name");
+ PolarisEntity polarisEntity = new PolarisEntity(baseEntity);
+ PolarisPrincipal polarisPrincipal = PolarisPrincipal.of("principal",
Map.of(), Set.of());
+ PolarisPrincipal anotherPolarisPrincipal =
+ PolarisPrincipal.of("anotherPrincipal", Map.of(), Set.of());
+
+ // add an item to the cache
+ storageCredentialCache.getOrGenerateSubScopeCreds(
+ storageCredentialsVendor,
+ polarisEntity,
+ true,
+ Set.of("s3://bucket1/path", "s3://bucket2/path"),
+ Set.of("s3://bucket3/path", "s3://bucket4/path"),
+ polarisPrincipal,
+ Optional.empty());
+
Assertions.assertThat(storageCredentialCache.getEstimatedSize()).isEqualTo(1);
+
+ // subscope for the same entity and same allowed locations, will hit the
cache
+ storageCredentialCache.getOrGenerateSubScopeCreds(
+ storageCredentialsVendor,
+ polarisEntity,
+ true,
+ Set.of("s3://bucket1/path", "s3://bucket2/path"),
+ Set.of("s3://bucket3/path", "s3://bucket4/path"),
+ anotherPolarisPrincipal,
+ Optional.empty());
+
Assertions.assertThat(storageCredentialCache.getEstimatedSize()).isEqualTo(1);
+ }
+
+ @Test
+ public void testCacheHitForAnotherPrincipal() {
+ Mockito.when(storageCredentialsVendor.getRealmConfig())
Review Comment:
nit: Since this test case and the one above it have the same core logic, it
would be best to extract a helper to make these tests more DRY. But not a big
deal if there are no further changes being made to this PR, a simple refactor
later can take care of this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
