diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index e6de7ab4d8..f9c2a6ff6c 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -3267,7 +3267,7 @@ pub async fn iceberg::writer::file_writer::ParquetWriter::close(self) -> iceberg pub async fn iceberg::writer::file_writer::ParquetWriter::write(&mut self, batch: &arrow_array::record_batch::RecordBatch) -> iceberg::Result<()> pub struct iceberg::writer::file_writer::ParquetWriterBuilder impl iceberg::writer::file_writer::ParquetWriterBuilder -pub fn iceberg::writer::file_writer::ParquetWriterBuilder::from_table_properties(table_props: &iceberg::spec::TableProperties, schema: iceberg::spec::SchemaRef) -> Self +pub fn iceberg::writer::file_writer::ParquetWriterBuilder::from_table_properties(table_props: &iceberg::spec::TableProperties, schema: iceberg::spec::SchemaRef) -> iceberg::Result pub fn iceberg::writer::file_writer::ParquetWriterBuilder::new(props: parquet::file::properties::WriterProperties, schema: iceberg::spec::SchemaRef) -> Self pub fn iceberg::writer::file_writer::ParquetWriterBuilder::new_with_match_mode(props: parquet::file::properties::WriterProperties, schema: iceberg::spec::SchemaRef, match_mode: iceberg::arrow::FieldMatchMode) -> Self pub fn iceberg::writer::file_writer::ParquetWriterBuilder::with_match_mode(self, match_mode: iceberg::arrow::FieldMatchMode) -> Self diff --git a/crates/iceberg/src/arrow/test_utils.rs b/crates/iceberg/src/arrow/test_utils.rs index 15ec0dc8da..b4e1beed34 100644 --- a/crates/iceberg/src/arrow/test_utils.rs +++ b/crates/iceberg/src/arrow/test_utils.rs @@ -19,7 +19,9 @@ use std::fs::File; use arrow_array::RecordBatch; use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder}; use parquet::basic::Compression; +use parquet::encryption::decrypt::FileDecryptionProperties; use parquet::encryption::encrypt::FileEncryptionProperties; use parquet::file::properties::WriterProperties; @@ -46,3 +48,26 @@ pub(crate) fn write_encrypted_parquet( writer.write(batch).expect("Writing batch"); writer.close().unwrap(); } + +/// Reads the Parquet file at `path` encrypted with `key` and `aad_prefix`, returning +/// all record batches. +pub(crate) fn read_encrypted_parquet( + path: &str, + key: &[u8], + aad_prefix: Option<&[u8]>, +) -> Vec { + let mut builder = FileDecryptionProperties::builder(key.to_vec()); + if let Some(aad) = aad_prefix { + builder = builder.with_aad_prefix(aad.to_vec()); + } + let options = + ArrowReaderOptions::new().with_file_decryption_properties(builder.build().unwrap()); + + let file = File::open(path).unwrap(); + ParquetRecordBatchReaderBuilder::try_new_with_options(file, options) + .unwrap() + .build() + .unwrap() + .map(|b| b.unwrap()) + .collect() +} diff --git a/crates/iceberg/src/encryption/key_metadata.rs b/crates/iceberg/src/encryption/key_metadata.rs index b271b32a19..b80782a0c1 100644 --- a/crates/iceberg/src/encryption/key_metadata.rs +++ b/crates/iceberg/src/encryption/key_metadata.rs @@ -20,7 +20,10 @@ use std::fmt; -use super::SecureKey; +use aes_gcm::aead::OsRng; +use aes_gcm::aead::rand_core::RngCore; + +use super::{AesKeySize, SecureKey}; use crate::{Error, ErrorKind, Result}; /// Standard key metadata for Iceberg table encryption. @@ -108,6 +111,22 @@ impl From for StandardKeyMetadata { } } +/// AAD prefix length in bytes. +const AAD_PREFIX_LENGTH: usize = 16; + +/// Generate a [`StandardKeyMetadata`] with a fresh random DEK and AAD prefix, +/// sized to `key_size`. +pub(crate) fn generate_standard_key_metadata(key_size: AesKeySize) -> StandardKeyMetadata { + let dek = SecureKey::generate(key_size); + StandardKeyMetadata::from(dek).with_aad_prefix(&generate_aad_prefix()) +} + +fn generate_aad_prefix() -> Box<[u8]> { + let mut prefix = vec![0u8; AAD_PREFIX_LENGTH]; + OsRng.fill_bytes(&mut prefix); + prefix.into_boxed_slice() +} + mod _serde { use std::io::Cursor; use std::sync::{Arc, LazyLock}; diff --git a/crates/iceberg/src/encryption/manager.rs b/crates/iceberg/src/encryption/manager.rs index e2294c2f2c..b026354a0a 100644 --- a/crates/iceberg/src/encryption/manager.rs +++ b/crates/iceberg/src/encryption/manager.rs @@ -28,8 +28,6 @@ use std::fmt; use std::sync::{Arc, RwLock}; use std::time::Duration; -use aes_gcm::aead::OsRng; -use aes_gcm::aead::rand_core::RngCore; use chrono::Utc; use moka::future::Cache; use uuid::Uuid; @@ -38,7 +36,7 @@ const MILLIS_IN_DAY: i64 = 24 * 60 * 60 * 1000; use super::crypto::{AesGcmCipher, AesKeySize, SecureKey, SensitiveBytes}; use super::io::EncryptedOutputFile; -use super::key_metadata::StandardKeyMetadata; +use super::key_metadata::{StandardKeyMetadata, generate_standard_key_metadata}; use super::kms::KeyManagementClient; use crate::io::OutputFile; use crate::spec::{EncryptedKey, FormatVersion, TableMetadataRef}; @@ -54,10 +52,6 @@ const DEFAULT_KEK_LIFESPAN_DAYS: i64 = 730; /// Default cache TTL for unwrapped KEKs. const DEFAULT_CACHE_TTL: Duration = Duration::from_secs(3600); -/// Default AAD prefix length in bytes. -/// Matches Java's `TableProperties.ENCRYPTION_AAD_LENGTH_DEFAULT`. -const AAD_PREFIX_LENGTH: usize = 16; - /// File-level encryption manager using two-layer envelope encryption. /// /// Uses an async cache for unwrapped KEK bytes to avoid repeated KMS calls. @@ -151,10 +145,7 @@ impl EncryptionManager { /// Returns an [`EncryptedOutputFile`] that transparently encrypts on /// write, along with key metadata for later decryption. pub fn encrypt(&self, raw_output: OutputFile) -> EncryptedOutputFile { - let dek = SecureKey::generate(self.key_size); - let aad_prefix = Self::generate_aad_prefix(); - let metadata = StandardKeyMetadata::from(dek).with_aad_prefix(&aad_prefix); - EncryptedOutputFile::new(raw_output, metadata) + EncryptedOutputFile::new(raw_output, generate_standard_key_metadata(self.key_size)) } /// Wrap a manifest list key metadata with a KEK for storage in table metadata. @@ -397,13 +388,6 @@ impl EncryptionManager { }) } - /// Generate a random AAD prefix for file encryption. - fn generate_aad_prefix() -> Box<[u8]> { - let mut prefix = vec![0u8; AAD_PREFIX_LENGTH]; - OsRng.fill_bytes(&mut prefix); - prefix.into_boxed_slice() - } - /// Wrap a DEK with a KEK using local AES-GCM. fn wrap_dek_with_kek( &self, diff --git a/crates/iceberg/src/writer/file_writer/parquet_writer.rs b/crates/iceberg/src/writer/file_writer/parquet_writer.rs index db9f170938..7fafa04639 100644 --- a/crates/iceberg/src/writer/file_writer/parquet_writer.rs +++ b/crates/iceberg/src/writer/file_writer/parquet_writer.rs @@ -27,6 +27,7 @@ use itertools::Itertools; use parquet::arrow::AsyncArrowWriter; use parquet::arrow::async_reader::AsyncFileReader; use parquet::arrow::async_writer::AsyncFileWriter as ArrowAsyncFileWriter; +use parquet::encryption::encrypt::FileEncryptionProperties; use parquet::file::metadata::ParquetMetaData; use parquet::file::properties::{CdcOptions, WriterProperties}; use parquet::file::statistics::Statistics; @@ -36,6 +37,8 @@ use crate::arrow::{ ArrowFileReader, DEFAULT_MAP_FIELD_NAME, FieldMatchMode, NanValueCountVisitor, get_parquet_stat_max_as_datum, get_parquet_stat_min_as_datum, }; +use crate::encryption::key_metadata::generate_standard_key_metadata; +use crate::encryption::{AesKeySize, StandardKeyMetadata}; use crate::io::{FileIO, FileWrite, OutputFile}; use crate::spec::{ DataContentType, DataFileBuilder, DataFileFormat, Datum, ListType, Literal, MapType, @@ -52,6 +55,7 @@ pub struct ParquetWriterBuilder { props: WriterProperties, schema: SchemaRef, match_mode: FieldMatchMode, + data_encryption_key_size: Option, } impl ParquetWriterBuilder { @@ -75,6 +79,7 @@ impl ParquetWriterBuilder { props, schema, match_mode, + data_encryption_key_size: None, } } @@ -85,19 +90,34 @@ impl ParquetWriterBuilder { /// Currently translates the content-defined-chunking keys /// (`write.parquet.content-defined-chunking.*`); other keys fall back to /// parquet-rs defaults. - pub fn from_table_properties(table_props: &TableProperties, schema: SchemaRef) -> Self { + /// + /// When `encryption.key-id` is set, records the DEK length. + pub fn from_table_properties(table_props: &TableProperties, schema: SchemaRef) -> Result { let cdc = table_props.cdc_enabled.then_some(CdcOptions { min_chunk_size: table_props.cdc_min_chunk_size, max_chunk_size: table_props.cdc_max_chunk_size, norm_level: table_props.cdc_norm_level, }); + // TODO: translate the remaining write.parquet.* keys (e.g. compression-codec, // row-group-size-bytes, page-size-bytes). // This constructor is intended to be the single place that maps them. let props = WriterProperties::builder() .set_content_defined_chunking(cdc) .build(); - Self::new_with_match_mode(props, schema, FieldMatchMode::Id) + + let data_encryption_key_size = table_props + .encryption_key_id + .is_some() + .then(|| AesKeySize::from_key_length(table_props.encryption_data_key_length)) + .transpose()?; + + Ok(Self { + props, + schema, + match_mode: FieldMatchMode::Id, + data_encryption_key_size, + }) } /// Set the field match mode used to map Arrow fields to Iceberg fields. @@ -114,13 +134,18 @@ impl FileWriterBuilder for ParquetWriterBuilder { type R = ParquetWriter; async fn build(&self, output_file: OutputFile) -> Result { + let key_metadata = self + .data_encryption_key_size + .map(generate_standard_key_metadata); + let writer_properties = resolve_writer_properties(&self.props, key_metadata.as_ref())?; Ok(ParquetWriter { schema: self.schema.clone(), inner_writer: None, - writer_properties: self.props.clone(), + writer_properties, current_row_num: 0, output_file, nan_value_count_visitor: NanValueCountVisitor::new_with_match_mode(self.match_mode), + key_metadata, }) } } @@ -258,6 +283,7 @@ pub struct ParquetWriter { writer_properties: WriterProperties, current_row_num: usize, nan_value_count_visitor: NanValueCountVisitor, + key_metadata: Option, } /// Used to aggregate min and max value of each column. @@ -349,7 +375,7 @@ impl MinMaxColAggregator { } impl ParquetWriter { - /// Converts parquet files to data files + /// Converts already-written parquet files into [`DataFile`]s. #[allow(dead_code)] pub(crate) async fn parquet_files_to_data_files( file_io: &FileIO, @@ -379,6 +405,7 @@ impl ParquetWriter { file_path, // TODO: Implement nan_value_counts here HashMap::new(), + None, )?; builder.partition_spec_id(table_metadata.default_partition_spec_id()); let data_file = builder.build().unwrap(); @@ -395,6 +422,7 @@ impl ParquetWriter { written_size: usize, file_path: String, nan_value_counts: HashMap, + key_metadata: Option, ) -> Result { let index_by_parquet_path = { let mut visitor = IndexByParquetPathName::new(); @@ -438,6 +466,11 @@ impl ParquetWriter { ) }; + let key_metadata = match key_metadata { + Some(m) => Some(m.encode()?.into_vec()), + None => None, + }; + let mut builder = DataFileBuilder::default(); builder .content(DataContentType::Data) @@ -454,6 +487,7 @@ impl ParquetWriter { // - We can ignore implementing distinct_counts due to this: https://lists.apache.org/thread/j52tsojv0x4bopxyzsp7m7bqt23n5fnd .lower_bounds(lower_bounds) .upper_bounds(upper_bounds) + .key_metadata(key_metadata) .split_offsets(Some( metadata .row_groups() @@ -514,6 +548,41 @@ impl ParquetWriter { } } +fn resolve_writer_properties( + writer_properties: &WriterProperties, + key_metadata: Option<&StandardKeyMetadata>, +) -> Result { + let Some(key_metadata) = key_metadata else { + return Ok(writer_properties.clone()); + }; + + if writer_properties.file_encryption_properties().is_some() { + return Err(Error::new( + ErrorKind::Unexpected, + "Parquet writer properties already have file encryption properties set", + )); + } + + let mut builder = + FileEncryptionProperties::builder(key_metadata.encryption_key().as_bytes().to_vec()); + if let Some(aad) = key_metadata.aad_prefix() { + builder = builder.with_aad_prefix(aad.to_vec()); + } + let file_encryption_properties = builder.build().map_err(|e| { + Error::new( + ErrorKind::Unexpected, + "Failed to build parquet file encryption properties", + ) + .with_source(e) + })?; + + Ok(writer_properties + .clone() + .into_builder() + .with_file_encryption_properties(file_encryption_properties) + .build()) +} + impl FileWriter for ParquetWriter { async fn write(&mut self, batch: &arrow_array::RecordBatch) -> Result<()> { // Skip empty batch @@ -588,6 +657,7 @@ impl FileWriter for ParquetWriter { written_size, self.output_file.location().to_string(), self.nan_value_count_visitor.nan_value_counts, + self.key_metadata, )?]) } } @@ -662,14 +732,18 @@ mod tests { }; use arrow_schema::{DataType, Field, Fields, SchemaRef as ArrowSchemaRef}; use arrow_select::concat::concat_batches; + use futures::TryStreamExt; use parquet::arrow::PARQUET_FIELD_ID_META_KEY; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use parquet::file::statistics::ValueStatistics; use tempfile::TempDir; use uuid::Uuid; use super::*; - use crate::arrow::schema_to_arrow_schema; + use crate::Runtime; + use crate::arrow::{ArrowReaderBuilder, schema_to_arrow_schema}; use crate::io::FileIO; + use crate::scan::{FileScanTask, FileScanTaskStream}; use crate::spec::decimal_utils::{decimal_mantissa, decimal_new, decimal_scale}; use crate::spec::{PrimitiveLiteral, Struct, *}; use crate::writer::file_writer::location_generator::{ @@ -918,6 +992,163 @@ mod tests { Ok(()) } + /// We use a helper for the read path to test write path only here. + #[tokio::test] + async fn test_parquet_writer_encrypted_write_path() -> Result<()> { + let temp_dir = TempDir::new().unwrap(); + let file_io = FileIO::new_with_fs(); + let location_gen = DefaultLocationGenerator::with_data_location( + temp_dir.path().to_str().unwrap().to_string(), + ); + let file_name_gen = + DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet); + + let arrow_schema = Arc::new(arrow_schema::Schema::new(vec![ + Field::new("col", DataType::Int64, true).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "0".to_string(), + )])), + ])); + let col = Arc::new(Int64Array::from_iter_values(0..1024)) as ArrayRef; + let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![col]).unwrap(); + + let output_file = file_io.new_output( + location_gen.generate_location(None, &file_name_gen.generate_file_name()), + )?; + + // `encryption.key-id` turns on table-managed encryption in the writer. + let table_properties = table_props(HashMap::from([( + TableProperties::PROPERTY_ENCRYPTION_KEY_ID.to_string(), + "test-key".to_string(), + )])); + let mut parquet_writer = ParquetWriterBuilder::from_table_properties( + &table_properties, + Arc::new(to_write.schema().as_ref().try_into().unwrap()), + )? + .build(output_file) + .await?; + parquet_writer.write(&to_write).await?; + let data_file = parquet_writer + .close() + .await? + .into_iter() + .next() + .unwrap() + .content(DataContentType::Data) + .partition(Struct::empty()) + .partition_spec_id(0) + .build() + .unwrap(); + + // The DEK must be recorded on the data file + assert!( + data_file.key_metadata().is_some(), + "encrypted data file must carry key metadata" + ); + + // A plain-text read with no decryption must fail. + let raw = file_io + .new_input(data_file.file_path.clone())? + .read() + .await?; + assert!( + ParquetRecordBatchReaderBuilder::try_new(raw).is_err(), + "an encrypted parquet file must not be readable without decryption" + ); + + // Recovering the DEK + AAD prefix from key_metadata reads the data back intact. + let key_metadata = StandardKeyMetadata::decode(data_file.key_metadata().unwrap()).unwrap(); + let batches = crate::arrow::test_utils::read_encrypted_parquet( + &data_file.file_path, + key_metadata.encryption_key().as_bytes(), + key_metadata.aad_prefix(), + ); + let res = concat_batches(&arrow_schema, &batches).unwrap(); + assert_eq!(to_write, res); + + Ok(()) + } + + /// Full roundtrip test: read and write + #[tokio::test] + async fn test_parquet_writer_encrypted_roundtrip() -> Result<()> { + let temp_dir = TempDir::new().unwrap(); + let file_io = FileIO::new_with_fs(); + let location_gen = DefaultLocationGenerator::with_data_location( + temp_dir.path().to_str().unwrap().to_string(), + ); + let file_name_gen = + DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet); + + let arrow_schema = Arc::new(arrow_schema::Schema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + ])); + let iceberg_schema: SchemaRef = Arc::new(arrow_schema.as_ref().try_into().unwrap()); + let batch = RecordBatch::try_new(arrow_schema, vec![Arc::new(Int32Array::from(vec![ + 10, 20, 30, + ])) as ArrayRef]) + .unwrap(); + + let output_file = file_io.new_output( + location_gen.generate_location(None, &file_name_gen.generate_file_name()), + )?; + let table_properties = table_props(HashMap::from([( + TableProperties::PROPERTY_ENCRYPTION_KEY_ID.to_string(), + "test-key".to_string(), + )])); + let mut parquet_writer = + ParquetWriterBuilder::from_table_properties(&table_properties, iceberg_schema.clone())? + .build(output_file) + .await?; + parquet_writer.write(&batch).await?; + let data_file = parquet_writer + .close() + .await? + .into_iter() + .next() + .unwrap() + .content(DataContentType::Data) + .partition(Struct::empty()) + .partition_spec_id(0) + .build() + .unwrap(); + + // Read back through the iceberg reader using the file's own key metadata. + let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build(); + let task = FileScanTask::builder() + .with_file_size_in_bytes(data_file.file_size_in_bytes()) + .with_start(0) + .with_length(0) + .with_data_file_path(data_file.file_path.clone()) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(iceberg_schema) + .with_project_field_ids(vec![1]) + .with_case_sensitive(false) + .with_key_metadata(data_file.key_metadata().map(Box::from)) + .build(); + let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream; + let batches: Vec = reader + .read(tasks) + .unwrap() + .stream() + .try_collect() + .await + .unwrap(); + + assert_eq!(batches.len(), 1); + let ids = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ids.values(), &[10, 20, 30]); + + Ok(()) + } + #[tokio::test] async fn test_parquet_writer_with_complex_schema() -> Result<()> { let temp_dir = TempDir::new().unwrap(); @@ -2366,7 +2597,7 @@ mod tests { #[test] fn test_from_table_properties_no_cdc_by_default() { let tp = table_props(HashMap::new()); - let builder = ParquetWriterBuilder::from_table_properties(&tp, cdc_test_schema()); + let builder = ParquetWriterBuilder::from_table_properties(&tp, cdc_test_schema()).unwrap(); assert!(builder.props.content_defined_chunking().is_none()); } @@ -2404,6 +2635,7 @@ mod tests { .new_output(format!("{}/cdc.parquet", tmp.path().to_str().unwrap())) .unwrap(); let writer = ParquetWriterBuilder::from_table_properties(&tp, cdc_test_schema()) + .unwrap() .build(output) .await .unwrap(); diff --git a/crates/iceberg/src/writer/file_writer/rolling_writer.rs b/crates/iceberg/src/writer/file_writer/rolling_writer.rs index b0b2d2f191..ba5469a31e 100644 --- a/crates/iceberg/src/writer/file_writer/rolling_writer.rs +++ b/crates/iceberg/src/writer/file_writer/rolling_writer.rs @@ -257,7 +257,7 @@ impl CurrentFi #[cfg(test)] mod tests { - use std::collections::HashMap; + use std::collections::{HashMap, HashSet}; use std::sync::Arc; use arrow_array::{ArrayRef, Int32Array, StringArray}; @@ -268,6 +268,8 @@ mod tests { use tempfile::TempDir; use super::*; + use crate::arrow::test_utils::read_encrypted_parquet; + use crate::encryption::StandardKeyMetadata; use crate::io::FileIO; use crate::spec::{DataFileFormat, NestedField, PrimitiveType, Schema, Type}; use crate::writer::base_writer::data_file_writer::DataFileWriterBuilder; @@ -439,4 +441,111 @@ mod tests { Ok(()) } + + /// Every rolled file is independently encrypted + #[tokio::test] + async fn test_rolling_writer_encrypted() -> Result<()> { + let temp_dir = TempDir::new()?; + let file_io = FileIO::new_with_fs(); + let location_gen = DefaultLocationGenerator::with_data_location( + temp_dir.path().to_str().unwrap().to_string(), + ); + let file_name_gen = + DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet); + + let schema = make_test_schema()?; + + let table_properties = TableProperties::try_from(&HashMap::from([( + TableProperties::PROPERTY_ENCRYPTION_KEY_ID.to_string(), + "test-key".to_string(), + )]))?; + let parquet_writer_builder = + ParquetWriterBuilder::from_table_properties(&table_properties, Arc::new(schema))?; + + // Set a very small target size to trigger rolling + let rolling_writer_builder = RollingFileWriterBuilder::new( + parquet_writer_builder, + 1024, + file_io.clone(), + location_gen, + file_name_gen, + ); + + let data_file_writer_builder = DataFileWriterBuilder::new(rolling_writer_builder); + + // Create writer + let mut writer = data_file_writer_builder.build(None).await?; + + // Create test data + let arrow_schema = make_test_arrow_schema(); + let arrow_schema_ref = Arc::new(arrow_schema.clone()); + + let names = vec![ + "Alice", "Bob", "Charlie", "Dave", "Eve", "Frank", "Grace", "Heidi", "Ivan", "Judy", + "Kelly", "Larry", "Mallory", "Shawn", + ]; + + let mut rng = rand::rng(); + let batch_num = 10; + let batch_rows = 100; + let expected_rows = batch_num * batch_rows; + + for i in 0..batch_num { + let int_values: Vec = (0..batch_rows).map(|row| i * batch_rows + row).collect(); + let str_values: Vec<&str> = (0..batch_rows) + .map(|_| *names.iter().choose(&mut rng).unwrap()) + .collect(); + + let int_array = Arc::new(Int32Array::from(int_values)) as ArrayRef; + let str_array = Arc::new(StringArray::from(str_values)) as ArrayRef; + + let batch = + RecordBatch::try_new(Arc::clone(&arrow_schema_ref), vec![int_array, str_array]) + .expect("Failed to create RecordBatch"); + + writer.write(batch).await?; + } + + let data_files = writer.close().await?; + + assert!( + data_files.len() > 4, + "Expected at least 4 data files to be created, but got {}", + data_files.len() + ); + + // Verify total record count across all files + let total_records: u64 = data_files.iter().map(|file| file.record_count).sum(); + assert_eq!( + total_records, expected_rows as u64, + "Expected {expected_rows} total records across all files" + ); + + // Every rolled file carries key metadata, and each file's DEK is distinct. + let mut distinct_keys = HashSet::new(); + for data_file in &data_files { + let key_metadata = StandardKeyMetadata::decode( + data_file + .key_metadata() + .expect("each rolled file must carry key metadata"), + )?; + assert!( + distinct_keys.insert(key_metadata.encryption_key().as_bytes().to_vec()), + "each rolled file must use a distinct DEK" + ); + } + + // One file decrypts and reads back with its own key metadata. + let data_file = &data_files[0]; + let key_metadata = StandardKeyMetadata::decode(data_file.key_metadata().unwrap())?; + let batches = read_encrypted_parquet( + &data_file.file_path, + key_metadata.encryption_key().as_bytes(), + key_metadata.aad_prefix(), + ); + let read_rows: u64 = batches.iter().map(|b| b.num_rows() as u64).sum(); + assert_eq!(read_rows, data_file.record_count); + + Ok(()) + } } diff --git a/crates/integrations/datafusion/src/physical_plan/write.rs b/crates/integrations/datafusion/src/physical_plan/write.rs index a7d771ec1b..3e5026dab0 100644 --- a/crates/integrations/datafusion/src/physical_plan/write.rs +++ b/crates/integrations/datafusion/src/physical_plan/write.rs @@ -226,6 +226,7 @@ impl ExecutionPlan for IcebergWriteExec { &table_props, self.table.metadata().current_schema().clone(), ) + .map_err(to_datafusion_error)? .with_match_mode(FieldMatchMode::Name); let target_file_size = table_props.write_target_file_size_bytes; @@ -317,12 +318,16 @@ mod tests { use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; use futures::{StreamExt, stream}; + use iceberg::encryption::StandardKeyMetadata; use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; use iceberg::spec::{ - DataFileFormat, NestedField, PrimitiveType, Schema, Type, deserialize_data_file_from_json, + DataFileFormat, NestedField, PrimitiveType, Schema, TableProperties, Type, + deserialize_data_file_from_json, }; use iceberg::{Catalog, CatalogBuilder, MemoryCatalog, NamespaceIdent, Result, TableCreation}; use parquet::arrow::PARQUET_FIELD_ID_META_KEY; + use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder}; + use parquet::encryption::decrypt::FileDecryptionProperties; use tempfile::TempDir; use super::*; @@ -601,4 +606,100 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn test_iceberg_write_exec_encrypted() -> Result<()> { + let iceberg_catalog = get_iceberg_catalog().await; + let namespace = NamespaceIdent::new("test_namespace_enc".to_string()); + iceberg_catalog + .create_namespace(&namespace, HashMap::new()) + .await?; + + // Create a table with encryption enabled via `encryption.key-id`. + let creation = TableCreation::builder() + .location(temp_path()) + .name("test_table_enc".to_string()) + .properties(HashMap::from([( + TableProperties::PROPERTY_ENCRYPTION_KEY_ID.to_string(), + "test-key".to_string(), + )])) + .schema(get_test_schema()?) + .build(); + let table = iceberg_catalog.create_table(&namespace, creation).await?; + + // Input data fed to the planner. + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + Field::new("name", DataType::Utf8, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "2".to_string(), + )])), + ])); + let id_array = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + let name_array = Arc::new(StringArray::from(vec!["Alice", "Bob", "Charlie"])) as ArrayRef; + let batch = RecordBatch::try_new(arrow_schema.clone(), vec![id_array, name_array]) + .map_err(|e| Error::new(ErrorKind::Unexpected, format!("record batch: {e}")))?; + + let input_plan = Arc::new(MockExecutionPlan::new(arrow_schema.clone(), vec![batch])); + let write_exec = IcebergWriteExec::new(table.clone(), input_plan, arrow_schema); + + // Execute the planner and collect the returned (serialized) data files. + let task_ctx = Arc::new(TaskContext::default()); + let mut stream = write_exec + .execute(0, task_ctx) + .map_err(|e| Error::new(ErrorKind::Unexpected, format!("execute: {e}")))?; + let mut results = vec![]; + while let Some(batch) = stream.next().await { + results + .push(batch.map_err(|e| Error::new(ErrorKind::Unexpected, format!("batch: {e}")))?); + } + assert_eq!(results.len(), 1, "expected one result batch"); + + let data_file_json = results[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("expected StringArray") + .value(0); + let data_file = deserialize_data_file_from_json( + data_file_json, + table.metadata().default_partition_spec_id(), + table.metadata().default_partition_type(), + table.metadata().current_schema(), + )?; + + // The planner encrypted the file. + let key_metadata_bytes = data_file + .key_metadata() + .expect("planner must record key metadata on an encrypted data file"); + let key_metadata = StandardKeyMetadata::decode(key_metadata_bytes)?; + let mut builder = + FileDecryptionProperties::builder(key_metadata.encryption_key().as_bytes().to_vec()); + if let Some(aad) = key_metadata.aad_prefix() { + builder = builder.with_aad_prefix(aad.to_vec()); + } + let options = + ArrowReaderOptions::new().with_file_decryption_properties(builder.build().unwrap()); + let bytes = table + .file_io() + .new_input(data_file.file_path()) + .unwrap() + .read() + .await + .unwrap(); + let batches: Vec = + ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, options) + .unwrap() + .build() + .unwrap() + .map(|b| b.unwrap()) + .collect(); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 3, "decrypted file must contain the 3 written rows"); + + Ok(()) + } }