refactor of allowed locations check

Author: adutra

integration-tests/src/main/java/org/apache/polaris/service/it/test/PolarisS3RemoteSigningIntegrationTest.java

@@ -20,15 +20,20 @@
 
 import static org.apache.polaris.core.storage.StorageAccessProperty.AWS_ENDPOINT;
 import static org.apache.polaris.core.storage.StorageAccessProperty.AWS_PATH_STYLE_ACCESS;
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.assertj.core.api.InstanceOfAssertFactories.type;
 
 import com.google.common.collect.ImmutableMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import org.apache.iceberg.BaseTable;
 import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
 import org.apache.iceberg.TableMetadata;
 import org.apache.iceberg.TableMetadataParser;
+import org.apache.iceberg.Transaction;
 import org.apache.iceberg.catalog.Namespace;
 import org.apache.iceberg.catalog.TableIdentifier;
 import org.apache.iceberg.exceptions.ForbiddenException;
@@ -87,6 +92,11 @@ protected Optional<String> endpoint() {
     return Optional.empty();
   }
 
+  /**
+   * A set of allowed locations to include in the {@linkplain #getStorageConfigInfo() storage
+   * configuration info}. The first allowed location will serve as the base for the catalog default
+   * location.
+   */
   protected abstract List<String> allowedLocations();
 
   protected boolean stsUnavailable() {
@@ -136,6 +146,48 @@ public void testExternalCatalogRemoteSigningDisabled() {
     }
   }
 
+  @CatalogConfig(
+      properties = {
+        "polaris.config.default-table-location-object-storage-prefix.enabled",
+        "true",
+        "polaris.config.allow.overlapping.table.location",
+        "true"
+      })
+  @Test
+  void testCreateTableWithObjectStoragePrefix() {
+    @SuppressWarnings("resource")
+    RESTCatalog restCatalog = catalog();
+    restCatalog.createNamespace(NS);
+    // Only direct table creation is supported with object storage prefix
+    Table tbl1 = restCatalog.buildTable(TABLE, SCHEMA).create();
+    // Will trigger write sign requests for manifests and snapshots, using object storage prefix
+    tbl1.newFastAppend().appendFile(FILE_A).commit();
+    // Will trigger many read sign requests for metadata and manifests
+    assertFiles(tbl1, FILE_A);
+    assertThat(tbl1).isNotNull();
+  }
+
+  @CatalogConfig(properties = {"polaris.config.allow.unstructured.table.location", "true"})
+  @Test
+  public void testCreateTableWithCustomLocation() {
+    @SuppressWarnings("resource")
+    RESTCatalog restCatalog = catalog();
+    restCatalog.createNamespace(NS);
+    String customLocation = allowedLocations().getFirst() + "/custom/tbl1";
+    Transaction create =
+        restCatalog.buildTable(TABLE, SCHEMA).withLocation(customLocation).createTransaction();
+    // Will trigger write sign requests for manifests and snapshots, before the table is created
+    create.newFastAppend().appendFile(FILE_A).commit();
+    // Will trigger table creation, then many read sign requests for metadata and manifests
+    create.commitTransaction();
+    Table tbl1 = restCatalog.loadTable(TABLE);
+    assertFiles(tbl1, FILE_A);
+    assertThat(tbl1)
+        .isNotNull()
+        .asInstanceOf(type(BaseTable.class))
+        .returns(customLocation, BaseTable::location);
+  }
+
   @Test
   @Override
   @Disabled("It's not possible to request an access delegation mode when registering a table.")

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

@@ -19,6 +19,7 @@
 package org.apache.polaris.core.storage;
 
 import jakarta.annotation.Nonnull;
+import java.util.Collection;
 import java.util.List;
 import java.util.Set;
 import org.apache.iceberg.catalog.TableIdentifier;
