Skip to content

Remove getCurrentContext from InMemoryStorageIntegration #2201

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dimas-b merged 1 commit into from
Jul 30, 2025
Merged

Remove getCurrentContext from InMemoryStorageIntegration #2201

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 9 additions & 11 deletions polaris-core/src/main/java/org/apache/polaris/core/storage/InMemoryStorageIntegration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
*/
Expand All @@ -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) {
Expand All @@ -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()
Expand Down Expand Up @@ -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);
}
}
5 changes: 4 additions & 1 deletion polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisStorageIntegration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

/**
Expand Down Expand Up @@ -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;
Expand Down
52 changes: 32 additions & 20 deletions ...is-core/src/test/java/org/apache/polaris/core/storage/InMemoryStorageIntegrationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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"),
Expand Down Expand Up @@ -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;
}
}
}
5 changes: 4 additions & 1 deletion service/common/src/main/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalog.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
2 changes: 2 additions & 0 deletions ...c/main/java/org/apache/polaris/service/storage/PolarisStorageIntegrationProviderImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down