Core: cleanExpiredMetadata cleans unreferenced encryption keys

core/src/main/java/org/apache/iceberg/RemoveSnapshots.java

@@ -38,10 +38,12 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
+import org.apache.iceberg.encryption.EncryptedKey;
 import org.apache.iceberg.exceptions.CommitFailedException;
 import org.apache.iceberg.exceptions.ValidationException;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
@@ -241,6 +243,11 @@ private TableMetadata internalApply() {
       reachableSpecs.add(base.defaultSpecId());
       Set<Integer> reachableSchemas = Sets.newConcurrentHashSet();
       reachableSchemas.add(base.currentSchemaId());
+      Set<String> reachableKeyIds =
+          base.encryptionKeys().stream()
+              .map(EncryptedKey::encryptedById)
+              .filter(Objects::nonNull)
+              .collect(Collectors.toCollection(Sets::newConcurrentHashSet));
 
       boolean mayHaveExpiredSpecs = base.specs().size() > 1;
 
@@ -255,6 +262,9 @@ private TableMetadata internalApply() {
                       .forEach(reachableSpecs::add);
                 }
                 reachableSchemas.add(snapshot.schemaId());
+                if (snapshot.keyId() != null) {
+                  reachableKeyIds.add(snapshot.keyId());
+                }
               });
 
       Set<Integer> specsToRemove =
@@ -270,6 +280,13 @@ private TableMetadata internalApply() {
               .filter(schemaId -> !reachableSchemas.contains(schemaId))
               .collect(Collectors.toSet());
       updatedMetaBuilder.removeSchemas(schemasToRemove);
+
+      Set<String> encryptionKeysToRemove =
+          base.encryptionKeys().stream()
+              .map(EncryptedKey::keyId)
+              .filter(keyId -> !reachableKeyIds.contains(keyId))
+              .collect(Collectors.toSet());
+      encryptionKeysToRemove.forEach(updatedMetaBuilder::removeEncryptionKey);
     }
 
     return updatedMetaBuilder.build();

core/src/main/java/org/apache/iceberg/encryption/EncryptionUtil.java

@@ -23,9 +23,13 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
+import java.util.stream.Collectors;
 import org.apache.iceberg.CatalogProperties;
 import org.apache.iceberg.ManifestListFile;
+import org.apache.iceberg.Snapshot;
+import org.apache.iceberg.TableMetadata;
 import org.apache.iceberg.TableProperties;
 import org.apache.iceberg.common.DynConstructors;
 import org.apache.iceberg.io.OutputFile;
@@ -181,6 +185,32 @@ public static Map<String, EncryptedKey> encryptionKeys(EncryptionManager em) {
     return sem.encryptionKeys();
   }
 
