Add FileEncryptionProperties to WriterProperties

Author: rok

parquet/src/column/writer/mod.rs

@@ -1535,6 +1535,7 @@ mod tests {
         page::PageReader,
         reader::{get_column_reader, get_typed_column_reader, ColumnReaderImpl},
     };
+    use crate::encryption::encryption::FileEncryptionProperties;
     use crate::file::writer::TrackedWrite;
     use crate::file::{
         properties::ReaderProperties, reader::SerializedPageReader, writer::SerializedPageWriter,
@@ -3379,6 +3380,31 @@ mod tests {
         );
     }
 
+    #[cfg(feature = "encryption")]
+    #[test]
+    fn test_encryption_writer() {
+        let message_type = "
+            message test_schema {
+                OPTIONAL BYTE_ARRAY a (UTF8);
+            }
+        ";
+        let schema = Arc::new(parse_message_type(message_type).unwrap());
+        let file: File = tempfile::tempfile().unwrap();
+
+        let builder = WriterProperties::builder();
+        let key_code: &[u8] = "0123456789012345".as_bytes();
+        let file_encryption_properties = FileEncryptionProperties::builder(key_code.to_vec())
+            .build()
+            .unwrap();
+
+        let props = Arc::new(
+            builder
+                .set_file_encryption_properties(file_encryption_properties)
+                .build(),
+        );
+        let mut writer = SerializedFileWriter::new(&file, schema, props).unwrap();
+    }
+
     #[test]
     fn test_increment_max_binary_chars() {
         let r = increment(vec![0xFF, 0xFE, 0xFD, 0xFF, 0xFF]);

parquet/src/encryption/encryption.rs

@@ -0,0 +1,57 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::collections::HashMap;
+
+#[derive(Debug, Clone)]
+pub struct FileEncryptionProperties {
+    encrypt_footer: bool,
+    footer_key: Vec<u8>,
+    column_keys: Option<HashMap<Vec<u8>, Vec<u8>>>,
+    aad_prefix: Option<Vec<u8>>,
+}
+
+impl FileEncryptionProperties {
+    pub fn builder(footer_key: Vec<u8>) -> EncryptionPropertiesBuilder {
+        EncryptionPropertiesBuilder::new(footer_key)
+    }
+}
+
+pub struct EncryptionPropertiesBuilder {
+    footer_key: Vec<u8>,
+    column_keys: Option<HashMap<Vec<u8>, Vec<u8>>>,
+    aad_prefix: Option<Vec<u8>>,
+}
+
+impl EncryptionPropertiesBuilder {
+    pub fn new(footer_key: Vec<u8>) -> EncryptionPropertiesBuilder {
+        Self {
+            footer_key,
+            column_keys: None,
+            aad_prefix: None,
+        }
+    }
+
+    pub fn build(self) -> crate::errors::Result<FileEncryptionProperties> {
+        Ok(FileEncryptionProperties {
+            encrypt_footer: true,
+            footer_key: self.footer_key,
+            column_keys: self.column_keys,
+            aad_prefix: self.aad_prefix,
+        })
+    }
+}

parquet/src/encryption/mod.rs

@@ -20,4 +20,5 @@
 
 pub mod ciphers;
 pub mod decryption;
+pub mod encryption;
 pub mod modules;

parquet/src/file/properties.rs

@@ -18,6 +18,7 @@
 //! Configuration via [`WriterProperties`] and [`ReaderProperties`]
 use crate::basic::{Compression, Encoding};
 use crate::compression::{CodecOptions, CodecOptionsBuilder};
+use crate::encryption::encryption::FileEncryptionProperties;
 use crate::file::metadata::KeyValue;
 use crate::format::SortingColumn;
 use crate::schema::types::ColumnPath;
@@ -169,6 +170,7 @@ pub struct WriterProperties {
     column_index_truncate_length: Option<usize>,
     statistics_truncate_length: Option<usize>,
     coerce_types: bool,
+    file_encryption_properties: Option<FileEncryptionProperties>,
 }
 
 impl Default for WriterProperties {
@@ -392,6 +394,7 @@ pub struct WriterPropertiesBuilder {
     column_index_truncate_length: Option<usize>,
     statistics_truncate_length: Option<usize>,
     coerce_types: bool,
+    file_encryption_properties: Option<FileEncryptionProperties>,
 }
 
 impl WriterPropertiesBuilder {
@@ -414,6 +417,7 @@ impl WriterPropertiesBuilder {
             column_index_truncate_length: DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH,
             statistics_truncate_length: DEFAULT_STATISTICS_TRUNCATE_LENGTH,
             coerce_types: DEFAULT_COERCE_TYPES,
+            file_encryption_properties: None,
         }
     }
 
@@ -436,6 +440,7 @@ impl WriterPropertiesBuilder {
             column_index_truncate_length: self.column_index_truncate_length,
             statistics_truncate_length: self.statistics_truncate_length,
             coerce_types: self.coerce_types,
+            file_encryption_properties: self.file_encryption_properties,
         }
     }
 
@@ -808,6 +813,16 @@ impl WriterPropertiesBuilder {
         self.coerce_types = coerce_types;
         self
     }
+
+    /// Sets FileEncryptionProperties.
+    /// Only applicable if file encryption is enabled.
+    pub fn set_file_encryption_properties(
+        mut self,
+        file_encryption_properties: FileEncryptionProperties,
+    ) -> Self {
+        self.file_encryption_properties = Some(file_encryption_properties);
+        self
+    }
 }
 
 /// Controls the level of statistics to be computed by the writer and stored in