Support encrypted table operations and server-side purge

polaris-core/src/main/java/org/apache/polaris/core/entity/PolarisTaskConstants.java

@@ -26,4 +26,7 @@ public class PolarisTaskConstants {
   public static final String TASK_DATA = "data";
   public static final String TASK_TYPE = "taskType";
   public static final String STORAGE_LOCATION = "storageLocation";
+
+  /** Serialized encryption context used by an Iceberg cleanup task. */
+  public static final String ENCRYPTION_CONTEXT = "encryptionContext";
 }

runtime/service/src/main/java/org/apache/polaris/service/catalog/iceberg/LocalIcebergCatalog.java

@@ -65,7 +65,12 @@
 import org.apache.iceberg.catalog.Namespace;
 import org.apache.iceberg.catalog.SupportsNamespaces;
 import org.apache.iceberg.catalog.TableIdentifier;
+import org.apache.iceberg.encryption.EncryptingFileIO;
 import org.apache.iceberg.encryption.EncryptionManager;
+import org.apache.iceberg.encryption.EncryptionUtil;
+import org.apache.iceberg.encryption.KeyManagementClient;
+import org.apache.iceberg.encryption.PlaintextEncryptionManager;
+import org.apache.iceberg.encryption.StandardEncryptionManager;
 import org.apache.iceberg.exceptions.AlreadyExistsException;
 import org.apache.iceberg.exceptions.BadRequestException;
 import org.apache.iceberg.exceptions.CommitFailedException;
@@ -136,6 +141,7 @@
 import org.apache.polaris.service.catalog.common.CatalogUtils;
 import org.apache.polaris.service.catalog.common.LocationUtils;
 import org.apache.polaris.service.catalog.io.FileIOFactory;
+import org.apache.polaris.service.catalog.io.PolarisEncryptionUtil;
 import org.apache.polaris.service.catalog.io.StorageAccessConfigProvider;
 import org.apache.polaris.service.catalog.validation.IcebergPropertiesValidation;
 import org.apache.polaris.service.events.EventAttributeMap;
@@ -191,6 +197,7 @@ public class LocalIcebergCatalog extends BaseMetastoreViewCatalog
   private FileIO catalogFileIO;
   private CloseableGroup closeableGroup;
   private Map<String, String> tableDefaultProperties;
+  private KeyManagementClient keyManagementClient;
 
   private final String catalogName;
   private final long catalogId;
@@ -321,6 +328,11 @@ public void initialize(String name, Map<String, String> properties) {
 
     this.closeableGroup = new CloseableGroup();
     closeableGroup.addCloseable(metricsReporter());
+    if (properties.containsKey(CatalogProperties.ENCRYPTION_KMS_TYPE)
+        || properties.containsKey(CatalogProperties.ENCRYPTION_KMS_IMPL)) {
+      keyManagementClient = EncryptionUtil.createKmsClient(properties);
+      closeableGroup.addCloseable(keyManagementClient);
+    }
     closeableGroup.setSuppressCloseFailure(true);
 
     tableDefaultProperties =
@@ -579,6 +591,8 @@ public boolean dropTable(TableIdentifier tableIdentifier, boolean purge) {
                   clone.put(CatalogProperties.FILE_IO_IMPL, ioImplClassName);
                   clone.putAll(properties);
                   clone.put(PolarisTaskConstants.STORAGE_LOCATION, lastMetadata.location());
+                  PolarisEncryptionUtil.addCleanupTaskEncryptionProperties(
+                      clone, catalogProperties, lastMetadata);
                   return clone;
                 })
             .orElse(Map.of());
@@ -1621,6 +1635,26 @@ public ViewBuilder withLocation(String newLocation) {
     }
   }
 
