rok commented on code in PR #7600:
URL: https://github.com/apache/arrow-rs/pull/7600#discussion_r2126976711
##########
parquet/tests/encryption/encryption.rs:
##########
@@ -256,6 +256,72 @@ fn
test_non_uniform_encryption_plaintext_footer_with_key_retriever() {
verify_encryption_test_file_read(file, decryption_properties);
}
+#[test]
+fn test_uniform_encryption_plaintext_footer_with_key_retriever() {
+ let test_data = arrow::util::test_util::parquet_test_data();
+ let path =
format!("{test_data}/encrypt_columns_plaintext_footer.parquet.encrypted");
+ let file = File::open(path).unwrap();
+ let temp_file = tempfile::tempfile().unwrap();
+
+ let key_retriever = TestKeyRetriever::new()
+ .with_key("kf".to_owned(), b"0123456789012345".to_vec())
+ .with_key("kc1".to_owned(), b"1234567890123450".to_vec())
+ .with_key("kc2".to_owned(), b"1234567890123451".to_vec());
+
+ let encryption_properties =
FileEncryptionProperties::builder(b"0123456789012345".to_vec())
+ .with_footer_key_metadata("kf".into())
+ .with_column_key_and_metadata("double_field",
b"1234567890123450".to_vec(), b"kc1".into())
+ .with_column_key_and_metadata("float_field",
b"1234567890123451".to_vec(), b"kc2".into())
+ .with_plaintext_footer(true)
+ .build()
+ .unwrap();
+
+ let decryption_properties =
+ FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever))
+ .build()
+ .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::<parquet::errors::Result<Vec<RecordBatch>, _>>()
+ .unwrap();
+
+ let props = WriterProperties::builder()
+ .with_file_encryption_properties(encryption_properties)
+ .build();
+
+ let mut writer = ArrowWriter::try_new(
+ temp_file.try_clone().unwrap(),
+ metadata.schema().clone(),
+ Some(props),
+ )
+ .unwrap();
+ for batch in batches {
+ writer.write(&batch).unwrap();
+ }
+
+ writer.close().unwrap();
+
+ let key_retriever = TestKeyRetriever::new()
Review Comment:
Oh, wasn't aware of that mechanism. Done.
--
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]
