feat: Dynamic Parquet encryption and decryption properties (#16779)

* Add encryption factory API for more flexible encryption configuration

* Remove extra DynEncryptionFactory trait

* Tidy up example

* Tidy ups

* Fix "unnecessary qualification" errors in example

* Fix toml format with taplo

* Fix broken link in docs

* Clippy fixes

* Update examples README

* Add extra validation of table encryption

* clippy fix

* Remove register_parquet_encryption_factory from SessionContext

* Remove new dependency from example

* Update examples readme

* Run taplo format

* Fix outdated method reference in comment

* Extra comment

---------

Co-authored-by: Andrew Lamb <[email protected]>

Author: adamreeve

Cargo.lock

@@ -65,6 +65,7 @@ async-trait = { workspace = true }
 bytes = { workspace = true }
 dashmap = { workspace = true }
 # note only use main datafusion crate for examples
+base64 = "0.22.1"
 datafusion = { workspace = true, default-features = true }
 datafusion-ffi = { workspace = true }
 datafusion-proto = { workspace = true }
@@ -74,6 +75,7 @@ log = { workspace = true }
 mimalloc = { version = "0.1", default-features = false }
 object_store = { workspace = true, features = ["aws", "http"] }
 prost = { workspace = true }
+rand = { workspace = true }
 serde_json = { workspace = true }
 tempfile = { workspace = true }
 test-utils = { path = "../test-utils" }

datafusion-examples/Cargo.toml

@@ -68,6 +68,7 @@ cargo run --example dataframe
 - [`optimizer_rule.rs`](examples/optimizer_rule.rs): Use a custom OptimizerRule to replace certain predicates
 - [`parquet_embedded_index.rs`](examples/parquet_embedded_index.rs): Store a custom index inside a Parquet file and use it to speed up queries
 - [`parquet_encrypted.rs`](examples/parquet_encrypted.rs): Read and write encrypted Parquet files using DataFusion
+- [`parquet_encrypted_with_kms.rs`](examples/parquet_encrypted_with_kms.rs): Read and write encrypted Parquet files using an encryption factory
 - [`parquet_index.rs`](examples/parquet_index.rs): Create an secondary index over several parquet files and use it to speed up queries
 - [`parquet_exec_visitor.rs`](examples/parquet_exec_visitor.rs): Extract statistics by visiting an ExecutionPlan after execution
 - [`parse_sql_expr.rs`](examples/parse_sql_expr.rs): Parse SQL text into DataFusion `Expr`.

datafusion-examples/README.md

@@ -0,0 +1,301 @@
+// 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 arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray};
+use arrow_schema::SchemaRef;
+use base64::Engine;
+use datafusion::common::extensions_options;
+use datafusion::config::{EncryptionFactoryOptions, TableParquetOptions};
+use datafusion::dataframe::DataFrameWriteOptions;
+use datafusion::datasource::file_format::parquet::ParquetFormat;
+use datafusion::datasource::listing::ListingOptions;
+use datafusion::error::Result;
+use datafusion::execution::parquet_encryption::EncryptionFactory;
+use datafusion::parquet::encryption::decrypt::KeyRetriever;
+use datafusion::parquet::encryption::{
+    decrypt::FileDecryptionProperties, encrypt::FileEncryptionProperties,
+};
+use datafusion::prelude::SessionContext;
+use futures::StreamExt;
+use object_store::path::Path;
+use rand::rand_core::{OsRng, TryRngCore};
+use std::collections::HashSet;
+use std::sync::Arc;
+use tempfile::TempDir;
+
+const ENCRYPTION_FACTORY_ID: &str = "example.mock_kms_encryption";
+
+/// This example demonstrates reading and writing Parquet files that
+/// are encrypted using Parquet Modular Encryption.
+///
+/// Compared to the `parquet_encrypted` example, where AES keys
+/// are specified directly, this example implements an `EncryptionFactory` that
+/// generates encryption keys dynamically per file.
+/// Encryption key metadata is stored inline in the Parquet files and is used to determine
+/// the decryption keys when reading the files.
+///
+/// In this example, encryption keys are simply stored base64 encoded in the Parquet metadata,
+/// which is not a secure way to store encryption keys.
+/// For production use, it is recommended to use a key-management service (KMS) to encrypt
+/// data encryption keys.
+#[tokio::main]
+async fn main() -> Result<()> {
+    let ctx = SessionContext::new();
+
+    // Register an `EncryptionFactory` implementation to be used for Parquet encryption
+    // in the runtime environment.
+    // `EncryptionFactory` instances are registered with a name to identify them so
+    // they can be later referenced in configuration options, and it's possible to register
+    // multiple different factories to handle different ways of encrypting Parquet.
+    let encryption_factory = TestEncryptionFactory::default();
+    ctx.runtime_env().register_parquet_encryption_factory(
+        ENCRYPTION_FACTORY_ID,
+        Arc::new(encryption_factory),
+    );
+
+    // Register some simple test data
+    let a: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d"]));
+    let b: ArrayRef = Arc::new(Int32Array::from(vec![1, 10, 10, 100]));
+    let c: ArrayRef = Arc::new(Int32Array::from(vec![2, 20, 20, 200]));
+    let batch = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)])?;
+    ctx.register_batch("test_data", batch)?;
+
+    {
+        // Write and read encrypted Parquet with the programmatic API
+        let tmpdir = TempDir::new()?;
+        let table_path = format!("{}/", tmpdir.path().to_str().unwrap());
+        write_encrypted(&ctx, &table_path).await?;
+        read_encrypted(&ctx, &table_path).await?;
+    }
+
+    {
+        // Write and read encrypted Parquet with the SQL API
+        let tmpdir = TempDir::new()?;
+        let table_path = format!("{}/", tmpdir.path().to_str().unwrap());
+        write_encrypted_with_sql(&ctx, &table_path).await?;
+        read_encrypted_with_sql(&ctx, &table_path).await?;
+    }
+
+    Ok(())
+}
+
+/// Write an encrypted Parquet file
+async fn write_encrypted(ctx: &SessionContext, table_path: &str) -> Result<()> {
+    let df = ctx.table("test_data").await?;
+
+    let mut parquet_options = TableParquetOptions::new();
+    // We specify that we want to use Parquet encryption by setting the identifier of the
+    // encryption factory to use and providing the factory-specific configuration.
+    // Our encryption factory only requires specifying the columns to encrypt.
+    let encryption_config = EncryptionConfig {
+        encrypted_columns: "b,c".to_owned(),
+    };
+    parquet_options
+        .crypto
+        .configure_factory(ENCRYPTION_FACTORY_ID, &encryption_config);
+
+    df.write_parquet(
+        table_path,
+        DataFrameWriteOptions::new(),
+        Some(parquet_options),
+    )
+    .await?;
+
+    println!("Encrypted Parquet written to {table_path}");
+    Ok(())
+}
+
+/// Read from an encrypted Parquet file
+async fn read_encrypted(ctx: &SessionContext, table_path: &str) -> Result<()> {
+    let mut parquet_options = TableParquetOptions::new();
+    // Specify the encryption factory to use for decrypting Parquet.
+    // In this example, we don't require any additional configuration options when reading
+    // as we only need the key metadata from the Parquet files to determine the decryption keys.
+    parquet_options
+        .crypto
+        .configure_factory(ENCRYPTION_FACTORY_ID, &EncryptionConfig::default());
+
+    let file_format = ParquetFormat::default().with_options(parquet_options);
+    let listing_options = ListingOptions::new(Arc::new(file_format));
+
+    ctx.register_listing_table(
+        "encrypted_parquet_table",
+        &table_path,
+        listing_options.clone(),
+        None,
+        None,
+    )
+    .await?;
+
+    let mut batch_stream = ctx
+        .table("encrypted_parquet_table")
+        .await?
+        .execute_stream()
+        .await?;
+    println!("Reading encrypted Parquet as a RecordBatch stream");
+    while let Some(batch) = batch_stream.next().await {
+        let batch = batch?;
+        println!("Read batch with {} rows", batch.num_rows());
+    }
+
+    println!("Finished reading");
+    Ok(())
+}
+
+/// Write an encrypted Parquet file using only SQL syntax with string configuration
+async fn write_encrypted_with_sql(ctx: &SessionContext, table_path: &str) -> Result<()> {
+    let query = format!(
+        "COPY test_data \
+        TO '{table_path}' \
+        STORED AS parquet
+        OPTIONS (\
+            'format.crypto.factory_id' '{ENCRYPTION_FACTORY_ID}', \
+            'format.crypto.factory_options.encrypted_columns' 'b,c' \
+        )"
+    );
+    let _ = ctx.sql(&query).await?.collect().await?;
+
+    println!("Encrypted Parquet written to {table_path}");
+    Ok(())
+}
+
+/// Read from an encrypted Parquet file using only the SQL API and string-based configuration
+async fn read_encrypted_with_sql(ctx: &SessionContext, table_path: &str) -> Result<()> {
+    let ddl = format!(
+        "CREATE EXTERNAL TABLE encrypted_parquet_table_2 \
+        STORED AS PARQUET LOCATION '{table_path}' OPTIONS (\
+        'format.crypto.factory_id' '{ENCRYPTION_FACTORY_ID}' \
+        )"
+    );
+    ctx.sql(&ddl).await?;
+    let df = ctx.sql("SELECT * FROM encrypted_parquet_table_2").await?;
+    let mut batch_stream = df.execute_stream().await?;
+
+    println!("Reading encrypted Parquet as a RecordBatch stream");
+    while let Some(batch) = batch_stream.next().await {
+        let batch = batch?;
+        println!("Read batch with {} rows", batch.num_rows());
+    }
+    println!("Finished reading");
+    Ok(())
+}
+
+// Options used to configure our example encryption factory
+extensions_options! {
+    struct EncryptionConfig {
+        /// Comma-separated list of columns to encrypt
+        pub encrypted_columns: String, default = "".to_owned()
+    }
+}
+
+/// Mock implementation of an `EncryptionFactory` that stores encryption keys
+/// base64 encoded in the Parquet encryption metadata.
+/// For production use, integrating with a key-management service to encrypt
+/// data encryption keys is recommended.
+#[derive(Default, Debug)]
+struct TestEncryptionFactory {}
+
+/// `EncryptionFactory` is a DataFusion trait for types that generate
+/// file encryption and decryption properties.
+impl EncryptionFactory for TestEncryptionFactory {
+    /// Generate file encryption properties to use when writing a Parquet file.
+    /// The `schema` is provided so that it may be used to dynamically configure
+    /// per-column encryption keys.
+    /// The file path is also available. We don't use the path in this example,
+    /// but other implementations may want to use this to compute an
+    /// AAD prefix for the file, or to allow use of external key material
+    /// (where key metadata is stored in a JSON file alongside Parquet files).
+    fn get_file_encryption_properties(
+        &self,
+        options: &EncryptionFactoryOptions,
+        schema: &SchemaRef,
+        _file_path: &Path,
+    ) -> Result<Option<FileEncryptionProperties>> {
+        let config: EncryptionConfig = options.to_extension_options()?;
+
+        // Generate a random encryption key for this file.
+        let mut key = vec![0u8; 16];
+        OsRng.try_fill_bytes(&mut key).unwrap();
+
+        // Generate the key metadata that allows retrieving the key when reading the file.
+        let key_metadata = wrap_key(&key);
+
+        let mut builder = FileEncryptionProperties::builder(key.to_vec())
+            .with_footer_key_metadata(key_metadata.clone());
+
+        let encrypted_columns: HashSet<&str> =
+            config.encrypted_columns.split(",").collect();
+        if !encrypted_columns.is_empty() {
+            // Set up per-column encryption.
+            for field in schema.fields().iter() {
+                if encrypted_columns.contains(field.name().as_str()) {
+                    // Here we re-use the same key for all encrypted columns,
+                    // but new keys could also be generated per column.
+                    builder = builder.with_column_key_and_metadata(
+                        field.name().as_str(),
+                        key.clone(),
+                        key_metadata.clone(),
+                    );
+                }
+            }
+        }
+
+        let encryption_properties = builder.build()?;
+
+        Ok(Some(encryption_properties))
+    }
+
+    /// Generate file decryption properties to use when reading a Parquet file.
+    /// Rather than provide the AES keys directly for decryption, we set a `KeyRetriever`
+    /// that can determine the keys using the encryption metadata.
+    fn get_file_decryption_properties(
+        &self,
+        _options: &EncryptionFactoryOptions,
+        _file_path: &Path,
+    ) -> Result<Option<FileDecryptionProperties>> {
+        let decryption_properties =
+            FileDecryptionProperties::with_key_retriever(Arc::new(TestKeyRetriever {}))
+                .build()?;
+        Ok(Some(decryption_properties))
+    }
+}
+
+/// Mock implementation of encrypting a key that simply base64 encodes the key.
+/// Note that this is not a secure way to store encryption keys,
+/// and for production use keys should be encrypted with a KMS.
+fn wrap_key(key: &[u8]) -> Vec<u8> {
+    base64::prelude::BASE64_STANDARD
+        .encode(key)
+        .as_bytes()
+        .to_vec()
+}
+
+struct TestKeyRetriever {}
+
+impl KeyRetriever for TestKeyRetriever {
+    /// Get a data encryption key using the metadata stored in the Parquet file.
+    fn retrieve_key(
+        &self,
+        key_metadata: &[u8],
+    ) -> datafusion::parquet::errors::Result<Vec<u8>> {
+        let key_metadata = std::str::from_utf8(key_metadata)?;
+        let key = base64::prelude::BASE64_STANDARD
+            .decode(key_metadata)
+            .unwrap();
+        Ok(key)
+    }
+}