diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheEntry.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheEntry.java index ae799457fd..2649ee99c3 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheEntry.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheEntry.java @@ -23,6 +23,7 @@ import java.util.Map; import org.apache.polaris.core.persistence.dao.entity.ScopedCredentialsResult; import org.apache.polaris.core.storage.PolarisCredentialProperty; +import org.apache.polaris.core.storage.azure.AzureLocation; /** A storage credential cached entry. */ public class StorageCredentialCacheEntry { @@ -51,23 +52,47 @@ public long getExpirationTime() { return Long.MAX_VALUE; } + /** + * Azure needs special handling, the credential key is dynamically generated based on the storage + * account endpoint + */ + private void handleAzureCredential( + HashMap results, PolarisCredentialProperty credentialProperty, String value) { + if (credentialProperty.equals(PolarisCredentialProperty.AZURE_SAS_TOKEN)) { + String host = credsMap.get(PolarisCredentialProperty.AZURE_ACCOUNT_HOST); + results.put(credentialProperty.getPropertyName() + host, value); + + // Iceberg 1.7.x may expect the credential key to _not_ be suffixed with endpoint + if (host.endsWith(AzureLocation.ADLS_ENDPOINT)) { + int suffixIndex = host.lastIndexOf(AzureLocation.ADLS_ENDPOINT) - 1; + if (suffixIndex > 0) { + String withSuffixStripped = host.substring(0, suffixIndex); + results.put(credentialProperty.getPropertyName() + withSuffixStripped, value); + } + } + + if (host.endsWith(AzureLocation.BLOB_ENDPOINT)) { + int suffixIndex = host.lastIndexOf(AzureLocation.BLOB_ENDPOINT) - 1; + if (suffixIndex > 0) { + String withSuffixStripped = host.substring(0, suffixIndex); + results.put(credentialProperty.getPropertyName() + withSuffixStripped, value); + } + } + } + } + /** * Get the map of string creds that is needed for the query engine. * * @return a map of string representing the subscoped creds info. */ public Map convertToMapOfString() { - Map resCredsMap = new HashMap<>(); + HashMap resCredsMap = new HashMap<>(); if (!credsMap.isEmpty()) { credsMap.forEach( (key, value) -> { - // only Azure needs special handle, the target key is dynamically with storageaccount - // endpoint appended if (key.equals(PolarisCredentialProperty.AZURE_SAS_TOKEN)) { - resCredsMap.put( - key.getPropertyName() - + credsMap.get(PolarisCredentialProperty.AZURE_ACCOUNT_HOST), - value); + handleAzureCredential(resCredsMap, key, value); } else if (!key.equals(PolarisCredentialProperty.AZURE_ACCOUNT_HOST)) { resCredsMap.put(key.getPropertyName(), value); } diff --git a/polaris-core/src/test/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheTest.java b/polaris-core/src/test/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheTest.java index f3b4fc2f62..54c31cf101 100644 --- a/polaris-core/src/test/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheTest.java @@ -45,6 +45,7 @@ import org.apache.polaris.core.persistence.transactional.TreeMapMetaStore; import org.apache.polaris.core.persistence.transactional.TreeMapTransactionalPersistenceImpl; import org.apache.polaris.core.storage.PolarisCredentialProperty; +import org.apache.polaris.core.storage.azure.AzureLocation; import org.assertj.core.api.Assertions; import org.junit.jupiter.api.RepeatedTest; import org.junit.jupiter.api.Test; @@ -438,4 +439,93 @@ private static List getPolarisEntities() { return Arrays.asList(polarisEntity1, polarisEntity2, polarisEntity3); } + + @Test + public void testAzureCredentialFormatting() { + storageCredentialCache = new StorageCredentialCache(); + List mockedScopedCreds = + List.of( + new ScopedCredentialsResult( + new EnumMap<>( + ImmutableMap.builder() + .put(PolarisCredentialProperty.AZURE_SAS_TOKEN, "sas_token_azure_1") + .put(PolarisCredentialProperty.AZURE_ACCOUNT_HOST, "some_account") + .put( + PolarisCredentialProperty.EXPIRATION_TIME, + String.valueOf(Long.MAX_VALUE)) + .buildOrThrow())), + new ScopedCredentialsResult( + new EnumMap<>( + ImmutableMap.builder() + .put(PolarisCredentialProperty.AZURE_SAS_TOKEN, "sas_token_azure_2") + .put( + PolarisCredentialProperty.AZURE_ACCOUNT_HOST, + "some_account." + AzureLocation.ADLS_ENDPOINT) + .put( + PolarisCredentialProperty.EXPIRATION_TIME, + String.valueOf(Long.MAX_VALUE)) + .buildOrThrow())), + new ScopedCredentialsResult( + new EnumMap<>( + ImmutableMap.builder() + .put(PolarisCredentialProperty.AZURE_SAS_TOKEN, "sas_token_azure_3") + .put( + PolarisCredentialProperty.AZURE_ACCOUNT_HOST, + "some_account." + AzureLocation.BLOB_ENDPOINT) + .put( + PolarisCredentialProperty.EXPIRATION_TIME, + String.valueOf(Long.MAX_VALUE)) + .buildOrThrow()))); + + Mockito.when( + metaStoreManager.getSubscopedCredsForEntity( + Mockito.any(), + Mockito.anyLong(), + Mockito.anyLong(), + Mockito.any(), + Mockito.anyBoolean(), + Mockito.anySet(), + Mockito.anySet())) + .thenReturn(mockedScopedCreds.get(0)) + .thenReturn(mockedScopedCreds.get(1)) + .thenReturn(mockedScopedCreds.get(2)); + List entityList = getPolarisEntities(); + + Map noSuffixResult = + storageCredentialCache.getOrGenerateSubScopeCreds( + metaStoreManager, + callCtx, + entityList.get(0), + true, + new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), + new HashSet<>(Arrays.asList("s3://bucket3/path", "s3://bucket4/path"))); + Assertions.assertThat(noSuffixResult.size()).isEqualTo(2); + Assertions.assertThat(noSuffixResult).containsKey("adls.sas-token.some_account"); + + Map adlsSuffixResult = + storageCredentialCache.getOrGenerateSubScopeCreds( + metaStoreManager, + callCtx, + entityList.get(1), + true, + new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), + new HashSet<>(Arrays.asList("s3://bucket3/path", "s3://bucket4/path"))); + Assertions.assertThat(adlsSuffixResult.size()).isEqualTo(3); + Assertions.assertThat(adlsSuffixResult).containsKey("adls.sas-token.some_account"); + Assertions.assertThat(adlsSuffixResult) + .containsKey("adls.sas-token.some_account." + AzureLocation.ADLS_ENDPOINT); + + Map blobSuffixResult = + storageCredentialCache.getOrGenerateSubScopeCreds( + metaStoreManager, + callCtx, + entityList.get(2), + true, + new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), + new HashSet<>(Arrays.asList("s3://bucket3/path", "s3://bucket4/path"))); + Assertions.assertThat(blobSuffixResult.size()).isEqualTo(3); + Assertions.assertThat(blobSuffixResult).containsKey("adls.sas-token.some_account"); + Assertions.assertThat(blobSuffixResult) + .containsKey("adls.sas-token.some_account." + AzureLocation.BLOB_ENDPOINT); + } }