+  /** Adds referenced encryption keys and the current key encryption key to table metadata. */
+  public static TableMetadata addEmKeysToMetadata(
+      TableMetadata metadata, EncryptionManager encryptionManager) {
+    if (!(encryptionManager instanceof StandardEncryptionManager standardEncryptionManager)) {
+      return metadata;
+    }
+
+    Set<String> referencedKeyIds =
+        Sets.union(
+            metadata.snapshots().stream()
+                .map(Snapshot::keyId)
+                .filter(Objects::nonNull)
+                .collect(Collectors.toUnmodifiableSet()),
+            metadata.encryptionKeys().stream()
+                .map(EncryptedKey::encryptedById)
+                .filter(Objects::nonNull)
+                .collect(Collectors.toUnmodifiableSet()));
+    String keyEncryptionKeyId = standardEncryptionManager.keyEncryptionKeyID();
+    TableMetadata.Builder builder = TableMetadata.buildFrom(metadata);
+    standardEncryptionManager.encryptionKeys().values().stream()
+        .filter(
+            key -> referencedKeyIds.contains(key.keyId()) || keyEncryptionKeyId.equals(key.keyId()))
+        .forEach(builder::addEncryptionKey);
+    return builder.build();
+  }
+
   /**
    * Encrypts the key metadata for a manifest list.
    *

core/src/test/java/org/apache/iceberg/TestRemoveSnapshots.java

@@ -41,6 +41,8 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 import org.apache.iceberg.ManifestEntry.Status;
+import org.apache.iceberg.encryption.BaseEncryptedKey;
+import org.apache.iceberg.encryption.EncryptedKey;
 import org.apache.iceberg.exceptions.ValidationException;
 import org.apache.iceberg.expressions.Expressions;
 import org.apache.iceberg.io.BulkDeletionFailureException;
@@ -1796,6 +1798,33 @@ public void testRemoveSchemas() {
     assertThat(table.schemas().values()).containsExactly(table.schema());
   }
 
+  @TestTemplate
+  public void testRemoveEncryptionKeys() {
+    table.newAppend().appendFile(FILE_A).commit();
+    waitUntilAfter(table.currentSnapshot().timestampMillis());
+
+    EncryptedKey kek =
+        new BaseEncryptedKey("kek", ByteBuffer.wrap(new byte[] {1}), "remote-kms-key", null);
+    EncryptedKey wrapped =
+        new BaseEncryptedKey("wrapped", ByteBuffer.wrap(new byte[] {2}), "kek", null);
+    TableMetadata base = table.ops().current();
+    table
+        .ops()
+        .commit(
+            base,
+            TableMetadata.buildFrom(base).addEncryptionKey(kek).addEncryptionKey(wrapped).build());
+
+    removeSnapshots(table)
+        .expireOlderThan(System.currentTimeMillis())
+        .cleanExpiredMetadata(true)
+        .commit();
+
+    assertThat(table.ops().current().encryptionKeys())
+        .as("Orphaned encryption keys are removed. The kek is kept as it encrypts the other keys.")
+        .extracting(EncryptedKey::keyId)
+        .containsExactly("kek");
+  }
+
   @TestTemplate
   public void testNoSchemasOrSpecsToRemove() {
     String tableName = "test_no_schemas_or_specs_to_remove";

hive-metastore/src/main/java/org/apache/iceberg/hive/HiveTableOperations.java

@@ -47,7 +47,6 @@
 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.CommitFailedException;
 import org.apache.iceberg.exceptions.CommitStateUnknownException;
@@ -240,25 +239,9 @@ the table key parameter (along with existing snapshots) in the file, making the
   @Override
   protected void doCommit(TableMetadata base, TableMetadata metadata) {
     boolean newTable = base == null;
-    final TableMetadata tableMetadata;
     encryptionPropsFromMetadata(metadata.properties());
-
-    String newMetadataLocation;
-    EncryptionManager encrManager = encryption();
-    if (encrManager instanceof StandardEncryptionManager) {
-      // Add new encryption keys to the metadata
-      TableMetadata.Builder builder = TableMetadata.buildFrom(metadata);
-      for (Map.Entry<String, EncryptedKey> entry :
-          EncryptionUtil.encryptionKeys(encrManager).entrySet()) {
-        builder.addEncryptionKey(entry.getValue());
-      }
-
-      tableMetadata = builder.build();
-    } else {
-      tableMetadata = metadata;
-    }
-
-    newMetadataLocation = writeNewMetadataIfRequired(newTable, tableMetadata);
+    TableMetadata tableMetadata = EncryptionUtil.addEmKeysToMetadata(metadata, encryption());
+    String newMetadataLocation = writeNewMetadataIfRequired(newTable, tableMetadata);
 
     boolean hiveEngineEnabled = hiveEngineEnabled(tableMetadata, conf);
     boolean keepHiveStats = conf.getBoolean(ConfigProperties.KEEP_HIVE_STATS, false);

spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/sql/TestTableEncryption.java

@@ -42,11 +42,13 @@
 import org.apache.iceberg.MetadataTableType;
 import org.apache.iceberg.Parameters;
 import org.apache.iceberg.Schema;
+import org.apache.iceberg.Snapshot;
 import org.apache.iceberg.Table;
 import org.apache.iceberg.TableMetadata;
 import org.apache.iceberg.Transaction;
 import org.apache.iceberg.actions.RewriteManifests;
 import org.apache.iceberg.encryption.Ciphers;
+import org.apache.iceberg.encryption.EncryptedKey;
 import org.apache.iceberg.encryption.UnitestKMS;
 import org.apache.iceberg.io.InputFile;
 import org.apache.iceberg.io.SeekableInputStream;
@@ -425,6 +427,45 @@ public void testDropTableWithPurge() {
         .allSatisfy(filePath -> assertThat(localInput(filePath).exists()).isFalse());
   }
 
+  @TestTemplate
+  void readAfterRemovingExpiredSnapshotKey() {
+    validationCatalog.initialize(catalogName, catalogConfig);
+    Table table = validationCatalog.loadTable(tableIdent);
+    Snapshot snapshotToExpire = table.currentSnapshot();
+
+    sql("INSERT INTO %s VALUES (4, 'd', 4.0)", tableName);
+    table.refresh();
+    Snapshot currentSnapshot = table.currentSnapshot();
+
+    assertThat(snapshotToExpire.keyId()).isNotNull().isNotEqualTo(currentSnapshot.keyId());
+    assertThat(currentSnapshot.keyId()).isNotNull();
+    assertThat(((HasTableOperations) table).operations().current().encryptionKeys())
+        .extracting(EncryptedKey::keyId)
+        .contains(snapshotToExpire.keyId(), currentSnapshot.keyId());
+
+    SparkActions.get()
+        .expireSnapshots(table)
+        .expireSnapshotId(snapshotToExpire.snapshotId())
+        .cleanExpiredMetadata(true)
+        .execute();
+
+    Table reloaded = validationCatalog.loadTable(tableIdent);
+    assertThat(reloaded.snapshot(snapshotToExpire.snapshotId())).isNull();
+    assertThat(reloaded.io().newInputFile(snapshotToExpire.manifestListLocation()).exists())
+        .isFalse();
+    assertThat(((HasTableOperations) reloaded).operations().current().encryptionKeys())
+        .extracting(EncryptedKey::keyId)
+        .doesNotContain(snapshotToExpire.keyId())
+        .contains(currentSnapshot.keyId());
+
+    sql("REFRESH TABLE %s", tableName);
+    assertEquals(
+        "Should read all rows after removing the expired snapshot's manifest list key",
+        ImmutableList.of(
+            row(1L, "a", 1.0F), row(2L, "b", 2.0F), row(3L, "c", Float.NaN), row(4L, "d", 4.0F)),
+        sql("SELECT * FROM %s ORDER BY id", tableName));
+  }
+
   private void checkMetadataFileEncryption(InputFile file) throws IOException {
     SeekableInputStream stream = file.newStream();
     byte[] magic = new byte[4];