Vend Azure credentials compatible with Iceberg 1.7

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<String, String> 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<String, String> convertToMapOfString() {
-    Map<String, String> resCredsMap = new HashMap<>();
+    HashMap<String, String> 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);
             }

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<PolarisEntity> getPolarisEntities() {
 
     return Arrays.asList(polarisEntity1, polarisEntity2, polarisEntity3);
   }
+
+  @Test
+  public void testAzureCredentialFormatting() {
+    storageCredentialCache = new StorageCredentialCache();
+    List<ScopedCredentialsResult> mockedScopedCreds =
+        List.of(
+            new ScopedCredentialsResult(
+                new EnumMap<>(
+                    ImmutableMap.<PolarisCredentialProperty, String>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.<PolarisCredentialProperty, String>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.<PolarisCredentialProperty, String>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<PolarisEntity> entityList = getPolarisEntities();
+
+    Map<String, String> 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<String, String> 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<String, String> 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);
+  }
 }