singhpk234 commented on code in PR #17984: URL: https://github.com/apache/iceberg/pull/17984#discussion_r3990989278
########## core/src/test/java/org/apache/iceberg/TestManifestListKeyPersistence.java: ########## @@ -0,0 +1,229 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.io.File; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.apache.iceberg.encryption.EncryptedKey; +import org.apache.iceberg.encryption.EncryptingFileIO; +import org.apache.iceberg.encryption.EncryptionManager; +import org.apache.iceberg.encryption.EncryptionTestHelpers; +import org.apache.iceberg.encryption.EncryptionUtil; +import org.apache.iceberg.encryption.KeyManagementClient; +import org.apache.iceberg.encryption.UnitestKMS; +import org.apache.iceberg.exceptions.CommitFailedException; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +/** + * Tests that a snapshot and the keys required to read its manifest list are committed in the same + * metadata update. This ensures that a concurrent refresh cannot cause the snapshot to be committed + * without those keys. + */ +class TestManifestListKeyPersistence { Review Comment: do we need a dedicated class for this test ? why not add this test with other test ? we just need more assertion that keys are present with the reproducer test, right ? ########## core/src/main/java/org/apache/iceberg/SnapshotProducer.java: ########## @@ -354,6 +360,11 @@ public Snapshot apply() { replacedRecords); } + // Each call creates a new key; use this result for both the snapshot and key lookup. Review Comment: please elaborate what key is this ? this would be MLK right ########## core/src/main/java/org/apache/iceberg/SnapshotProducer.java: ########## @@ -499,10 +540,14 @@ public void commit() { if (base.snapshot(newSnapshot.snapshotId()) != null) { // this is a rollback operation update.setBranchSnapshot(newSnapshot.snapshotId(), targetBranch); - } else if (stageOnly) { - update.addSnapshot(newSnapshot); } else { - update.setBranchSnapshot(newSnapshot, targetBranch); + // Ensure the snapshot's manifest list key and its wrapping key are present. + encryptionKeysForManifestList.forEach(update::addEncryptionKey); + if (stageOnly) { + update.addSnapshot(newSnapshot); + } else { + update.setBranchSnapshot(newSnapshot, targetBranch); + } } Review Comment: This is a very hot code path ... must be careful, i wonder if we should always apply this update ? because in rollback it will be a non-op ? ########## core/src/main/java/org/apache/iceberg/SnapshotProducer.java: ########## @@ -365,7 +376,37 @@ public Snapshot apply() { manifestList.location(), nextRowId, assignedRows, - writer.toManifestListFile().encryptionKeyID()); + manifestListFile.encryptionKeyID()); + } + + /** + * Returns the encryption key records required to read a manifest list. + * + * <p>The manifest list key metadata is encrypted by a key encryption key, so both records must be + * present in table metadata. Unencrypted manifest lists require no keys. + */ + private static List<EncryptedKey> findEncryptionKeysForManifestList( + EncryptionManager encryption, String manifestListKeyID) { + if (manifestListKeyID == null) { + return List.of(); + } + + Map<String, EncryptedKey> encryptionKeys = EncryptionUtil.encryptionKeys(encryption); + EncryptedKey manifestListKey = encryptionKeys.get(manifestListKeyID); + if (manifestListKey == null) { + throw new CommitFailedException( + "Cannot find manifest list key with ID %s", manifestListKeyID); + } + + String keyEncryptionKeyID = manifestListKey.encryptedById(); + EncryptedKey keyEncryptionKey = encryptionKeys.get(keyEncryptionKeyID); + if (keyEncryptionKey == null) { + throw new CommitFailedException( + "Cannot find key encryption key with ID %s, which wraps manifest list key %s", + keyEncryptionKeyID, manifestListKeyID); Review Comment: i would also add snapshot id here ? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
