ggershinsky commented on code in PR #16724:
URL: https://github.com/apache/iceberg/pull/16724#discussion_r3381351684


##########
data/src/test/java/org/apache/iceberg/data/BaseFormatModelTests.java:
##########
@@ -1706,6 +1715,98 @@ void testDataWriterMetaMap(FileFormat fileFormat) throws 
IOException {
         });
   }
 
+  @ParameterizedTest
+  @FieldSource("FILE_FORMATS")
+  void testDataWriterEncryption(FileFormat fileFormat) throws IOException {
+    assumeSupports(fileFormat, FEATURE_ENCRYPTION);
+
+    EncryptionManager encryptionManager = 
EncryptionTestHelpers.createEncryptionManager();
+    EncryptingFileIO encryptingFileIO = EncryptingFileIO.combine(fileIO, 
encryptionManager);
+    EncryptedOutputFile encryptedOutputFile =
+        encryptingFileIO.newEncryptingOutputFile("test-file-" + 
UUID.randomUUID());
+
+    FileWriterBuilder<DataWriter<T>, ?> writerBuilder =
+        FormatModelRegistry.dataWriteBuilder(fileFormat, engineType(), 
encryptedOutputFile)
+            .keyMetadata(encryptedOutputFile.keyMetadata());
+
+    writeAndAssertEncryptedDataWriter(fileFormat, encryptingFileIO, 
writerBuilder);
+  }
+
+  @ParameterizedTest
+  @FieldSource("FILE_FORMATS")
+  void testDataWriterWithFileEncryptionKeyAndWithAADPrefix(FileFormat 
fileFormat)
+      throws IOException {
+    assumeSupports(fileFormat, FEATURE_ENCRYPTION);
+
+    EncryptionManager encryptionManager = 
EncryptionTestHelpers.createEncryptionManager();
+    EncryptingFileIO encryptingFileIO = EncryptingFileIO.combine(fileIO, 
encryptionManager);
+    String location = "test-file-" + UUID.randomUUID();
+    NativeEncryptionKeyMetadata keyMetadata =
+        (NativeEncryptionKeyMetadata)
+            encryptingFileIO.newEncryptingOutputFile(location).keyMetadata();
+
+    // Use a plain encrypted output file so Parquet cannot auto-inject the 
native encryption key
+    // and AAD prefix from the output file metadata. This ensures encryption 
is driven only by
+    // withFileEncryptionKey and withAADPrefix below.
+    EncryptedOutputFile encryptedOutputFile =
+        EncryptedFiles.plainAsEncryptedOutput(fileIO.newOutputFile(location));
+
+    FileWriterBuilder<DataWriter<T>, ?> writerBuilder =
+        FormatModelRegistry.dataWriteBuilder(fileFormat, engineType(), 
encryptedOutputFile)
+            .keyMetadata(keyMetadata)
+            .withFileEncryptionKey(keyMetadata.encryptionKey().duplicate())
+            .withAADPrefix(keyMetadata.aadPrefix().duplicate());
+
+    writeAndAssertEncryptedDataWriter(fileFormat, encryptingFileIO, 
writerBuilder);
+  }
+
+  @SuppressWarnings("checkstyle:AssertThatThrownByWithMessageCheck")
+  private void writeAndAssertEncryptedDataWriter(
+      FileFormat fileFormat,
+      EncryptingFileIO encryptingFileIO,
+      FileWriterBuilder<DataWriter<T>, ?> writerBuilder)
+      throws IOException {
+    DataGenerator dataGenerator = new DataGenerators.DefaultSchema();
+    Schema schema = dataGenerator.schema();
+    List<Record> genericRecords = dataGenerator.generateRecords();
+    List<T> engineRecords = convertToEngineRecords(genericRecords, schema);
+
+    DataWriter<T> writer = 
writerBuilder.schema(schema).spec(PartitionSpec.unpartitioned()).build();
+
+    try (writer) {
+      engineRecords.forEach(writer::write);
+    }
+
+    DataFile dataFile = writer.toDataFile();
+    assertThat(dataFile).isNotNull();
+    assertThat(dataFile.recordCount()).isEqualTo(engineRecords.size());
+    assertThat(dataFile.format()).isEqualTo(fileFormat);
+    assertThat(dataFile.keyMetadata()).isNotNull();
+
+    assertThatThrownBy(

Review Comment:
   SGTM for ORC.
   Encryption of tables with Avro data format is supported,  
https://github.com/apache/iceberg/pull/9436



-- 
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]

Reply via email to