+  private record TableEncryptionConfig(String tableKeyId, int dataKeyLength) {
+    private static @Nullable TableEncryptionConfig from(TableMetadata metadata) {
+      if (metadata == null) {
+        return null;
+      }
+
+      String tableKeyId = metadata.properties().get(TableProperties.ENCRYPTION_TABLE_KEY);
+      if (tableKeyId == null) {
+        return null;
+      }
+
+      return new TableEncryptionConfig(
+          tableKeyId,
+          PropertyUtil.propertyAsInt(
+              metadata.properties(),
+              TableProperties.ENCRYPTION_DEK_LENGTH,
+              TableProperties.ENCRYPTION_DEK_LENGTH_DEFAULT));
+    }
+  }
+
   /**
    * An implementation of {@link TableOperations} that integrates with {@link LocalIcebergCatalog}.
    * Much of this code was originally copied from {@link
@@ -1635,6 +1669,9 @@ public class BasePolarisTableOperations extends PolarisOperationsBase<TableMetad
     private final boolean makeMetadataCurrentOnCommit;
 
     private FileIO tableFileIO;
+    private EncryptionManager encryptionManager;
+    private TableEncryptionConfig encryptionManagerConfig;
+    private EncryptingFileIO encryptingFileIO;
 
     BasePolarisTableOperations(
         FileIO defaultFileIO,
@@ -1670,6 +1707,7 @@ public TableMetadata refresh() {
         currentMetadata = null;
         currentMetadataLocation = null;
         version = -1;
+        resetEncryptionState();
         throw e;
       }
       return current();
@@ -1707,7 +1745,65 @@ public void commit(TableMetadata base, TableMetadata metadata) {
 
     @Override
     public FileIO io() {
-      return tableFileIO;
+      return io(currentMetadata);
+    }
+
+    private FileIO io(TableMetadata metadata) {
+      Preconditions.checkState(
+          tableFileIO != null, "FileIO has not been initialized for table %s", fullTableName);
+
+      EncryptionManager manager = encryption(metadata);
+      if (manager == PlaintextEncryptionManager.instance()) {
+        return tableFileIO;
+      }
+
+      if (encryptingFileIO == null) {
+        encryptingFileIO = EncryptingFileIO.combine(tableFileIO, manager);
+      }
+      return encryptingFileIO;
+    }
+
+    @Override
+    public EncryptionManager encryption() {
+      return encryption(currentMetadata);
+    }
+
+    private EncryptionManager encryption(TableMetadata metadata) {
+      TableEncryptionConfig requestedConfig = TableEncryptionConfig.from(metadata);
+      if (encryptionManager != null) {
+        // A successful new-table commit may leave currentMetadata null until the next refresh, but
+        // transaction cleanup still calls the public io() method. Preserve the manager that was
+        // bound to the committed metadata in that case. A non-null plaintext metadata object is a
+        // real configuration change and is rejected below.
+        if (metadata == null) {
+          return encryptionManager;
+        }
+
+        Preconditions.checkState(
+            requestedConfig != null && requestedConfig.equals(encryptionManagerConfig),
+            "Cannot reuse encryption manager for table %s: initialized with %s but metadata requires %s",
+            fullTableName,
+            encryptionManagerConfig,
+            requestedConfig == null ? "plaintext" : requestedConfig);
+        return encryptionManager;
+      }
+
+      // Do not cache the plaintext singleton. In particular, io() may be called before metadata is
+      // loaded for a new table, and that must not prevent temporary operations from initializing
+      // encryption from uncommitted metadata later.
+      if (requestedConfig == null) {
+        return PlaintextEncryptionManager.instance();
+      }
+
+      encryptionManager = PolarisEncryptionUtil.encryptionManager(metadata, keyManagementClient);
+      encryptionManagerConfig = requestedConfig;
+      return encryptionManager;
+    }
+
+    private void resetEncryptionState() {
+      encryptionManager = null;
+      encryptionManagerConfig = null;
+      encryptingFileIO = null;
     }
 
     @Override
@@ -1771,7 +1867,10 @@ public void doRefresh() {
                       resolvedEntities,
                       new HashMap<>(tableDefaultProperties),
                       Set.of(PolarisStorageActions.READ, PolarisStorageActions.LIST));
-              return TableMetadataParser.read(fileIO, metadataLocation);
+              TableMetadata metadata = TableMetadataParser.read(fileIO, metadataLocation);
+              tableFileIO = fileIO;
+              resetEncryptionState();
+              return metadata;
             });
         if (polarisEventDispatcher.hasListeners(PolarisEventType.AFTER_REFRESH_TABLE)) {
           polarisEventDispatcher.dispatch(
@@ -1850,7 +1949,23 @@ public void doCommit(TableMetadata base, TableMetadata metadata) {
                   PolarisStorageActions.WRITE,
                   PolarisStorageActions.LIST));
 
-      String newLocation = writeNewMetadataIfRequired(base == null, metadata);
+      // SnapshotProducer may already have used this manager to generate and register encryption
+      // keys while writing encrypted manifests. Recreating it here would discard those keys
+      // before they are persisted in the new table metadata. A new-table commit has no current
+      // metadata and therefore still initializes the manager from the proposed metadata.
+      encryption(metadata);
+      encryptingFileIO = null;
+
+      TableMetadata tableMetadata = metadata;
+      if (encryptionManager instanceof StandardEncryptionManager) {
+        TableMetadata.Builder builder = TableMetadata.buildFrom(metadata);
+        for (var entry : EncryptionUtil.encryptionKeys(encryptionManager).entrySet()) {
+          builder.addEncryptionKey(entry.getValue());
+        }
+        tableMetadata = builder.build();
+      }
+
+      String newLocation = writeNewMetadataIfRequired(base == null, tableMetadata);
       String oldLocation = base == null ? null : base.metadataFileLocation();
 
       // TODO: Consider using the entity from doRefresh() directly to do the conflict detection
@@ -1923,21 +2038,22 @@ public void doCommit(TableMetadata base, TableMetadata metadata) {
             tableIdentifier, oldLocation, newLocation, existingLocation);
       }
 
-      // We diverge from `BaseMetastoreTableOperations` in the below code block
+      if (null == existingLocation) {
+        createTableLike(tableIdentifier, entity);
+      } else {
+        updateTableLike(tableIdentifier, entity);
+      }
+
+      // We diverge from `BaseMetastoreTableOperations` in the below code block. Only expose the
+      // new metadata through this operations instance after the catalog commit has succeeded.
       if (makeMetadataCurrentOnCommit) {
         currentMetadata =
-            TableMetadata.buildFrom(metadata)
+            TableMetadata.buildFrom(tableMetadata)
                 .withMetadataLocation(newLocation)
                 .discardChanges()
                 .build();
         currentMetadataLocation = newLocation;
       }
-
-      if (null == existingLocation) {
-        createTableLike(tableIdentifier, entity);
-      } else {
-        updateTableLike(tableIdentifier, entity);
-      }
     }
 
     private boolean requestedTableLocationsChanged(TableMetadata base, TableMetadata metadata) {
@@ -1995,12 +2111,12 @@ public LocationProvider locationProvider() {
 
         @Override
         public FileIO io() {
-          return BasePolarisTableOperations.this.io();
+          return BasePolarisTableOperations.this.io(uncommittedMetadata);
         }
 
         @Override
         public EncryptionManager encryption() {
-          return BasePolarisTableOperations.this.encryption();
+          return BasePolarisTableOperations.this.encryption(uncommittedMetadata);
         }
 
         @Override
@@ -2018,7 +2134,7 @@ protected String writeNewMetadataIfRequired(boolean newTable, TableMetadata meta
 
     protected String writeNewMetadata(TableMetadata metadata, int newVersion) {
       String newTableMetadataFilePath = newTableMetadataFilePath(metadata, newVersion);
-      OutputFile newMetadataLocation = io().newOutputFile(newTableMetadataFilePath);
+      OutputFile newMetadataLocation = io(metadata).newOutputFile(newTableMetadataFilePath);
 
       // write the new metadata
       // use overwrite to avoid negative caching in S3. this is safe because the metadata location