[AMORO-2401] Support reading encrypted iceberg data files (apache#2402)

Support reading iceberg encrypted files
huyuanfeng2018 · Mar 8, 2024 · a294917 · a294917
1 parent 168688b
commit a294917
Show file tree

Hide file tree

Showing 7 changed files with 39 additions and 7 deletions.
diff --git a/core/src/main/java/com/netease/arctic/io/reader/BaseIcebergPosDeleteReader.java b/core/src/main/java/com/netease/arctic/io/reader/BaseIcebergPosDeleteReader.java
@@ -27,6 +27,9 @@
 import org.apache.iceberg.data.Record;
 import org.apache.iceberg.data.orc.GenericOrcReader;
 import org.apache.iceberg.data.parquet.GenericParquetReaders;
+import org.apache.iceberg.encryption.EncryptedFiles;
+import org.apache.iceberg.encryption.EncryptedInputFile;
+import org.apache.iceberg.encryption.EncryptionManager;
 import org.apache.iceberg.io.CloseableIterable;
 import org.apache.iceberg.io.InputFile;
 import org.apache.iceberg.orc.ORC;
@@ -46,10 +49,13 @@ public class BaseIcebergPosDeleteReader {
       POS_DELETE_SCHEMA.accessorForField(MetadataColumns.DELETE_FILE_POS.fieldId());
 
   protected final ArcticFileIO fileIO;
+  protected final EncryptionManager encryptionManager;
   protected final List<DeleteFile> posDeleteFiles;
 
-  public BaseIcebergPosDeleteReader(ArcticFileIO fileIO, List<DeleteFile> posDeleteFiles) {
+  public BaseIcebergPosDeleteReader(
+      ArcticFileIO fileIO, EncryptionManager encryptionManager, List<DeleteFile> posDeleteFiles) {
     this.fileIO = fileIO;
+    this.encryptionManager = encryptionManager;
     this.posDeleteFiles = posDeleteFiles;
   }
 
@@ -67,7 +73,11 @@ public Long readPos(Record record) {
   }
 
   private CloseableIterable<Record> readDelete(DeleteFile deleteFile) {
-    InputFile input = fileIO.newInputFile(deleteFile.path().toString());
+    EncryptedInputFile encryptedInput =
+        EncryptedFiles.encryptedInput(
+            fileIO.newInputFile(deleteFile.path().toString()), deleteFile.keyMetadata());
+    InputFile input = encryptionManager.decrypt(encryptedInput);
+
     switch (deleteFile.format()) {
       case PARQUET:
         Parquet.ReadBuilder builder =

diff --git a/core/src/main/java/com/netease/arctic/io/reader/CombinedDeleteFilter.java b/core/src/main/java/com/netease/arctic/io/reader/CombinedDeleteFilter.java
@@ -172,7 +172,7 @@ public boolean isFilterEqDelete() {
     return filterEqDelete;
   }
 
-  protected abstract InputFile getInputFile(String location);
+  protected abstract InputFile getInputFile(ContentFile<?> contentFile);
 
   protected abstract ArcticFileIO getArcticFileIo();
 
@@ -387,7 +387,7 @@ private CloseableIterable<Record> openPosDeletes(DeleteFile file) {
   }
 
   private CloseableIterable<Record> openFile(ContentFile<?> contentFile, Schema deleteSchema) {
-    InputFile input = getInputFile(contentFile.path().toString());
+    InputFile input = getInputFile(contentFile);
     switch (contentFile.format()) {
       case AVRO:
         return Avro.read(input)

diff --git a/core/src/main/java/com/netease/arctic/io/reader/GenericCombinedIcebergDataReader.java b/core/src/main/java/com/netease/arctic/io/reader/GenericCombinedIcebergDataReader.java
@@ -23,6 +23,7 @@
 import com.netease.arctic.optimizing.RewriteFilesInput;
 import com.netease.arctic.scan.CombinedIcebergScanTask;
 import com.netease.arctic.utils.map.StructLikeCollections;
+import org.apache.iceberg.ContentFile;
 import org.apache.iceberg.DataFile;
 import org.apache.iceberg.MetadataColumns;
 import org.apache.iceberg.PartitionSpec;
@@ -32,6 +33,9 @@
 import org.apache.iceberg.data.avro.DataReader;
 import org.apache.iceberg.data.orc.GenericOrcReader;
 import org.apache.iceberg.data.parquet.GenericParquetReaders;
+import org.apache.iceberg.encryption.EncryptedFiles;
+import org.apache.iceberg.encryption.EncryptedInputFile;
+import org.apache.iceberg.encryption.EncryptionManager;
 import org.apache.iceberg.io.CloseableIterable;
 import org.apache.iceberg.io.InputFile;
 import org.apache.iceberg.parquet.Parquet;
@@ -61,6 +65,7 @@ public class GenericCombinedIcebergDataReader implements OptimizingDataReader {
   protected final String nameMapping;
   protected final boolean caseSensitive;
   protected final ArcticFileIO fileIO;
+  protected final EncryptionManager encryptionManager;
   protected final BiFunction<Type, Object, Object> convertConstant;
   protected final boolean reuseContainer;
   protected CombinedDeleteFilter<Record> deleteFilter;
@@ -73,6 +78,7 @@ public GenericCombinedIcebergDataReader(
       ArcticFileIO fileIO,
       Schema tableSchema,
       PartitionSpec spec,
+      EncryptionManager encryptionManager,
       String nameMapping,
       boolean caseSensitive,
       BiFunction<Type, Object, Object> convertConstant,
@@ -81,6 +87,7 @@ public GenericCombinedIcebergDataReader(
       RewriteFilesInput rewriteFilesInput) {
     this.tableSchema = tableSchema;
     this.spec = spec;
+    this.encryptionManager = encryptionManager;
     this.nameMapping = nameMapping;
     this.caseSensitive = caseSensitive;
     this.fileIO = fileIO;
@@ -164,7 +171,10 @@ private CloseableIterable<Record> openFile(
 
   private CloseableIterable<Record> openFile(
       DataFile dataFile, Schema fileProjection, Map<Integer, ?> idToConstant) {
-    InputFile input = fileIO.newInputFile(dataFile.path().toString());
+    EncryptedInputFile encryptedInput =
+        EncryptedFiles.encryptedInput(
+            fileIO.newInputFile(dataFile.path().toString()), dataFile.keyMetadata());
+    InputFile input = encryptionManager.decrypt(encryptedInput);
 
     switch (dataFile.format()) {
       case AVRO:
@@ -296,8 +306,11 @@ public GenericDeleteFilter(
     }
 
     @Override
-    protected InputFile getInputFile(String location) {
-      return fileIO.newInputFile(location);
+    protected InputFile getInputFile(ContentFile<?> contentFile) {
+      EncryptedInputFile encryptedInput =
+          EncryptedFiles.encryptedInput(
+              fileIO.newInputFile(contentFile.path().toString()), contentFile.keyMetadata());
+      return encryptionManager.decrypt(encryptedInput);
     }
 
     @Override

diff --git a/core/src/main/java/com/netease/arctic/optimizing/IcebergRewriteExecutor.java b/core/src/main/java/com/netease/arctic/optimizing/IcebergRewriteExecutor.java
@@ -64,6 +64,7 @@ protected OptimizingDataReader dataReader() {
         io,
         table.schema(),
         table.spec(),
+        table.asUnkeyedTable().encryption(),
         table.properties().get(TableProperties.DEFAULT_NAME_MAPPING),
         false,
         IdentityPartitionConverters::convertConstant,

diff --git a/core/src/test/java/com/netease/arctic/io/TestIcebergCombinedReader.java b/core/src/test/java/com/netease/arctic/io/TestIcebergCombinedReader.java
@@ -191,6 +191,7 @@ public void readAllData() throws IOException {
             getArcticTable().io(),
             getArcticTable().schema(),
             getArcticTable().spec(),
+            getArcticTable().asUnkeyedTable().encryption(),
             null,
             false,
             IdentityPartitionConverters::convertConstant,
@@ -212,6 +213,7 @@ public void readAllDataNegate() throws IOException {
             getArcticTable().io(),
             getArcticTable().schema(),
             getArcticTable().spec(),
+            getArcticTable().asUnkeyedTable().encryption(),
             null,
             false,
             IdentityPartitionConverters::convertConstant,
@@ -235,6 +237,7 @@ public void readOnlyData() throws IOException {
             getArcticTable().io(),
             getArcticTable().schema(),
             getArcticTable().spec(),
+            getArcticTable().asUnkeyedTable().encryption(),
             null,
             false,
             IdentityPartitionConverters::convertConstant,
@@ -254,6 +257,7 @@ public void readOnlyDataNegate() throws IOException {
             getArcticTable().io(),
             getArcticTable().schema(),
             getArcticTable().spec(),
+            getArcticTable().asUnkeyedTable().encryption(),
             null,
             false,
             IdentityPartitionConverters::convertConstant,
@@ -274,6 +278,7 @@ public void readDataEnableFilterEqDelete() throws IOException {
             getArcticTable().io(),
             getArcticTable().schema(),
             getArcticTable().spec(),
+            getArcticTable().asUnkeyedTable().encryption(),
             null,
             false,
             IdentityPartitionConverters::convertConstant,

diff --git a/core/src/test/java/com/netease/arctic/io/TestIcebergCombinedReaderVariousTypes.java b/core/src/test/java/com/netease/arctic/io/TestIcebergCombinedReaderVariousTypes.java
@@ -139,6 +139,7 @@ public void valid() throws IOException {
                 table.io(),
                 table.schema(),
                 table.spec(),
+                table.encryption(),
                 null,
                 false,
                 IdentityPartitionConverters::convertConstant,
@@ -189,6 +190,7 @@ public void readDataEnableFilterEqDelete() throws IOException {
             table.io(),
             table.schema(),
             table.spec(),
+            table.encryption(),
             null,
             false,
             IdentityPartitionConverters::convertConstant,

diff --git a/core/src/test/java/com/netease/arctic/io/TestTaskReader.java b/core/src/test/java/com/netease/arctic/io/TestTaskReader.java
@@ -197,6 +197,7 @@ public void testReadPosDelete() {
     BaseIcebergPosDeleteReader baseIcebergPosDeleteReader =
         new BaseIcebergPosDeleteReader(
             getArcticTable().asKeyedTable().io(),
+            getArcticTable().asKeyedTable().baseTable().encryption(),
             Collections.singletonList(deleteFileOfPositionDelete));
     ImmutableList.Builder<Record> builder = ImmutableList.builder();
     baseIcebergPosDeleteReader.readDeletes().forEach(record -> builder.add(record.copy()));