From 27f7aea56ddaff0ddda62b13f76b04b736a0c3d8 Mon Sep 17 00:00:00 2001 From: Yuya Ebihara Date: Fri, 22 May 2026 16:38:37 +0900 Subject: [PATCH] Core: Add missing encryption keys when commiting snapshot --- .../org/apache/iceberg/SnapshotProducer.java | 11 ++ .../TestSnapshotProducerWithEncryption.java | 117 ++++++++++++++++++ 2 files changed, 128 insertions(+) create mode 100644 core/src/test/java/org/apache/iceberg/TestSnapshotProducerWithEncryption.java diff --git a/core/src/main/java/org/apache/iceberg/SnapshotProducer.java b/core/src/main/java/org/apache/iceberg/SnapshotProducer.java index 6ba10e8049f6..b2eb60286eb6 100644 --- a/core/src/main/java/org/apache/iceberg/SnapshotProducer.java +++ b/core/src/main/java/org/apache/iceberg/SnapshotProducer.java @@ -48,8 +48,12 @@ import java.util.concurrent.atomic.AtomicReferenceArray; import java.util.function.Consumer; import java.util.function.Function; +import org.apache.iceberg.encryption.EncryptedKey; import org.apache.iceberg.encryption.EncryptedOutputFile; import org.apache.iceberg.encryption.EncryptingFileIO; +import org.apache.iceberg.encryption.EncryptionManager; +import org.apache.iceberg.encryption.EncryptionUtil; +import org.apache.iceberg.encryption.StandardEncryptionManager; import org.apache.iceberg.events.CreateSnapshotEvent; import org.apache.iceberg.events.Listeners; import org.apache.iceberg.exceptions.CleanableFailure; @@ -484,6 +488,13 @@ public void commit() { update.setBranchSnapshot(newSnapshot, targetBranch); } + EncryptionManager encryptionManager = ops.encryption(); + if (encryptionManager instanceof StandardEncryptionManager) { + Map keys = + EncryptionUtil.encryptionKeys(encryptionManager); + keys.values().forEach(update::addEncryptionKey); + } + TableMetadata updated = update.build(); if (updated.changes().isEmpty()) { // do not commit if the metadata has not changed. for example, this may happen diff --git a/core/src/test/java/org/apache/iceberg/TestSnapshotProducerWithEncryption.java b/core/src/test/java/org/apache/iceberg/TestSnapshotProducerWithEncryption.java new file mode 100644 index 000000000000..b21d1c201fed --- /dev/null +++ b/core/src/test/java/org/apache/iceberg/TestSnapshotProducerWithEncryption.java @@ -0,0 +1,117 @@ +/* + * 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.apache.iceberg.types.Types.NestedField.required; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatCode; + +import java.io.File; +import java.util.Map; +import org.apache.iceberg.encryption.EncryptedKey; +import org.apache.iceberg.encryption.EncryptingFileIO; +import org.apache.iceberg.encryption.EncryptionManager; +import org.apache.iceberg.encryption.EncryptionUtil; +import org.apache.iceberg.encryption.PlaintextEncryptionManager; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; +import org.apache.iceberg.types.Types; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class TestSnapshotProducerWithEncryption { + + @TempDir private File tableDir; + + @Test + void propagateEncryptionKeysToMetadata() { + Schema schema = new Schema(required(1, "id", Types.LongType.get())); + PartitionSpec spec = PartitionSpec.unpartitioned(); + + EncryptedTableOps ops = new EncryptedTableOps("test", tableDir); + TestTables.TestTable table = + TestTables.create(tableDir, "test", schema, spec, SortOrder.unsorted(), 3, ops); + + table.updateProperties().set(TableProperties.ENCRYPTION_TABLE_KEY, "keyA").commit(); + + DataFile dataFile = + DataFiles.builder(spec) + .withPath(table.location() + "/data/file1.parquet") + .withFileSizeInBytes(100) + .withRecordCount(1) + .build(); + table.newAppend().appendFile(dataFile).commit(); + + Snapshot snapshot = table.currentSnapshot(); + String keyId = snapshot.keyId(); + assertThat(keyId).isNotNull(); + + TableMetadata metadata = ops.current(); + assertThat(metadata.encryptionKeys()).extracting(EncryptedKey::keyId).contains(keyId); + + assertThatCode(() -> table.newScan().planFiles()).doesNotThrowAnyException(); + } + + static class EncryptedTableOps extends TestTables.TestTableOperations { + private EncryptionManager encryptionManager; + + EncryptedTableOps(String tableName, File location) { + super(tableName, location); + updateEncryptionManager(); + } + + @Override + public EncryptionManager encryption() { + return encryptionManager; + } + + @Override + public FileIO io() { + return EncryptingFileIO.combine(super.io(), encryptionManager); + } + + @Override + public TableMetadata refresh() { + TableMetadata metadata = super.refresh(); + updateEncryptionManager(); + return metadata; + } + + private void updateEncryptionManager() { + TableMetadata metadata = current(); + if (metadata == null) { + encryptionManager = PlaintextEncryptionManager.instance(); + return; + } + + Map properties = metadata.properties(); + String tableKeyId = properties.get(TableProperties.ENCRYPTION_TABLE_KEY); + if (tableKeyId == null) { + encryptionManager = PlaintextEncryptionManager.instance(); + return; + } + + Map catalogProps = + ImmutableMap.of("encryption.kms-impl", "org.apache.iceberg.encryption.UnitestKMS"); + encryptionManager = + EncryptionUtil.createEncryptionManager( + metadata.encryptionKeys(), properties, EncryptionUtil.createKmsClient(catalogProps)); + } + } +}