Add uniform encryption test

Author: rok

parquet/src/arrow/arrow_writer/mod.rs

@@ -1290,16 +1290,14 @@ mod tests {
 
     use std::fs::File;
 
-    #[cfg(feature = "encryption")]
-    use crate::arrow::arrow_reader::ArrowReaderMetadata;
     #[cfg(feature = "encryption")]
     use crate::arrow::arrow_reader::ArrowReaderOptions;
     use crate::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
     use crate::arrow::ARROW_SCHEMA_META_KEY;
     #[cfg(feature = "encryption")]
     use crate::encryption::encrypt::EncryptionKey;
     #[cfg(feature = "encryption")]
-    use crate::util::test_common::encryption_util::verify_encryption_test_file_read;
+    use crate::util::test_common::encryption_util::read_and_roundtrip_to_encrypted_file;
     use arrow::datatypes::ToByteSlice;
     use arrow::datatypes::{DataType, Schema};
     use arrow::error::Result as ArrowResult;
@@ -3808,7 +3806,6 @@ mod tests {
     fn test_non_uniform_encryption() {
         let testdata = arrow::util::test_util::parquet_test_data();
         let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted");
-        let temp_file = tempfile::tempfile().unwrap();
 
         let footer_key = "0123456789012345".as_bytes(); // 128bit/16
         let column_1_key = "1234567890123450".as_bytes();
@@ -3820,62 +3817,39 @@ mod tests {
             .build()
             .unwrap();
 
-        // read example data
-        let (batches, schema) = read_encrypted_file(&path, decryption_properties.clone());
-
-        // write example data
         let column_1_key = EncryptionKey::new(column_1_key.as_bytes().to_vec());
         let column_2_key = EncryptionKey::new(column_2_key.as_bytes().to_vec());
         let file_encryption_properties = FileEncryptionProperties::builder(footer_key.to_vec())
             .with_column_key("double_field".into(), column_1_key)
             .with_column_key("float_field".into(), column_2_key)
             .build();
 
-        let temp_file =
-            write_batches_to_encrypted_file(temp_file, file_encryption_properties, batches, schema);
-
-        // check re-written example data
-        verify_encryption_test_file_read(temp_file, decryption_properties);
+        read_and_roundtrip_to_encrypted_file(
+            &path,
+            decryption_properties,
+            file_encryption_properties,
+        );
     }
 
+    #[test]
     #[cfg(feature = "encryption")]
-    fn write_batches_to_encrypted_file(
-        temp_file: File,
-        file_encryption_properties: FileEncryptionProperties,
-        batches: Vec<RecordBatch>,
-        schema: SchemaRef,
-    ) -> File {
-        let props = WriterProperties::builder()
-            .with_file_encryption_properties(file_encryption_properties)
-            .build();
+    fn test_uniform_encryption() {
+        let testdata = arrow::util::test_util::parquet_test_data();
+        let path = format!("{testdata}/uniform_encryption.parquet.encrypted");
 
-        let mut writer =
-            ArrowWriter::try_new(temp_file.try_clone().unwrap(), schema, Some(props)).unwrap();
-        for batch in batches {
-            writer.write(&batch).unwrap();
-        }
+        let footer_key = "0123456789012345".as_bytes(); // 128bit/16
 
-        writer.close().unwrap();
-        temp_file
-    }
+        let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec())
+            .build()
+            .unwrap();
 
-    #[cfg(feature = "encryption")]
-    fn read_encrypted_file(
-        path: &str,
-        decryption_properties: FileDecryptionProperties,
-    ) -> (Vec<RecordBatch>, SchemaRef) {
-        let file = File::open(path).unwrap();
-        let options =
-            ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties);
-        let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap();
+        let file_encryption_properties =
+            FileEncryptionProperties::builder(footer_key.to_vec()).build();
 
-        let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap();
-        let batch_reader = builder.build().unwrap();
-        (
-            batch_reader
-                .collect::<Result<Vec<RecordBatch>, _>>()
-                .unwrap(),
-            metadata.schema,
-        )
+        read_and_roundtrip_to_encrypted_file(
+            &path,
+            decryption_properties,
+            file_encryption_properties,
+        );
     }
 }

parquet/src/file/metadata/writer.rs

@@ -317,7 +317,7 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> {
         column_index: usize,
     ) -> Result<ColumnChunk> {
         // Column crypto metadata should have already been set when the column was created.
-        // Here we apply the encryption by encrypted the column metadata if required.
+        // Here we apply the encryption by encrypting the column metadata if required.
         match column_chunk.crypto_metadata.as_ref() {
             None => {}
             Some(ColumnCryptoMetaData::ENCRYPTIONWITHFOOTERKEY(_)) => {

parquet/src/util/test_common/encryption_util.rs

@@ -26,6 +26,9 @@ use arrow_array::cast::AsArray;
 use arrow_array::{types, RecordBatch};
 use futures::TryStreamExt;
 use std::fs::File;
+use crate::arrow::ArrowWriter;
+use crate::encryption::encrypt::FileEncryptionProperties;
+use crate::file::properties::WriterProperties;
 
 /// Tests reading an encrypted file from the parquet-testing repository
 pub(crate) fn verify_encryption_test_file_read(
@@ -134,3 +137,34 @@ fn verify_encryption_test_data(
 
     assert_eq!(row_count, file_metadata.num_rows() as usize);
 }
+
+#[cfg(feature = "encryption")]
+pub fn read_and_roundtrip_to_encrypted_file(path: &str, decryption_properties: FileDecryptionProperties, encryption_properties: FileEncryptionProperties) {
+    let temp_file = tempfile::tempfile().unwrap();
+
+    // read example data
+    let file = File::open(path).unwrap();
+    let options =
+        ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties.clone());
+    let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap();
+
+    let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap();
+    let batch_reader = builder.build().unwrap();
+    let batches = batch_reader.collect::<crate::errors::Result<Vec<RecordBatch>, _>>().unwrap();
+
+    // write example data
+    let props = WriterProperties::builder()
+        .with_file_encryption_properties(encryption_properties)
+        .build();
+
+    let mut writer =
+        ArrowWriter::try_new(temp_file.try_clone().unwrap(), metadata.schema, Some(props)).unwrap();
+    for batch in batches {
+        writer.write(&batch).unwrap();
+    }
+
+    writer.close().unwrap();
+
+    // check re-written example data
+    verify_encryption_test_file_read(temp_file, decryption_properties);
+}