diff --git a/crates/iceberg/src/transaction/snapshot.rs b/crates/iceberg/src/transaction/snapshot.rs index d200c5ba9c..ecbc3f1fae 100644 --- a/crates/iceberg/src/transaction/snapshot.rs +++ b/crates/iceberg/src/transaction/snapshot.rs @@ -249,16 +249,25 @@ impl<'a> SnapshotProducer<'a> { DataFileFormat::Avro ); let output_file = self.table.file_io().new_output(new_manifest_path)?; - let builder = ManifestWriterBuilder::new( - output_file, - Some(self.snapshot_id), - self.table.metadata().current_schema().clone(), - self.table - .metadata() - .default_partition_spec() - .as_ref() - .clone(), - ); + let partition_spec = self + .table + .metadata() + .default_partition_spec() + .as_ref() + .clone(); + let schema = self.table.metadata().current_schema().clone(); + + let builder = if let Some(em) = self.table.encryption_manager() { + ManifestWriterBuilder::new_from_encrypted( + em.encrypt(output_file), + Some(self.snapshot_id), + schema, + partition_spec, + )? + } else { + ManifestWriterBuilder::new(output_file, Some(self.snapshot_id), schema, partition_spec) + }; + match self.table.metadata().format_version() { FormatVersion::V1 => Ok(builder.build_v1()), FormatVersion::V2 => match content { @@ -530,3 +539,179 @@ impl<'a> SnapshotProducer<'a> { Ok(ActionCommit::new(updates, requirements)) } } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::fs::File; + use std::io::BufReader; + use std::sync::Arc; + + use uuid::Uuid; + + use super::SnapshotProducer; + use crate::TableIdent; + use crate::encryption::kms::{KeyManagementClient, MemoryKeyManagementClient}; + use crate::encryption::{SensitiveBytes, StandardKeyMetadata}; + use crate::io::FileIO; + use crate::spec::{ + DataContentType, DataFileBuilder, DataFileFormat, Manifest, ManifestContentType, + ManifestStatus, Struct, TableMetadata, + }; + use crate::table::Table; + use crate::test_utils::test_runtime; + + /// Build a table backed by the V3 encryption fixture and an in-memory KMS, + /// so it has an [`EncryptionManager`](crate::encryption::EncryptionManager). + fn make_encrypted_table() -> Table { + let file = File::open(format!( + "{}/testdata/table_metadata/{}", + env!("CARGO_MANIFEST_DIR"), + "TableMetadataV3ValidEncryption.json" + )) + .unwrap(); + let reader = BufReader::new(file); + let metadata = serde_json::from_reader::<_, TableMetadata>(reader).unwrap(); + + let kms: Arc = { + let k = MemoryKeyManagementClient::new(); + k.add_master_key_bytes( + "master-1", + SensitiveBytes::new([ + 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, + 0x0d, 0x0e, 0x0f, + ]), + ) + .unwrap(); + Arc::new(k) + }; + + Table::builder() + .metadata(metadata) + .metadata_location("memory:///table/metadata/v1.json") + .identifier(TableIdent::from_strs(["ns1", "test1"]).unwrap()) + .file_io(FileIO::new_with_memory()) + .kms_client(kms) + .runtime(test_runtime()) + .build() + .unwrap() + } + + fn data_file() -> crate::spec::DataFile { + DataFileBuilder::default() + .content(DataContentType::Data) + .file_path("memory:///table/data/00000.parquet".to_string()) + .file_format(DataFileFormat::Parquet) + .partition(Struct::empty()) + .record_count(100) + .file_size_in_bytes(4096) + .partition_spec_id(0) + .build() + .unwrap() + } + + #[tokio::test] + async fn write_encrypted_manifest_roundtrips() { + let table = make_encrypted_table(); + assert!( + table.encryption_manager().is_some(), + "fixture table should have an EncryptionManager" + ); + + let mut producer = + SnapshotProducer::new(&table, Uuid::new_v4(), HashMap::new(), vec![data_file()]); + + let mut writer = producer + .new_manifest_writer(ManifestContentType::Data) + .unwrap(); + writer + .add_entry( + crate::spec::ManifestEntry::builder() + .status(ManifestStatus::Added) + .data_file(data_file()) + .build(), + ) + .unwrap(); + let manifest_file = writer.write_manifest_file().await.unwrap(); + + // The manifest list entry must carry decodable key metadata. + let key_metadata_bytes = manifest_file + .key_metadata + .as_ref() + .expect("encrypted manifest must record key metadata"); + assert!(!key_metadata_bytes.is_empty()); + StandardKeyMetadata::decode(key_metadata_bytes) + .expect("recorded key metadata must decode as StandardKeyMetadata"); + + // The bytes on disk must actually be encrypted: a plaintext Avro parse + // must fail. (Guards against silently taking the unencrypted branch.) + let raw = table + .file_io() + .new_input(&manifest_file.manifest_path) + .unwrap() + .read() + .await + .unwrap(); + assert!( + Manifest::parse_avro(&raw).is_err(), + "manifest bytes should not be readable as plaintext Avro" + ); + + // load_manifest self-decrypts using the recorded key metadata and must + // recover the entry we wrote. + let manifest = manifest_file.load_manifest(table.file_io()).await.unwrap(); + assert_eq!(manifest.entries().len(), 1); + assert_eq!( + manifest.entries()[0].data_file().file_path(), + "memory:///table/data/00000.parquet" + ); + } + + #[tokio::test] + async fn write_plaintext_manifest_when_no_encryption() { + let table = crate::transaction::tests::make_v2_table(); + assert!(table.encryption_manager().is_none()); + + let data_file = DataFileBuilder::default() + .content(DataContentType::Data) + .file_path("memory:///table/data/00000.parquet".to_string()) + .file_format(DataFileFormat::Parquet) + .partition(Struct::from_iter([Some(crate::spec::Literal::long(0))])) + .record_count(100) + .file_size_in_bytes(4096) + .partition_spec_id(0) + .build() + .unwrap(); + + let mut producer = SnapshotProducer::new(&table, Uuid::new_v4(), HashMap::new(), vec![ + data_file.clone(), + ]); + + let mut writer = producer + .new_manifest_writer(ManifestContentType::Data) + .unwrap(); + writer + .add_entry( + crate::spec::ManifestEntry::builder() + .status(ManifestStatus::Added) + .data_file(data_file) + .build(), + ) + .unwrap(); + let manifest_file = writer.write_manifest_file().await.unwrap(); + + assert!( + manifest_file.key_metadata.is_none(), + "plaintext manifest must not record key metadata" + ); + + let raw = table + .file_io() + .new_input(&manifest_file.manifest_path) + .unwrap() + .read() + .await + .unwrap(); + Manifest::parse_avro(&raw).expect("plaintext manifest bytes must parse as Avro"); + } +}