[IcebergIO] Ensure compliance with table encryption

.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
 }

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
 

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<DataFile> writer = ManifestFiles.write(spec, outputFile)) {
+      try (ManifestWriter<DataFile> 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<DataFile> createManifestWriter(
+        Table table, PartitionSpec spec, String manifestPath) {
+      EncryptedOutputFile outputFile =
+          EncryptingFileIO.combine(table.io(), table.encryption())
+              .newEncryptingOutputFile(manifestPath);
+      return ManifestFiles.write(1, spec, outputFile, null);
+    }
   }
 
   /**

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<FileWriteResult> fileWrit
         int specId = entry.getKey();
         List<DataFile> files = entry.getValue();
         PartitionSpec spec = Preconditions.checkStateNotNull(specs.get(specId));
-        FileIO io = table.io();
-        ManifestWriter<DataFile> writer = createManifestWriter(table.location(), uuid, spec, io);
+        ManifestWriter<DataFile> 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<FileWriteResult> fileWrit
       update.commit();
     }
 
+    @SuppressWarnings("argument")
     private ManifestWriter<DataFile> 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

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<Record> 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<Record> createReader(
       Table table,
       IcebergScanConfig scanConfig,
@@ -91,7 +116,10 @@ public static CloseableIterable<Record> 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<Integer, ?> idToConstants = PartitionUtils.constantsMap(spec, file, fileSequenceNumber);
 
     ParquetReadOptions.Builder optionsBuilder;
@@ -109,6 +137,9 @@ public static CloseableIterable<Record> 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 =

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)

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<Record> 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<Record> genericDeleteFilter(
       Table table, Schema outputSchema, String dataFilePath, List<SerializableDeleteFile> deletes) {
     return new GenericDeleteFilter(
-        table.io(),
+        encryptingIO(table),
         dataFilePath,
         table.schema(),
         outputSchema,
@@ -162,7 +180,7 @@ public static DeleteReader<Record> genericDeleteReader(
       List<SerializableDeleteFile> deletes,
       DeleteReader.PreloadedDeletes preloadedDeletes) {
     return new GenericDeleteReader(
-        table.io(),
+        encryptingIO(table),
         dataFilePath,
         table.schema(),
         outputSchema,
@@ -340,7 +358,7 @@ private static CloseableIterable<Record> 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<Record> 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<Record> {
-    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<Record> {
-    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);
+    }
   }
 
   /**