Remove getCurrentContext from InMemoryStorageIntegration (#2201)

Author: XN137

polaris-core/src/main/java/org/apache/polaris/core/storage/InMemoryStorageIntegration.java

@@ -22,18 +22,17 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
-import org.apache.polaris.core.context.CallContext;
+import org.apache.polaris.core.config.RealmConfig;
 
 /**
  * Base class for in-memory implementations of {@link PolarisStorageIntegration}. A basic
- * implementation of {@link #validateAccessToLocations(PolarisStorageConfigurationInfo, Set, Set)}
- * is provided that checks to see that the list of locations being accessed is among the list of
- * {@link PolarisStorageConfigurationInfo#getAllowedLocations()}. Locations being accessed must be
- * equal to or a subdirectory of at least one of the allowed locations.
+ * implementation of {@link #validateAccessToLocations(RealmConfig, PolarisStorageConfigurationInfo,
+ * Set, Set)} is provided that checks to see that the list of locations being accessed is among the
+ * list of {@link PolarisStorageConfigurationInfo#getAllowedLocations()}. Locations being accessed
+ * must be equal to or a subdirectory of at least one of the allowed locations.
  *
  * @param <T>
  */
@@ -56,6 +55,7 @@ public InMemoryStorageIntegration(String identifierOrId) {
    */
   public static Map<String, Map<PolarisStorageActions, ValidationResult>>
       validateSubpathsOfAllowedLocations(
+          @Nonnull RealmConfig realmConfig,
           @Nonnull PolarisStorageConfigurationInfo storageConfig,
           @Nonnull Set<PolarisStorageActions> actions,
           @Nonnull Set<String> locations) {
@@ -76,10 +76,7 @@ public InMemoryStorageIntegration(String identifierOrId) {
     List<StorageLocation> allowedLocations =
         allowedLocationStrings.stream().map(StorageLocation::of).collect(Collectors.toList());
 
-    boolean allowWildcardLocation =
-        Optional.ofNullable(CallContext.getCurrentContext())
-            .map(ctx -> ctx.getRealmConfig().getConfig("ALLOW_WILDCARD_LOCATION", false))
-            .orElse(false);
+    boolean allowWildcardLocation = realmConfig.getConfig("ALLOW_WILDCARD_LOCATION", false);
 
     if (allowWildcardLocation && allowedLocationStrings.contains("*")) {
       return locations.stream()
@@ -122,9 +119,10 @@ public InMemoryStorageIntegration(String identifierOrId) {
   @Override
   @Nonnull
   public Map<String, Map<PolarisStorageActions, ValidationResult>> validateAccessToLocations(
+      @Nonnull RealmConfig realmConfig,
       @Nonnull T storageConfig,
       @Nonnull Set<PolarisStorageActions> actions,
       @Nonnull Set<String> locations) {
-    return validateSubpathsOfAllowedLocations(storageConfig, actions, locations);
+    return validateSubpathsOfAllowedLocations(realmConfig, storageConfig, actions, locations);
   }
 }

polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisStorageIntegration.java

@@ -21,6 +21,7 @@
 import jakarta.annotation.Nonnull;
 import java.util.Map;
 import java.util.Set;
+import org.apache.polaris.core.config.RealmConfig;
 import org.apache.polaris.core.context.CallContext;
 
 /**
@@ -93,12 +94,14 @@ public abstract AccessConfig getSubscopedCreds(
   @Nonnull
   public abstract Map<String, Map<PolarisStorageActions, ValidationResult>>
       validateAccessToLocations(
+          @Nonnull RealmConfig realmConfig,
           @Nonnull T storageConfig,
           @Nonnull Set<PolarisStorageActions> actions,
           @Nonnull Set<String> locations);
 
   /**
-   * Result of calling {@link #validateAccessToLocations(PolarisStorageConfigurationInfo, Set, Set)}
+   * Result of calling {@link #validateAccessToLocations(RealmConfig,
+   * PolarisStorageConfigurationInfo, Set, Set)}
    */
   public static final class ValidationResult {
     private final boolean success;

polaris-core/src/test/java/org/apache/polaris/core/storage/InMemoryStorageIntegrationTest.java

@@ -20,13 +20,12 @@
 
 import jakarta.annotation.Nonnull;
 import jakarta.annotation.Nullable;
-import java.time.Clock;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import org.apache.polaris.core.PolarisCallContext;
-import org.apache.polaris.core.PolarisDefaultDiagServiceImpl;
 import org.apache.polaris.core.config.PolarisConfigurationStore;
+import org.apache.polaris.core.config.RealmConfig;
+import org.apache.polaris.core.config.RealmConfigImpl;
 import org.apache.polaris.core.context.CallContext;
 import org.apache.polaris.core.context.RealmContext;
 import org.apache.polaris.core.storage.aws.AwsStorageConfigurationInfo;
@@ -35,16 +34,20 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.CsvSource;
 import org.junit.jupiter.params.provider.ValueSource;
-import org.mockito.Mockito;
 
 class InMemoryStorageIntegrationTest {
 
+  private static final RealmContext REALM_CONTEXT = () -> "realm";
+
   @ParameterizedTest
   @CsvSource({"s3,s3", "s3,s3a", "s3a,s3", "s3a,s3a"})
   public void testValidateAccessToLocations(String allowedScheme, String locationScheme) {
+    RealmConfig realmConfig =
+        new RealmConfigImpl(new MockedConfigurationStore(Map.of()), REALM_CONTEXT);
     MockInMemoryStorageIntegration storage = new MockInMemoryStorageIntegration();
     Map<String, Map<PolarisStorageActions, PolarisStorageIntegration.ValidationResult>> result =
         storage.validateAccessToLocations(
+            realmConfig,
             new AwsStorageConfigurationInfo(
                 PolarisStorageConfigurationInfo.StorageType.S3,
                 List.of(
@@ -96,24 +99,12 @@ public void testAwsAccountIdParsing() {
   @ValueSource(strings = {"s3", "s3a"})
   public void testValidateAccessToLocationsWithWildcard(String s3Scheme) {
     MockInMemoryStorageIntegration storage = new MockInMemoryStorageIntegration();
-    Map<String, Boolean> config = Map.of("ALLOW_WILDCARD_LOCATION", true);
-    PolarisCallContext polarisCallContext =
-        new PolarisCallContext(
-            () -> "testRealm",
-            Mockito.mock(),
-            new PolarisDefaultDiagServiceImpl(),
-            new PolarisConfigurationStore() {
-              @SuppressWarnings("unchecked")
-              @Override
-              public <T> @Nullable T getConfiguration(
-                  @Nonnull RealmContext ctx, String configName) {
-                return (T) config.get(configName);
-              }
-            },
-            Clock.systemUTC());
-    CallContext.setCurrentContext(polarisCallContext);
+    Map<String, Object> config = Map.of("ALLOW_WILDCARD_LOCATION", true);
+    RealmConfig realmConfig =
+        new RealmConfigImpl(new MockedConfigurationStore(config), REALM_CONTEXT);
     Map<String, Map<PolarisStorageActions, PolarisStorageIntegration.ValidationResult>> result =
         storage.validateAccessToLocations(
+            realmConfig,
             new FileStorageConfigurationInfo(List.of("file://", "*")),
             Set.of(PolarisStorageActions.READ),
             Set.of(
@@ -151,8 +142,11 @@ public void testValidateAccessToLocationsWithWildcard(String s3Scheme) {
   @Test
   public void testValidateAccessToLocationsNoAllowedLocations() {
     MockInMemoryStorageIntegration storage = new MockInMemoryStorageIntegration();
+    RealmConfig realmConfig =
+        new RealmConfigImpl(new MockedConfigurationStore(Map.of()), REALM_CONTEXT);
     Map<String, Map<PolarisStorageActions, PolarisStorageIntegration.ValidationResult>> result =
         storage.validateAccessToLocations(
+            realmConfig,
             new AwsStorageConfigurationInfo(
                 PolarisStorageConfigurationInfo.StorageType.S3,
                 List.of(),
@@ -185,8 +179,11 @@ public void testValidateAccessToLocationsNoAllowedLocations() {
   @Test
   public void testValidateAccessToLocationsWithPrefixOfAllowedLocation() {
     MockInMemoryStorageIntegration storage = new MockInMemoryStorageIntegration();
+    RealmConfig realmConfig =
+        new RealmConfigImpl(new MockedConfigurationStore(Map.of()), REALM_CONTEXT);
     Map<String, Map<PolarisStorageActions, PolarisStorageIntegration.ValidationResult>> result =
         storage.validateAccessToLocations(
+            realmConfig,
             new AwsStorageConfigurationInfo(
                 PolarisStorageConfigurationInfo.StorageType.S3,
                 List.of("s3://bucket/path/to/warehouse"),
@@ -220,4 +217,19 @@ public AccessConfig getSubscopedCreds(
       return null;
     }
   }
+
+  private static class MockedConfigurationStore implements PolarisConfigurationStore {
+    private final Map<String, Object> defaults;
+
+    public MockedConfigurationStore(Map<String, Object> defaults) {
+      this.defaults = Map.copyOf(defaults);
+    }
+
+    @Override
+    public <T> @Nullable T getConfiguration(@Nonnull RealmContext realmContext, String configName) {
+      @SuppressWarnings("unchecked")
+      T confgValue = (T) defaults.get(configName);
+      return confgValue;
+    }
+  }
 }

service/common/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalog.java

@@ -1040,7 +1040,10 @@ private void validateLocationsForTableLike(
           Map<String, Map<PolarisStorageActions, PolarisStorageIntegration.ValidationResult>>
               validationResults =
                   InMemoryStorageIntegration.validateSubpathsOfAllowedLocations(
-                      storageConfigInfo, Set.of(PolarisStorageActions.ALL), locations);
+                      callContext.getRealmConfig(),
+                      storageConfigInfo,
+                      Set.of(PolarisStorageActions.ALL),
+                      locations);
           validationResults
               .values()
               .forEach(

service/common/src/main/java/org/apache/polaris/service/storage/PolarisStorageIntegrationProviderImpl.java

@@ -30,6 +30,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Supplier;
+import org.apache.polaris.core.config.RealmConfig;
 import org.apache.polaris.core.context.CallContext;
 import org.apache.polaris.core.storage.AccessConfig;
 import org.apache.polaris.core.storage.PolarisStorageActions;
@@ -110,6 +111,7 @@ public AccessConfig getSubscopedCreds(
               @Override
               public @Nonnull Map<String, Map<PolarisStorageActions, ValidationResult>>
                   validateAccessToLocations(
+                      @Nonnull RealmConfig realmConfig,
                       @Nonnull T storageConfig,
                       @Nonnull Set<PolarisStorageActions> actions,
                       @Nonnull Set<String> locations) {