@@ -41,7 +42,7 @@ public class LocationRestrictions {
    * <p>All locations in this list have been validated to conform to the storage type's URI scheme
    * requirements during construction.
    */
-  private final List<String> allowedLocations;
+  private final Collection<String> allowedLocations;
 
   /**
    * The parent location for structured table enforcement.
@@ -54,15 +55,23 @@ public class LocationRestrictions {
 
   public LocationRestrictions(
       @Nonnull PolarisStorageConfigurationInfo storageConfigurationInfo, String parentLocation) {
-    this.allowedLocations = storageConfigurationInfo.getAllowedLocations();
+    this(storageConfigurationInfo.getAllowedLocations(), parentLocation);
     allowedLocations.forEach(storageConfigurationInfo::validatePrefixForStorageType);
-    this.parentLocation = parentLocation;
   }
 
   public LocationRestrictions(@Nonnull PolarisStorageConfigurationInfo storageConfigurationInfo) {
     this(storageConfigurationInfo, null);
   }
 
+  public LocationRestrictions(Collection<String> allowedLocations, String parentLocation) {
+    this.allowedLocations = allowedLocations;
+    this.parentLocation = parentLocation;
+  }
+
+  public LocationRestrictions(Collection<String> allowedLocations) {
+    this(allowedLocations, null);
+  }
+
   /**
    * Validates that the requested storage locations are permitted for the given table identifier.
    *
@@ -94,7 +103,7 @@ public void validate(
 
   private void validateLocations(
       RealmConfig realmConfig,
-      List<String> allowedLocations,
+      Collection<String> allowedLocations,
       Set<String> requestLocations,
       TableIdentifier identifier) {
     var validationResults =

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

@@ -114,7 +114,7 @@ public static PolarisStorageConfigurationInfo deserialize(final @Nonnull String
     }
   }
 
-  public static Optional<LocationRestrictions> forEntityPath(
+  public static Optional<LocationRestrictions> getLocationRestrictionsForEntityPath(
       RealmConfig realmConfig, List<PolarisEntity> entityPath) {
     return findStorageInfoFromHierarchy(entityPath)
         .map(

runtime/service/src/intTest/java/org/apache/polaris/service/it/S3RemoteSigningMinIOIT.java

@@ -71,7 +71,9 @@ static void setup(
 
   @Override
   protected List<String> allowedLocations() {
-    return List.of(storageBase.toString());
+    return List.of(
+        storageBase.resolve(BUCKET_URI_PREFIX + "/allowed-location1").toString(),
+        storageBase.resolve(BUCKET_URI_PREFIX + "/allowed-location2").toString());
   }
 
   @Override

runtime/service/src/main/java/org/apache/polaris/service/catalog/common/CatalogUtils.java

@@ -71,7 +71,7 @@ public static void validateLocationsForTableLike(
       Set<String> locations,
       PolarisResolvedPathWrapper resolvedStorageEntity) {
 
-    PolarisStorageConfigurationInfo.forEntityPath(
+    PolarisStorageConfigurationInfo.getLocationRestrictionsForEntityPath(
             realmConfig, resolvedStorageEntity.getRawFullPath())
         .ifPresentOrElse(
             restrictions -> restrictions.validate(realmConfig, identifier, locations),

runtime/service/src/main/java/org/apache/polaris/service/storage/s3/sign/S3RemoteSigningCatalogHandler.java

@@ -18,13 +18,8 @@
  */
 package org.apache.polaris.service.storage.s3.sign;
 
-import java.util.Collection;
 import java.util.EnumSet;
-import java.util.Optional;
 import java.util.Set;
-import org.apache.iceberg.BaseTable;
-import org.apache.iceberg.Table;
-import org.apache.iceberg.catalog.Catalog;
 import org.apache.iceberg.catalog.TableIdentifier;
 import org.apache.iceberg.exceptions.ForbiddenException;
 import org.apache.polaris.core.PolarisDiagnostics;
@@ -35,17 +30,16 @@
 import org.apache.polaris.core.config.RealmConfig;
 import org.apache.polaris.core.context.CallContext;
 import org.apache.polaris.core.entity.CatalogEntity;
-import org.apache.polaris.core.entity.PolarisEntity;
-import org.apache.polaris.core.entity.PolarisEntityConstants;
+import org.apache.polaris.core.entity.PolarisEntitySubType;
+import org.apache.polaris.core.entity.PolarisEntityType;
+import org.apache.polaris.core.entity.table.IcebergTableLikeEntity;
+import org.apache.polaris.core.entity.table.TableLikeEntity;
 import org.apache.polaris.core.persistence.PolarisResolvedPathWrapper;
 import org.apache.polaris.core.persistence.resolver.ResolutionManifestFactory;
-import org.apache.polaris.core.storage.InMemoryStorageIntegration;
-import org.apache.polaris.core.storage.PolarisStorageActions;
-import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo;
+import org.apache.polaris.core.storage.LocationRestrictions;
 import org.apache.polaris.core.storage.StorageUtil;
 import org.apache.polaris.service.catalog.common.CatalogHandler;
 import org.apache.polaris.service.catalog.common.CatalogUtils;
-import org.apache.polaris.service.catalog.io.FileIOUtil;
 import org.apache.polaris.service.context.catalog.CallContextCatalogFactory;
 import org.apache.polaris.service.s3.sign.model.PolarisS3SignRequest;
 import org.apache.polaris.service.s3.sign.model.PolarisS3SignResponse;
@@ -60,7 +54,6 @@ public class S3RemoteSigningCatalogHandler extends CatalogHandler implements Aut
   private final S3RequestSigner s3RequestSigner;
 
   private CatalogEntity catalogEntity;
-  private Catalog baseCatalog;
 
   public S3RemoteSigningCatalogHandler(
       PolarisDiagnostics diagnostics,
@@ -92,7 +85,7 @@ protected void initializeCatalog() {
     if (catalogEntity.isExternal()) {
       throw new ForbiddenException("Cannot use S3 remote signing with federated catalogs.");
     }
-    baseCatalog = catalogFactory.createCallContextCatalog(resolutionManifest);
+    // no need to materialize the catalog here, as we only need the catalog entity
   }
 
   public PolarisS3SignResponse signS3Request(
@@ -107,76 +100,43 @@ public PolarisS3SignResponse signS3Request(
 
     authorizeRemoteSigningOrThrow(EnumSet.of(authzOp), tableIdentifier);
 
-    // Must be done after the authorization check, as the auth check creates the catalog entity
+    // Must be done after the authorization check, as the auth check creates the catalog entity;
+    // also, materializing the catalog here could hurt performance.
     throwIfRemoteSigningNotEnabled(callContext.getRealmConfig(), catalogEntity);
 
-    var result =
-        InMemoryStorageIntegration.validateAllowedLocations(
-            callContext.getRealmConfig(),
-            getAllowedLocations(tableIdentifier),
-            getStorageActions(s3SignRequest),
-            getTargetLocations(s3SignRequest));
-
-    if (result.values().stream().anyMatch(r -> r.values().stream().anyMatch(v -> !v.isSuccess()))) {
-      throw new ForbiddenException("Requested S3 location is not allowed.");
-    }
+    validateLocations(s3SignRequest, tableIdentifier);
 
     PolarisS3SignResponse s3SignResponse = s3RequestSigner.signRequest(s3SignRequest);
     LOGGER.debug("S3 signing response: {}", s3SignResponse);
 
     return s3SignResponse;
   }
 
-  // TODO M2 computing allowed locations is expensive. We should cache it.
-  private Collection<String> getAllowedLocations(TableIdentifier tableIdentifier) {
-
-    if (baseCatalog.tableExists(tableIdentifier)) {
-
-      // If the table exists, get allowed locations from the table metadata
-      Table table = baseCatalog.loadTable(tableIdentifier);
-      if (table instanceof BaseTable baseTable) {
-        return StorageUtil.getLocationsUsedByTable(baseTable.operations().current());
-      }
-
-      throw new ForbiddenException("No storage configuration found for table.");
+  private void validateLocations(
+      PolarisS3SignRequest s3SignRequest, TableIdentifier tableIdentifier) {
 
+    // Will point to the table entity if it exists, otherwise the namespace entity.
+    PolarisResolvedPathWrapper tableOrNamespace =
+        CatalogUtils.findResolvedStorageEntity(resolutionManifest, tableIdentifier);
+
+    Set<String> targetLocations = getTargetLocations(s3SignRequest);
+
+    // If the table exists already, validate the target locations against the table's locations;
+    // otherwise, validate against the namespace's locations using the entity path hierarchy.
+    if (tableOrNamespace.getRawLeafEntity().getType() == PolarisEntityType.TABLE_LIKE
+        && tableOrNamespace.getRawLeafEntity().getSubType() == PolarisEntitySubType.ICEBERG_TABLE) {
+      TableLikeEntity tableEntity = new IcebergTableLikeEntity(tableOrNamespace.getRawLeafEntity());
+      Set<String> allowedLocations =
+          StorageUtil.getLocationsUsedByTable(
+              tableEntity.getBaseLocation(), tableEntity.getPropertiesAsMap());
+      new LocationRestrictions(allowedLocations)
+          .validate(callContext.getRealmConfig(), tableIdentifier, targetLocations);
     } else {
-
-      // If the table or view doesn't exist, the engine might be writing the manifests before the
-      // table creation is committed. In this case, we still need to check allowed locations from
-      // the parent entities.
-
-      PolarisResolvedPathWrapper resolvedPath =
-          CatalogUtils.findResolvedStorageEntity(resolutionManifest, tableIdentifier);
-
-      Optional<PolarisEntity> storageInfo = FileIOUtil.findStorageInfoFromHierarchy(resolvedPath);
-
-      var configurationInfo =
-          storageInfo
-              .map(PolarisEntity::getInternalPropertiesAsMap)
-              .map(info -> info.get(PolarisEntityConstants.getStorageConfigInfoPropertyName()))
-              .map(PolarisStorageConfigurationInfo::deserialize);
-
-      if (configurationInfo.isEmpty()) {
-        throw new ForbiddenException("No storage configuration found for table.");
-      }
-
-      return configurationInfo.get().getAllowedLocations();
+      CatalogUtils.validateLocationsForTableLike(
+          callContext.getRealmConfig(), tableIdentifier, targetLocations, tableOrNamespace);
     }
   }
 
-  private Set<PolarisStorageActions> getStorageActions(PolarisS3SignRequest s3SignRequest) {
-    // TODO M2: better mapping of request URIs to storage actions.
-    // Disambiguate LIST vs READ or WRITE vs DELETE is not possible based on the HTTP method alone.
-    // Examples:
-    // - ListObjects is conceptually a LIST operation, and GetObject is conceptually a READ. But
-    // both requests use the GET method.
-    // - DeleteObject uses the DELETE method, but the DeleteObjects operation uses the POST method.
-    return s3SignRequest.write()
-        ? Set.of(PolarisStorageActions.WRITE)
-        : Set.of(PolarisStorageActions.READ);
-  }
-
   private Set<String> getTargetLocations(PolarisS3SignRequest s3SignRequest) {
     // TODO M2: map http URI to s3 URI
     return Set.of();

runtime/service/src/test/java/org/apache/polaris/service/it/S3RemoteSigningMinIOIntegrationTest.java

@@ -72,7 +72,9 @@ static void setup(
 
   @Override
   protected List<String> allowedLocations() {
-    return List.of(storageBase.toString());
+    return List.of(
+        storageBase.resolve(BUCKET_URI_PREFIX + "/allowed-location1").toString(),
+        storageBase.resolve(BUCKET_URI_PREFIX + "/allowed-location2").toString());
   }
 
   @Override