From 954786cc5e070408a40326eb18397713389a7582 Mon Sep 17 00:00:00 2001 From: Ahmed Abualsaud Date: Tue, 18 Aug 2026 12:18:43 -0700 Subject: [PATCH 1/3] encryption --- .../apache/beam/sdk/io/iceberg/AddFiles.java | 17 +- .../sdk/io/iceberg/AppendFilesToTables.java | 17 +- .../apache/beam/sdk/io/iceberg/ReadUtils.java | 33 +- .../beam/sdk/io/iceberg/RecordWriter.java | 9 +- .../beam/sdk/io/iceberg/cdc/CdcReadUtils.java | 56 +- .../sdk/io/iceberg/EncryptedTestCatalog.java | 189 ++++++ .../sdk/io/iceberg/IcebergEncryptionTest.java | 585 ++++++++++++++++++ .../apache/beam/sdk/io/iceberg/TestKms.java | 89 +++ .../io/iceberg/cdc/TestChangelogTasks.java | 61 ++ 9 files changed, 1029 insertions(+), 27 deletions(-) create mode 100644 sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/EncryptedTestCatalog.java create mode 100644 sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergEncryptionTest.java create mode 100644 sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TestKms.java create mode 100644 sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/TestChangelogTasks.java diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AddFiles.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AddFiles.java index 18b95ca50b1f..6f5fe245819a 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AddFiles.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AddFiles.java @@ -98,10 +98,11 @@ import org.apache.iceberg.TableProperties; import org.apache.iceberg.avro.Avro; import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.encryption.EncryptedOutputFile; +import org.apache.iceberg.encryption.EncryptingFileIO; import org.apache.iceberg.exceptions.AlreadyExistsException; import org.apache.iceberg.exceptions.NoSuchTableException; import org.apache.iceberg.io.InputFile; -import org.apache.iceberg.io.OutputFile; import org.apache.iceberg.mapping.MappingUtil; import org.apache.iceberg.mapping.NameMapping; import org.apache.iceberg.mapping.NameMappingParser; @@ -692,11 +693,9 @@ public void process( String manifestPath = String.format( "%s/metadata/%s-%s-m0.avro", table.location(), MANIFEST_PREFIX, UUID.randomUUID()); - OutputFile outputFile = table.io().newOutputFile(manifestPath); - int numDataFiles = 0; ManifestFile manifestFile; - try (ManifestWriter writer = ManifestFiles.write(spec, outputFile)) { + try (ManifestWriter writer = createManifestWriter(table, spec, manifestPath)) { for (SerializableDataFile sdf : batch.getValue()) { DataFile df = sdf.createDataFile(table.specs()); writer.add(df); @@ -713,6 +712,16 @@ public void process( output.output(KV.of(identifier, ManifestFiles.encode(manifestFile))); numDataFilesAdded.inc(numDataFiles); } + + /** Encrypts the manifest when the table is encrypted. */ + @SuppressWarnings("argument") + private static ManifestWriter createManifestWriter( + Table table, PartitionSpec spec, String manifestPath) { + EncryptedOutputFile outputFile = + EncryptingFileIO.combine(table.io(), table.encryption()) + .newEncryptingOutputFile(manifestPath); + return ManifestFiles.write(1, spec, outputFile, null); + } } /** diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AppendFilesToTables.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AppendFilesToTables.java index 917b087e39e0..8100c687f421 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AppendFilesToTables.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/AppendFilesToTables.java @@ -47,7 +47,8 @@ import org.apache.iceberg.Snapshot; import org.apache.iceberg.Table; import org.apache.iceberg.catalog.Catalog; -import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.encryption.EncryptedOutputFile; +import org.apache.iceberg.encryption.EncryptingFileIO; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -189,8 +190,7 @@ private void appendManifestFiles(Table table, Iterable fileWrit int specId = entry.getKey(); List files = entry.getValue(); PartitionSpec spec = Preconditions.checkStateNotNull(specs.get(specId)); - FileIO io = table.io(); - ManifestWriter writer = createManifestWriter(table.location(), uuid, spec, io); + ManifestWriter writer = createManifestWriter(table, uuid, spec); for (DataFile file : files) { writer.add(file); committedDataFileByteSize.update(file.fileSizeInBytes()); @@ -202,14 +202,19 @@ private void appendManifestFiles(Table table, Iterable fileWrit update.commit(); } + @SuppressWarnings("argument") private ManifestWriter createManifestWriter( - String tableLocation, String uuid, PartitionSpec spec, FileIO io) { + Table table, String uuid, PartitionSpec spec) { String location = FileFormat.AVRO.addExtension( String.format( "%s/metadata/%s-%s-%s.manifest", - tableLocation, manifestFilePrefix, uuid, spec.specId())); - return ManifestFiles.write(spec, io.newOutputFile(location)); + table.location(), manifestFilePrefix, uuid, spec.specId())); + // Encrypts the manifest when the table is encrypted + EncryptedOutputFile outputFile = + EncryptingFileIO.combine(table.io(), table.encryption()) + .newEncryptingOutputFile(location); + return ManifestFiles.write(1, spec, outputFile, null); } // If the process call fails immediately after a successful commit, it gets diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadUtils.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadUtils.java index 3e70ae126877..8e4b75b5e363 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadUtils.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadUtils.java @@ -41,6 +41,7 @@ import org.apache.iceberg.data.parquet.GenericParquetReaders; import org.apache.iceberg.encryption.EncryptedFiles; import org.apache.iceberg.encryption.EncryptedInputFile; +import org.apache.iceberg.encryption.NativeEncryptionInputFile; import org.apache.iceberg.expressions.Evaluator; import org.apache.iceberg.expressions.Expression; import org.apache.iceberg.hadoop.HadoopInputFile; @@ -49,9 +50,11 @@ import org.apache.iceberg.mapping.NameMapping; import org.apache.iceberg.mapping.NameMappingParser; import org.apache.iceberg.parquet.ParquetReader; +import org.apache.iceberg.util.ByteBuffers; import org.apache.iceberg.util.SnapshotUtil; import org.apache.parquet.HadoopReadOptions; import org.apache.parquet.ParquetReadOptions; +import org.apache.parquet.crypto.FileDecryptionProperties; import org.checkerframework.checker.nullness.qual.Nullable; /** Helper class for source operations. */ @@ -79,6 +82,28 @@ public static CloseableIterable createReader( task.residual()); } + /** + * Returns the file a Parquet reader should actually open for {@code decrypted}, the result of + * {@link org.apache.iceberg.encryption.EncryptionManager#decrypt}. + */ + public static InputFile parquetInputFile(InputFile decrypted) { + return decrypted instanceof NativeEncryptionInputFile + ? ((NativeEncryptionInputFile) decrypted).encryptedInputFile() + : decrypted; + } + + /** Returns the Parquet decryption properties, or null if the file is not encrypted. */ + public static @Nullable FileDecryptionProperties parquetDecryption(InputFile decrypted) { + if (!(decrypted instanceof NativeEncryptionInputFile)) { + return null; + } + NativeEncryptionInputFile nativeFile = (NativeEncryptionInputFile) decrypted; + return FileDecryptionProperties.builder() + .withFooterKey(ByteBuffers.toByteArray(nativeFile.keyMetadata().encryptionKey())) + .withAADPrefix(ByteBuffers.toByteArray(nativeFile.keyMetadata().aadPrefix())) + .build(); + } + public static CloseableIterable createReader( Table table, IcebergScanConfig scanConfig, @@ -91,7 +116,10 @@ public static CloseableIterable createReader( Expression residual) { EncryptedInputFile encryptedInput = EncryptedFiles.encryptedInput(table.io().newInputFile(file.location()), file.keyMetadata()); - InputFile inputFile = table.encryption().decrypt(encryptedInput); + InputFile decrypted = table.encryption().decrypt(encryptedInput); + InputFile inputFile = parquetInputFile(decrypted); + @Nullable FileDecryptionProperties decryptionProperties = parquetDecryption(decrypted); + Map idToConstants = PartitionUtils.constantsMap(spec, file, fileSequenceNumber); ParquetReadOptions.Builder optionsBuilder; @@ -109,6 +137,9 @@ public static CloseableIterable createReader( optionsBuilder .withRange(start, start + length) .withMaxAllocationInBytes(MAX_FILE_BUFFER_SIZE); + if (decryptionProperties != null) { + optionsBuilder = optionsBuilder.withDecryption(decryptionProperties); + } @Nullable String nameMapping = table.properties().get(TableProperties.DEFAULT_NAME_MAPPING); NameMapping mapping = diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriter.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriter.java index c3b63b2a336f..8196c9306470 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriter.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriter.java @@ -83,19 +83,16 @@ class RecordWriter { fileFormat.addExtension( table.locationProvider().newDataLocation(table.spec(), partitionKey, filename)); } - OutputFile outputFile; - EncryptionKeyMetadata keyMetadata; // table.io() may return a shared FileIO instance. // FileIO lifecycle is managed by RecordWriterManager.close(). OutputFile tmpFile = table.io().newOutputFile(absoluteFilename); EncryptedOutputFile encryptedOutputFile = table.encryption().encrypt(tmpFile); - outputFile = encryptedOutputFile.encryptingOutputFile(); - keyMetadata = encryptedOutputFile.keyMetadata(); + EncryptionKeyMetadata keyMetadata = encryptedOutputFile.keyMetadata(); switch (fileFormat) { case AVRO: icebergDataWriter = - Avro.writeData(outputFile) + Avro.writeData(encryptedOutputFile) .forTable(table) .createWriterFunc(org.apache.iceberg.data.avro.DataWriter::create) .withPartition(partitionKey) @@ -105,7 +102,7 @@ class RecordWriter { break; case PARQUET: Parquet.DataWriteBuilder parquetBuilder = - Parquet.writeData(outputFile) + Parquet.writeData(encryptedOutputFile) .forTable(table) .createWriterFunc(GenericParquetWriter::create) .withPartition(partitionKey) diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtils.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtils.java index b8c18d1a4e53..640fb24079b3 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtils.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/cdc/CdcReadUtils.java @@ -44,6 +44,7 @@ import org.apache.iceberg.data.InternalRecordWrapper; import org.apache.iceberg.data.Record; import org.apache.iceberg.deletes.PositionDeleteIndex; +import org.apache.iceberg.encryption.EncryptingFileIO; import org.apache.iceberg.expressions.Expression; import org.apache.iceberg.expressions.Expressions; import org.apache.iceberg.io.CloseableIterable; @@ -54,6 +55,8 @@ import org.apache.iceberg.types.TypeUtil; import org.apache.iceberg.types.Types; import org.apache.iceberg.util.StructLikeSet; +import org.apache.parquet.ParquetReadOptions; +import org.apache.parquet.crypto.FileDecryptionProperties; import org.apache.parquet.hadoop.ParquetFileReader; import org.apache.parquet.hadoop.metadata.BlockMetaData; import org.apache.parquet.hadoop.metadata.ParquetMetadata; @@ -141,11 +144,26 @@ public static CloseableIterable createReader( combined); } + /** + * Wraps the table's {@link FileIO} so files in an encrypted table are decrypted on open. Returns + * the plain IO unchanged for unencrypted tables. + */ + private static EncryptingFileIO encryptingIO(Table table) { + return EncryptingFileIO.combine(table.io(), table.encryption()); + } + + /** Applies Parquet decryption properties to {@code builder} if {@code decrypted} needs them. */ + private static ParquetReadOptions.Builder withDecryption( + ParquetReadOptions.Builder builder, InputFile decrypted) { + @Nullable FileDecryptionProperties decryption = ReadUtils.parquetDecryption(decrypted); + return decryption == null ? builder : builder.withDecryption(decryption); + } + /** Returns a filter that skips records marked for deletion. */ public static DeleteFilter genericDeleteFilter( Table table, Schema outputSchema, String dataFilePath, List deletes) { return new GenericDeleteFilter( - table.io(), + encryptingIO(table), dataFilePath, table.schema(), outputSchema, @@ -162,7 +180,7 @@ public static DeleteReader genericDeleteReader( List deletes, DeleteReader.PreloadedDeletes preloadedDeletes) { return new GenericDeleteReader( - table.io(), + encryptingIO(table), dataFilePath, table.schema(), outputSchema, @@ -340,7 +358,7 @@ private static CloseableIterable deletedRowsForTask( // 1. pre-load the position index for this data file. PositionDeleteIndex posIndex; try { - DeleteLoader loader = new BaseDeleteLoader(df -> table.io().newInputFile(df.location())); + DeleteLoader loader = new BaseDeleteLoader(encryptingIO(table)::newInputFile); posIndex = loader.loadPositionDeletes(posFiles, dataFilePath); } catch (RuntimeException e) { LOG.info( @@ -381,8 +399,14 @@ private static CloseableIterable deletedRowsForTask( try { long[] sortedDeletePositions = sortedDeletePositions(posIndex); - InputFile inputFile = table.io().newInputFile(dataFilePath); - try (ParquetFileReader reader = ParquetFileReader.open(asParquetInputFile(inputFile))) { + // decrypt via the table's encryption manager + InputFile decrypted = + encryptingIO(table).newInputFile(task.getDataFile().createDataFile(table.specs())); + ParquetReadOptions readOptions = + withDecryption(ParquetReadOptions.builder(), decrypted).build(); + try (ParquetFileReader reader = + ParquetFileReader.open( + asParquetInputFile(ReadUtils.parquetInputFile(decrypted)), readOptions)) { ParquetMetadata footer = reader.getFooter(); MessageType parquetSchema = footer.getFileMetaData().getSchema(); @@ -505,7 +529,7 @@ private static EqualityPushdownResult buildEqualityDeletePushdown( } Schema deleteSchema = TypeUtil.select(table.schema(), sharedIds); - DeleteLoader loader = new BaseDeleteLoader(df -> table.io().newInputFile(df.location())); + DeleteLoader loader = new BaseDeleteLoader(encryptingIO(table)::newInputFile); StructLikeSet set; try { set = loader.loadEqualityDeletes(eqFiles, deleteSchema); @@ -597,12 +621,12 @@ private DeleteReader.PreloadedDeletes preloadedDeletes( } public static class GenericDeleteFilter extends DeleteFilter { - private final FileIO io; + private final EncryptingFileIO io; private final InternalRecordWrapper asStructLike; @SuppressWarnings("method.invocation") public GenericDeleteFilter( - FileIO io, + EncryptingFileIO io, String dataFilePath, Schema tableSchema, Schema requiredSchema, @@ -621,15 +645,21 @@ protected StructLike asStructLike(Record record) { protected InputFile getInputFile(String location) { return io.newInputFile(location); } + + /** Overridden so delete files in an encrypted table are decrypted. */ + @Override + protected InputFile loadInputFile(DeleteFile deleteFile) { + return io.newInputFile(deleteFile); + } } public static class GenericDeleteReader extends DeleteReader { - private final FileIO io; + private final EncryptingFileIO io; private final InternalRecordWrapper asStructLike; @SuppressWarnings("method.invocation") public GenericDeleteReader( - FileIO io, + EncryptingFileIO io, String dataFilePath, Schema tableSchema, Schema requiredSchema, @@ -649,6 +679,12 @@ protected StructLike asStructLike(Record record) { protected InputFile getInputFile(String location) { return io.newInputFile(location); } + + /** Overridden so delete files in an encrypted table are decrypted. */ + @Override + protected InputFile loadInputFile(DeleteFile deleteFile) { + return io.newInputFile(deleteFile); + } } /** diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/EncryptedTestCatalog.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/EncryptedTestCatalog.java new file mode 100644 index 000000000000..662c37f8ecb8 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/EncryptedTestCatalog.java @@ -0,0 +1,189 @@ +/* + * 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.beam.sdk.io.iceberg; + +import java.io.IOException; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.apache.iceberg.TableMetadata; +import org.apache.iceberg.TableOperations; +import org.apache.iceberg.catalog.TableIdentifier; +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.KeyManagementClient; +import org.apache.iceberg.encryption.PlaintextEncryptionManager; +import org.apache.iceberg.encryption.StandardEncryptionManager; +import org.apache.iceberg.hadoop.HadoopCatalog; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.io.LocationProvider; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * A {@link HadoopCatalog} that supports Iceberg v3 table encryption, for tests. + * + *

This catalog reproduces that wiring on top of {@link HadoopCatalog} so encryption can be + * tested against a local warehouse. It builds an encryption manager from the table's {@code + * encryption.key-id} property, and persists the keys it generates into the table metadata's {@code + * encryption-keys} list on commit. + * + *

Enable it by pointing {@code catalog-impl} at this class and {@code encryption.kms-impl} at a + * {@link KeyManagementClient} such as {@link TestKms}. + */ +public class EncryptedTestCatalog extends HadoopCatalog { + private @Nullable KeyManagementClient kmsClient; + + @Override + public void initialize(String name, Map properties) { + super.initialize(name, properties); + this.kmsClient = EncryptionUtil.createKmsClient(properties); + } + + @Override + protected TableOperations newTableOps(TableIdentifier identifier) { + if (kmsClient == null) { + return super.newTableOps(identifier); + } + return new EncryptingTableOperations(super.newTableOps(identifier), kmsClient); + } + + @Override + public void close() throws IOException { + super.close(); + if (kmsClient != null) { + kmsClient.close(); + kmsClient = null; + } + } + + /** + * Delegating {@link TableOperations} that adds an encryption manager and carries the keys it + * mints into table metadata. + */ + private static class EncryptingTableOperations implements TableOperations { + private final TableOperations delegate; + private final KeyManagementClient kmsClient; + + /** + * Keys minted by encryption managers built by these operations. Kept across rebuilds so that a + * key generated before a refresh is not lost before it is committed. + */ + private final Map knownKeys = new LinkedHashMap<>(); + + private @Nullable EncryptionManager encryptionManager; + private @Nullable FileIO encryptingFileIO; + + EncryptingTableOperations(TableOperations delegate, KeyManagementClient kmsClient) { + this.delegate = delegate; + this.kmsClient = kmsClient; + } + + @Override + public EncryptionManager encryption() { + EncryptionManager existing = encryptionManager; + if (existing != null) { + return existing; + } + + TableMetadata metadata = current(); + if (metadata == null) { + // table does not exist yet; no properties to build an encryption manager from + return PlaintextEncryptionManager.instance(); + } + + for (EncryptedKey key : metadata.encryptionKeys()) { + knownKeys.putIfAbsent(key.keyId(), key); + } + EncryptionManager created = + EncryptionUtil.createEncryptionManager( + List.copyOf(knownKeys.values()), metadata.properties(), kmsClient); + encryptionManager = created; + return created; + } + + @Override + public void commit(TableMetadata base, TableMetadata metadata) { + EncryptionManager encryption = encryption(); + TableMetadata toCommit = metadata; + if (encryption instanceof StandardEncryptionManager) { + knownKeys.putAll(EncryptionUtil.encryptionKeys(encryption)); + TableMetadata.Builder builder = TableMetadata.buildFrom(metadata); + knownKeys.values().forEach(builder::addEncryptionKey); + toCommit = builder.build(); + } + + delegate.commit(base, toCommit); + this.encryptionManager = null; + this.encryptingFileIO = null; + } + + @Override + public TableMetadata refresh() { + TableMetadata refreshed = delegate.refresh(); + this.encryptionManager = null; + this.encryptingFileIO = null; + return refreshed; + } + + @Override + public TableMetadata current() { + return delegate.current(); + } + + /** + * Returns an {@link EncryptingFileIO} for encrypted tables, so manifests and manifest lists are + * transparently decrypted. This mirrors {@code HiveTableOperations#io()}. + */ + @Override + public FileIO io() { + EncryptionManager encryption = encryption(); + if (!(encryption instanceof StandardEncryptionManager)) { + return delegate.io(); + } + + FileIO existing = encryptingFileIO; + if (existing == null) { + existing = EncryptingFileIO.combine(delegate.io(), encryption); + encryptingFileIO = existing; + } + return existing; + } + + @Override + public String metadataFileLocation(String fileName) { + return delegate.metadataFileLocation(fileName); + } + + @Override + public LocationProvider locationProvider() { + return delegate.locationProvider(); + } + + @Override + public long newSnapshotId() { + return delegate.newSnapshotId(); + } + + @Override + public boolean requireStrictCleanup() { + return delegate.requireStrictCleanup(); + } + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergEncryptionTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergEncryptionTest.java new file mode 100644 index 000000000000..b1b33407e09b --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/IcebergEncryptionTest.java @@ -0,0 +1,585 @@ +/* + * 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.beam.sdk.io.iceberg; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.io.IOException; +import java.io.Serializable; +import java.net.URI; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.UUID; +import java.util.stream.Collectors; +import org.apache.beam.sdk.io.iceberg.cdc.CdcReadUtils; +import org.apache.beam.sdk.io.iceberg.cdc.TestChangelogTasks; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.Path; +import org.apache.iceberg.CatalogProperties; +import org.apache.iceberg.CatalogUtil; +import org.apache.iceberg.ContentFile; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.HasTableOperations; +import org.apache.iceberg.ManifestFile; +import org.apache.iceberg.ManifestFiles; +import org.apache.iceberg.PartitionKey; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Snapshot; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableMetadata; +import org.apache.iceberg.TableProperties; +import org.apache.iceberg.catalog.Catalog; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.GenericAppenderFactory; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.IcebergGenerics; +import org.apache.iceberg.data.InternalRecordWrapper; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.deletes.EqualityDeleteWriter; +import org.apache.iceberg.deletes.PositionDelete; +import org.apache.iceberg.deletes.PositionDeleteWriter; +import org.apache.iceberg.encryption.EncryptedFiles; +import org.apache.iceberg.encryption.EncryptedInputFile; +import org.apache.iceberg.encryption.EncryptedKey; +import org.apache.iceberg.encryption.EncryptedOutputFile; +import org.apache.iceberg.encryption.EncryptingFileIO; +import org.apache.iceberg.encryption.NativeEncryptionInputFile; +import org.apache.iceberg.encryption.StandardEncryptionManager; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.DataWriter; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.io.InputFile; +import org.apache.iceberg.io.OutputFileFactory; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.ByteBuffers; +import org.apache.parquet.ParquetReadOptions; +import org.apache.parquet.crypto.FileDecryptionProperties; +import org.apache.parquet.crypto.ParquetCryptoRuntimeException; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.metadata.FileMetaData; +import org.apache.parquet.hadoop.util.HadoopInputFile; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.hamcrest.Matchers; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.rules.TestName; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Verifies that IcebergIO honors Iceberg v3 table encryption. + * + *

These are self-contained: {@link TestKms} stands in for a KMS service, so no cloud credentials + * are needed. + * + *

What the spec requires of an engine working with a table that has {@code encryption.key-id} + * set: + * + *

    + *
  • data files are encrypted, and each manifest entry carries the file's {@code key_metadata} + *
  • manifests and manifest lists are encrypted; the snapshot records the manifest list's {@code + * key-id} + *
  • the keys minted along the way are tracked in the table metadata's {@code encryption-keys} + *
  • the encrypted files are readable again + *
+ */ +@RunWith(JUnit4.class) +public class IcebergEncryptionTest implements Serializable { + /** Same shape as {@link TestFixtures#SCHEMA}, plus the identifier field CDC needs. */ + private static final Schema CDC_SCHEMA = + new Schema( + ImmutableList.of( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.optional(2, "data", Types.StringType.get())), + ImmutableSet.of(1)); + + @Rule public transient TemporaryFolder temporaryFolder = new TemporaryFolder(); + @Rule public transient TestPipeline writePipeline = TestPipeline.create(); + @Rule public transient TestPipeline readPipeline = TestPipeline.create(); + @Rule public transient TestName testName = new TestName(); + + private String warehouse; + private IcebergCatalogConfig catalogConfig; + private TableIdentifier tableId; + + @Before + public void setUp() throws IOException { + warehouse = "file:" + temporaryFolder.newFolder("warehouse"); + catalogConfig = + IcebergCatalogConfig.builder() + .setCatalogName("encryption-test") + .setCatalogProperties( + ImmutableMap.builder() + .put(CatalogProperties.CATALOG_IMPL, EncryptedTestCatalog.class.getName()) + .put(CatalogProperties.WAREHOUSE_LOCATION, warehouse) + .put(CatalogProperties.ENCRYPTION_KMS_IMPL, TestKms.class.getName()) + .build()) + .build(); + tableId = + TableIdentifier.of( + "default", + testName.getMethodName() + "_" + UUID.randomUUID().toString().substring(0, 8)); + } + + /** + * Checks the fixture itself: the test catalog really does hand out a {@link + * StandardEncryptionManager}. Without this, every other test here could pass vacuously against a + * plaintext table. + */ + @Test + public void testFixtureProducesAnEncryptedTable() { + Table table = createEncryptedTable(TestFixtures.SCHEMA); + + assertThat(table.encryption(), Matchers.instanceOf(StandardEncryptionManager.class)); + assertEquals( + TestKms.MASTER_KEY_ID, table.properties().get(TableProperties.ENCRYPTION_TABLE_KEY)); + assertEquals(3, ((HasTableOperations) table).operations().current().formatVersion()); + } + + /** + * Spec: data files in an encrypted table are encrypted, and the manifest entry carries the + * per-file {@code key_metadata} needed to decrypt them. + */ + @Test + public void testWrittenDataFilesAreEncryptedAndCarryKeyMetadata() throws IOException { + createEncryptedTable(TestFixtures.SCHEMA); + runWritePipeline(); + + Table table = loadTable(); + List dataFiles = + ImmutableList.copyOf(table.currentSnapshot().addedDataFiles(table.io())); + assertFalse("expected the write to produce data files", dataFiles.isEmpty()); + + for (DataFile dataFile : dataFiles) { + assertNotNull( + "data file " + dataFile.location() + " is missing key_metadata", dataFile.keyMetadata()); + assertDataFileIsEncrypted(table, dataFile); + } + } + + /** + * Spec: manifests and manifest lists in an encrypted table are encrypted, and the snapshot + * records the key ID used for its manifest list. + */ + @Test + public void testManifestsAndManifestListAreEncrypted() throws IOException { + createEncryptedTable(TestFixtures.SCHEMA); + runWritePipeline(); + + Table table = loadTable(); + Snapshot snapshot = table.currentSnapshot(); + + assertNotNull("snapshot is missing the manifest list key-id", snapshot.keyId()); + // Iceberg refuses to open an encrypted manifest list through a plain FileIO + assertThrows( + "manifest list is readable without the encryption manager", + IllegalArgumentException.class, + () -> snapshot.allManifests(plaintextIo())); + + List manifests = snapshot.allManifests(table.io()); + assertFalse("expected the write to produce manifests", manifests.isEmpty()); + assertManifestsAreEncrypted(table, manifests); + } + + /** + * Covers the manifest-writing path in {@link AppendFilesToTables}, which only runs when a batch + * of files spans more than one partition spec. + */ + @Test + public void testManifestsWrittenForMultiSpecBatchAreEncrypted() throws IOException { + Table table = createEncryptedTable(TestFixtures.SCHEMA); + + // one file under the initial (unpartitioned) spec... + FileWriteResult first = writeDataFileWithBeam(table, TestFixtures.FILE1SNAPSHOT1); + + // ...and one under a second spec, so the batch spans two specs + table.updateSpec().addField("data").commit(); + table.refresh(); + FileWriteResult second = writeDataFileWithBeam(table, TestFixtures.FILE1SNAPSHOT2); + + writePipeline + .apply(Create.of(first, second)) + .apply(new AppendFilesToTables(catalogConfig, "test-manifest")); + writePipeline.run().waitUntilFinish(); + + Table committed = loadTable(); + List manifests = committed.currentSnapshot().allManifests(committed.io()); + + // sanity check: we really did exercise the manifest path, one manifest per spec + assertEquals("expected one manifest per partition spec", 2, manifests.size()); + assertManifestsAreEncrypted(committed, manifests); + } + + @Test + public void testEncryptionKeysAreTrackedInTableMetadata() { + createEncryptedTable(TestFixtures.SCHEMA); + runWritePipeline(); + + Table table = loadTable(); + TableMetadata metadata = ((HasTableOperations) table).operations().current(); + + assertFalse( + "table metadata has no encryption-keys after a write", metadata.encryptionKeys().isEmpty()); + for (EncryptedKey key : metadata.encryptionKeys()) { + assertNotNull(key.keyId()); + ByteBuffer encryptedKeyMetadata = key.encryptedKeyMetadata(); + assertNotNull(encryptedKeyMetadata); + assertTrue(encryptedKeyMetadata.remaining() > 0); + } + + // the snapshot's key-id must resolve to one of the tracked keys + String snapshotKeyId = table.currentSnapshot().keyId(); + assertThat( + "snapshot key-id is not tracked in encryption-keys", + metadata.encryptionKeys().stream().map(EncryptedKey::keyId).collect(Collectors.toList()), + Matchers.hasItem(snapshotKeyId)); + } + + @Test + public void testEncryptedFilesAreReadableByIceberg() { + createEncryptedTable(TestFixtures.SCHEMA); + runWritePipeline(); + + List records = ImmutableList.copyOf(IcebergGenerics.read(loadTable()).build()); + + assertThat(records, Matchers.containsInAnyOrder(TestFixtures.FILE1SNAPSHOT1.toArray())); + } + + @Test + public void testEncryptedTableRoundTripsThroughBeam() { + createEncryptedTable(TestFixtures.SCHEMA); + runWritePipeline(); + + PCollection rows = + readPipeline.apply("Read encrypted table", IcebergIO.readRows(catalogConfig).from(tableId)); + + PAssert.that(rows).containsInAnyOrder(TestFixtures.asRows(TestFixtures.FILE1SNAPSHOT1)); + readPipeline.run().waitUntilFinish(); + } + + @Test + public void testReadsFilesWrittenByIcebergLibraries() throws IOException { + Table table = createEncryptedTable(TestFixtures.SCHEMA); + appendWithIceberg(table, TestFixtures.FILE1SNAPSHOT1); + + PCollection rows = + readPipeline.apply("Read encrypted table", IcebergIO.readRows(catalogConfig).from(tableId)); + + PAssert.that(rows).containsInAnyOrder(TestFixtures.asRows(TestFixtures.FILE1SNAPSHOT1)); + readPipeline.run().waitUntilFinish(); + } + + @Test + public void testTableIsUnreadableWithoutTheKms() { + createEncryptedTable(TestFixtures.SCHEMA); + runWritePipeline(); + + Table table = loadTableWithoutKms(); + + assertThrows( + "encrypted table was readable without the KMS", + IllegalArgumentException.class, + () -> ImmutableList.copyOf(IcebergGenerics.read(table).build())); + } + + @Test + public void testCdcReadsEncryptedDataFile() throws IOException { + Table table = createEncryptedTable(CDC_SCHEMA); + DataFile dataFile = + writeDataFileWithIceberg(table, cdcRecord(0L, "alpha"), cdcRecord(1L, "beta")); + assertNotNull("data file should carry key metadata", dataFile.keyMetadata()); + + CloseableIterable records = + CdcReadUtils.changelogRecordsForTask( + TestChangelogTasks.addedRows(table, dataFile, ImmutableList.of()), + table, + cdcScanConfig(table), + true); + + assertEquals(ImmutableList.of(0L, 1L), idsOf(records)); + } + + @Test + public void testCdcAppliesEncryptedDeleteFiles() throws IOException { + Table table = createEncryptedTable(CDC_SCHEMA); + DataFile dataFile = + writeDataFileWithIceberg( + table, + cdcRecord(0L, "keep-0"), + cdcRecord(1L, "drop-by-pos"), + cdcRecord(2L, "drop-by-data"), + cdcRecord(3L, "keep-3")); + + DeleteFile positionDelete = writePositionDelete(table, dataFile, 1L); + DeleteFile equalityDelete = writeEqualityDelete(table, dataFile, "drop-by-data"); + assertNotNull("position delete should carry key metadata", positionDelete.keyMetadata()); + assertNotNull("equality delete should carry key metadata", equalityDelete.keyMetadata()); + + CloseableIterable records = + CdcReadUtils.changelogRecordsForTask( + TestChangelogTasks.addedRows( + table, dataFile, ImmutableList.of(positionDelete, equalityDelete)), + table, + cdcScanConfig(table), + true); + + assertEquals(ImmutableList.of(0L, 3L), idsOf(records)); + } + + /** Creates a v3 table with a table encryption key, i.e. an encrypted table. */ + private Table createEncryptedTable(Schema schema) { + return catalog() + .buildTable(tableId, schema) + .withProperty(TableProperties.FORMAT_VERSION, "3") + .withProperty(TableProperties.ENCRYPTION_TABLE_KEY, TestKms.MASTER_KEY_ID) + .create(); + } + + private Catalog catalog() { + return catalogConfig.catalog(); + } + + private Table loadTable() { + return catalog().loadTable(tableId); + } + + private void runWritePipeline() { + writePipeline + .apply("Records to add", Create.of(TestFixtures.asRows(TestFixtures.FILE1SNAPSHOT1))) + .setRowSchema(IcebergUtils.icebergSchemaToBeamSchema(TestFixtures.SCHEMA)) + .apply("Append to table", IcebergIO.writeRows(catalogConfig).to(tableId)); + writePipeline.run().waitUntilFinish(); + } + + /** + * Asserts that each manifest is encrypted: it carries {@code key_metadata}, Iceberg refuses to + * open it without an encryption manager. + */ + private void assertManifestsAreEncrypted(Table table, List manifests) + throws IOException { + for (ManifestFile manifest : manifests) { + assertNotNull( + "manifest " + manifest.path() + " is missing key_metadata", manifest.keyMetadata()); + assertThrows( + "manifest " + manifest.path() + " is readable without the encryption manager", + IllegalArgumentException.class, + () -> plaintextIo().newInputFile(manifest)); + + List dataFilePaths = + ImmutableList.copyOf(ManifestFiles.readPaths(manifest, table.io(), table.specs())); + assertFalse( + "manifest " + manifest.path() + " decrypted to no entries", dataFilePaths.isEmpty()); + } + } + + /** Writes a data file through Beam's {@link RecordWriter}. */ + private FileWriteResult writeDataFileWithBeam(Table table, List records) + throws IOException { + PartitionKey partitionKey = new PartitionKey(table.spec(), table.schema()); + if (!table.spec().isUnpartitioned()) { + partitionKey.partition( + new InternalRecordWrapper(table.schema().asStruct()).wrap(records.get(0))); + } + + RecordWriter writer = + new RecordWriter(table, FileFormat.PARQUET, UUID.randomUUID().toString(), partitionKey); + for (Record record : records) { + writer.write(record); + } + writer.close(); + + return FileWriteResult.builder() + .setTableIdentifier(tableId) + .setSerializableDataFile(SerializableDataFile.from(writer.getDataFile(), table.spec())) + .build(); + } + + /** Writes a data file through Iceberg's own writer stack. */ + private static DataFile writeDataFileWithIceberg(Table table, Record... records) + throws IOException { + return writeDataFileWithIceberg(table, ImmutableList.copyOf(records)); + } + + private static DataFile writeDataFileWithIceberg(Table table, List records) + throws IOException { + OutputFileFactory fileFactory = + OutputFileFactory.builderFor(table, 1, 1).format(FileFormat.PARQUET).build(); + GenericAppenderFactory appenderFactory = + new GenericAppenderFactory(table.schema(), table.spec()); + + DataWriter writer = + appenderFactory.newDataWriter(fileFactory.newOutputFile(), FileFormat.PARQUET, null); + try (DataWriter toClose = writer) { + for (Record record : records) { + toClose.write(record); + } + } + return writer.toDataFile(); + } + + /** Appends records to a table using Iceberg's own writer stack, and commits. */ + private static void appendWithIceberg(Table table, List records) throws IOException { + table.newAppend().appendFile(writeDataFileWithIceberg(table, records)).commit(); + } + + private static DeleteFile writePositionDelete(Table table, DataFile dataFile, long... positions) + throws IOException { + GenericAppenderFactory appenderFactory = + new GenericAppenderFactory(table.schema(), table.spec()); + PositionDeleteWriter writer = + appenderFactory.newPosDeleteWriter( + encryptedFile(table, dataFile.location() + ".pos-delete.parquet"), + FileFormat.PARQUET, + null); + try (PositionDeleteWriter toClose = writer) { + for (long position : positions) { + toClose.write(PositionDelete.create().set(dataFile.location(), position)); + } + } + return writer.toDeleteFile(); + } + + private static DeleteFile writeEqualityDelete( + Table table, DataFile dataFile, @Nullable String data) throws IOException { + Schema deleteSchema = table.schema().select("data"); + GenericAppenderFactory appenderFactory = + new GenericAppenderFactory(table.schema(), table.spec(), new int[] {2}, deleteSchema, null); + EqualityDeleteWriter writer = + appenderFactory.newEqDeleteWriter( + encryptedFile(table, dataFile.location() + ".eq-delete.parquet"), + FileFormat.PARQUET, + null); + try (EqualityDeleteWriter toClose = writer) { + GenericRecord deleteRecord = GenericRecord.create(deleteSchema); + deleteRecord.setField("data", data); + toClose.write(deleteRecord); + } + return writer.toDeleteFile(); + } + + private static EncryptedOutputFile encryptedFile(Table table, String location) { + return EncryptingFileIO.combine(table.io(), table.encryption()) + .newEncryptingOutputFile(location); + } + + private static Record cdcRecord(long id, String data) { + GenericRecord record = GenericRecord.create(CDC_SCHEMA); + record.setField("id", id); + record.setField("data", data); + return record; + } + + private IcebergScanConfig cdcScanConfig(Table table) { + return IcebergScanConfig.builder() + .setCatalogConfig(catalogConfig) + .setTableIdentifier(tableId) + .setSchema(IcebergUtils.icebergSchemaToBeamSchema(table.schema())) + .setKeepFields(ImmutableList.of("id")) + .build(); + } + + private static List idsOf(CloseableIterable records) { + return ImmutableList.copyOf(records).stream() + .map(record -> (Long) record.getField("id")) + .collect(Collectors.toList()); + } + + /** + * Asserts that {@code dataFile} uses Parquet modular encryption, which is the layout every + * Iceberg engine expects of an encrypted Parquet file. + * + *

The footer is unreadable without a key, and reading it with the key from {@code + * key_metadata} reports an encrypted footer. + */ + private static void assertDataFileIsEncrypted(Table table, ContentFile dataFile) + throws IOException { + assertThrows( + "data file " + dataFile.location() + " is readable without a key", + ParquetCryptoRuntimeException.class, + () -> ParquetFileReader.open(parquetInputFile(dataFile.location())).close()); + + EncryptedInputFile encrypted = + EncryptedFiles.encryptedInput( + table.io().newInputFile(dataFile.location()), dataFile.keyMetadata()); + InputFile decrypted = table.encryption().decrypt(encrypted); + assertThat(decrypted, Matchers.instanceOf(NativeEncryptionInputFile.class)); + NativeEncryptionInputFile nativeFile = (NativeEncryptionInputFile) decrypted; + + ParquetReadOptions options = + ParquetReadOptions.builder() + .withDecryption( + FileDecryptionProperties.builder() + .withFooterKey( + ByteBuffers.toByteArray(nativeFile.keyMetadata().encryptionKey())) + .withAADPrefix(ByteBuffers.toByteArray(nativeFile.keyMetadata().aadPrefix())) + .build()) + .build(); + try (ParquetFileReader reader = + ParquetFileReader.open(parquetInputFile(dataFile.location()), options)) { + assertEquals( + "data file " + dataFile.location() + " does not use Parquet modular encryption", + FileMetaData.EncryptionType.ENCRYPTED_FOOTER, + reader.getFooter().getFileMetaData().getEncryptionType()); + } + } + + /** + * The same warehouse seen through a catalog with no KMS configured, i.e. what an engine without + * access to the encryption keys has. + */ + private Table loadTableWithoutKms() { + return IcebergCatalogConfig.builder() + .setCatalogName("no-kms") + .setCatalogProperties( + ImmutableMap.of( + "type", + CatalogUtil.ICEBERG_CATALOG_TYPE_HADOOP, + CatalogProperties.WAREHOUSE_LOCATION, + warehouse)) + .build() + .catalog() + .loadTable(tableId); + } + + /** A {@link FileIO} with no encryption manager attached. */ + private FileIO plaintextIo() { + return loadTableWithoutKms().io(); + } + + private static org.apache.parquet.io.InputFile parquetInputFile(String location) + throws IOException { + return HadoopInputFile.fromPath(new Path(URI.create(location)), new Configuration()); + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TestKms.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TestKms.java new file mode 100644 index 000000000000..4f2be0aef33f --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/TestKms.java @@ -0,0 +1,89 @@ +/* + * 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.beam.sdk.io.iceberg; + +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.security.SecureRandom; +import java.util.Map; +import javax.crypto.Cipher; +import javax.crypto.spec.GCMParameterSpec; +import javax.crypto.spec.SecretKeySpec; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.iceberg.encryption.KeyManagementClient; +import org.apache.iceberg.util.ByteBuffers; + +/** + * An in-memory {@link KeyManagementClient} for tests, so that table encryption can be exercised + * without reaching a real KMS service. + * + *

Keys are wrapped with AES-GCM under a hard-coded master key, mirroring what Iceberg's own + * {@code UnitestKMS} test fixture does. + */ +public class TestKms implements KeyManagementClient { + /** Table key ID to use as the {@code encryption.key-id} table property. */ + public static final String MASTER_KEY_ID = "beam-test-master-key"; + + private static final Map MASTER_KEYS = + ImmutableMap.of(MASTER_KEY_ID, "0123456789012345".getBytes(StandardCharsets.UTF_8)); + + private static final String TRANSFORMATION = "AES/GCM/NoPadding"; + private static final int NONCE_LENGTH = 12; + private static final int TAG_LENGTH_BITS = 128; + + @Override + public ByteBuffer wrapKey(ByteBuffer key, String wrappingKeyId) { + byte[] nonce = new byte[NONCE_LENGTH]; + new SecureRandom().nextBytes(nonce); + byte[] ciphertext = + crypt(Cipher.ENCRYPT_MODE, wrappingKeyId, nonce, ByteBuffers.toByteArray(key)); + + // prepend the nonce so unwrapKey can recover it + return ByteBuffer.allocate(nonce.length + ciphertext.length).put(nonce).put(ciphertext).flip(); + } + + @Override + public ByteBuffer unwrapKey(ByteBuffer wrappedKey, String wrappingKeyId) { + byte[] wrapped = ByteBuffers.toByteArray(wrappedKey); + byte[] nonce = new byte[NONCE_LENGTH]; + System.arraycopy(wrapped, 0, nonce, 0, NONCE_LENGTH); + byte[] ciphertext = new byte[wrapped.length - NONCE_LENGTH]; + System.arraycopy(wrapped, NONCE_LENGTH, ciphertext, 0, ciphertext.length); + + return ByteBuffer.wrap(crypt(Cipher.DECRYPT_MODE, wrappingKeyId, nonce, ciphertext)); + } + + @Override + public void initialize(Map properties) {} + + private static byte[] crypt(int mode, String wrappingKeyId, byte[] nonce, byte[] input) { + byte[] masterKey = MASTER_KEYS.get(wrappingKeyId); + if (masterKey == null) { + throw new IllegalArgumentException("Unknown master key ID: " + wrappingKeyId); + } + try { + Cipher cipher = Cipher.getInstance(TRANSFORMATION); + cipher.init( + mode, new SecretKeySpec(masterKey, "AES"), new GCMParameterSpec(TAG_LENGTH_BITS, nonce)); + return cipher.doFinal(input); + } catch (Exception e) { + throw new RuntimeException( + "Failed to " + (mode == Cipher.ENCRYPT_MODE ? "wrap" : "unwrap"), e); + } + } +} diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/TestChangelogTasks.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/TestChangelogTasks.java new file mode 100644 index 000000000000..3d92ead147c4 --- /dev/null +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/TestChangelogTasks.java @@ -0,0 +1,61 @@ +/* + * 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.beam.sdk.io.iceberg.cdc; + +import java.util.List; +import java.util.stream.Collectors; +import org.apache.beam.sdk.io.iceberg.SerializableDeleteFile; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.iceberg.ChangelogOperation; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.Table; +import org.apache.iceberg.expressions.ExpressionParser; +import org.apache.iceberg.expressions.Expressions; + +/** Builds {@link SerializableChangelogTask}s for tests outside this package. */ +public class TestChangelogTasks { + private TestChangelogTasks() {} + + public static SerializableChangelogTask addedRows( + Table table, DataFile dataFile, List addedDeletes) { + return SerializableChangelogTask.builder() + .setType(SerializableChangelogTask.Type.ADDED_ROWS) + .setDataFile(dataFile, table.spec(), true) + .setAddedDeletes(serializableDeletes(table, addedDeletes)) + .setExistingDeletes(ImmutableList.of()) + .setSpecId(table.spec().specId()) + .setOperation(ChangelogOperation.INSERT) + .setOrdinal(0) + .setCommitSnapshotId(1L) + .setStart(0L) + .setLength(dataFile.fileSizeInBytes()) + .setJsonExpression(ExpressionParser.toJson(Expressions.alwaysTrue())) + .build(); + } + + private static List serializableDeletes( + Table table, List deletes) { + return deletes.stream() + .map( + delete -> + SerializableDeleteFile.from( + delete, table.spec().partitionToPath(delete.partition()), true)) + .collect(Collectors.toList()); + } +} From ad8aa168beed5a32ae7e67c4c9f9689d2d5cb643 Mon Sep 17 00:00:00 2001 From: Ahmed Abualsaud Date: Tue, 18 Aug 2026 12:22:08 -0700 Subject: [PATCH 2/3] changes and trigger ITs --- .github/trigger_files/IO_Iceberg_Integration_Tests.json | 2 +- CHANGES.md | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/trigger_files/IO_Iceberg_Integration_Tests.json b/.github/trigger_files/IO_Iceberg_Integration_Tests.json index 7ab7bcd9a9c6..b73af5e61a43 100644 --- a/.github/trigger_files/IO_Iceberg_Integration_Tests.json +++ b/.github/trigger_files/IO_Iceberg_Integration_Tests.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", - "modification": 2 + "modification": 1 } diff --git a/CHANGES.md b/CHANGES.md index 73966a48313c..c377aaa54242 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -65,6 +65,7 @@ ## I/Os * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). +* [IcebergIO] is now compliant with encrypted tables ([39808](https://github.com/apache/beam/issues/39808)). ## New Features / Improvements From cff4aeb61f0c15ceb5593bfb50e43c3554f19be4 Mon Sep 17 00:00:00 2001 From: Ahmed Abualsaud Date: Tue, 18 Aug 2026 12:35:08 -0700 Subject: [PATCH 3/3] format changes --- CHANGES.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGES.md b/CHANGES.md index c377aaa54242..c6b2d8515fab 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -65,7 +65,7 @@ ## I/Os * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). -* [IcebergIO] is now compliant with encrypted tables ([39808](https://github.com/apache/beam/issues/39808)). +* [IcebergIO] is now compliant with encrypted tables ([#39808](https://github.com/apache/beam/issues/39808)). ## New Features / Improvements