iceberg/spec/manifest_list/
writer.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::collections::HashMap;
19
20use apache_avro::Writer;
21use bytes::Bytes;
22
23use super::_const_schema::{
24    MANIFEST_LIST_AVRO_SCHEMA_V1, MANIFEST_LIST_AVRO_SCHEMA_V2, MANIFEST_LIST_AVRO_SCHEMA_V3,
25};
26use super::_serde::{ManifestFileV1, ManifestFileV2, ManifestFileV3};
27use super::{FormatVersion, ManifestContentType, ManifestFile, UNASSIGNED_SEQUENCE_NUMBER};
28use crate::error::Result;
29use crate::io::FileWrite;
30use crate::{Error, ErrorKind};
31
32/// A manifest list writer.
33pub struct ManifestListWriter {
34    format_version: FormatVersion,
35    writer: Box<dyn FileWrite>,
36    avro_writer: Writer<'static, Vec<u8>>,
37    sequence_number: i64,
38    snapshot_id: i64,
39    next_row_id: Option<u64>,
40}
41
42impl std::fmt::Debug for ManifestListWriter {
43    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44        f.debug_struct("ManifestListWriter")
45            .field("format_version", &self.format_version)
46            .field("avro_writer", &self.avro_writer.schema())
47            .finish_non_exhaustive()
48    }
49}
50
51impl ManifestListWriter {
52    /// Get the next row ID that will be assigned to the next data manifest added.
53    pub fn next_row_id(&self) -> Option<u64> {
54        self.next_row_id
55    }
56
57    /// Construct a v1 [`ManifestListWriter`] that writes to a provided [`FileWrite`].
58    pub fn v1(
59        writer: Box<dyn FileWrite>,
60        snapshot_id: i64,
61        parent_snapshot_id: Option<i64>,
62    ) -> Self {
63        let mut metadata = HashMap::from_iter([
64            ("snapshot-id".to_string(), snapshot_id.to_string()),
65            ("format-version".to_string(), "1".to_string()),
66        ]);
67        if let Some(parent_snapshot_id) = parent_snapshot_id {
68            metadata.insert(
69                "parent-snapshot-id".to_string(),
70                parent_snapshot_id.to_string(),
71            );
72        }
73        Self::new(FormatVersion::V1, writer, metadata, 0, snapshot_id, None)
74    }
75
76    /// Construct a v2 [`ManifestListWriter`] that writes to a provided [`FileWrite`].
77    pub fn v2(
78        writer: Box<dyn FileWrite>,
79        snapshot_id: i64,
80        parent_snapshot_id: Option<i64>,
81        sequence_number: i64,
82    ) -> Self {
83        let mut metadata = HashMap::from_iter([
84            ("snapshot-id".to_string(), snapshot_id.to_string()),
85            ("sequence-number".to_string(), sequence_number.to_string()),
86            ("format-version".to_string(), "2".to_string()),
87        ]);
88        metadata.insert(
89            "parent-snapshot-id".to_string(),
90            parent_snapshot_id
91                .map(|v| v.to_string())
92                .unwrap_or("null".to_string()),
93        );
94        Self::new(
95            FormatVersion::V2,
96            writer,
97            metadata,
98            sequence_number,
99            snapshot_id,
100            None,
101        )
102    }
103
104    /// Construct a v3 [`ManifestListWriter`] that writes to a provided [`FileWrite`].
105    pub fn v3(
106        writer: Box<dyn FileWrite>,
107        snapshot_id: i64,
108        parent_snapshot_id: Option<i64>,
109        sequence_number: i64,
110        first_row_id: Option<u64>, // Always None for delete manifests
111    ) -> Self {
112        let mut metadata = HashMap::from_iter([
113            ("snapshot-id".to_string(), snapshot_id.to_string()),
114            ("sequence-number".to_string(), sequence_number.to_string()),
115            ("format-version".to_string(), "3".to_string()),
116        ]);
117        metadata.insert(
118            "parent-snapshot-id".to_string(),
119            parent_snapshot_id
120                .map(|v| v.to_string())
121                .unwrap_or("null".to_string()),
122        );
123        metadata.insert(
124            "first-row-id".to_string(),
125            first_row_id
126                .map(|v| v.to_string())
127                .unwrap_or("null".to_string()),
128        );
129        Self::new(
130            FormatVersion::V3,
131            writer,
132            metadata,
133            sequence_number,
134            snapshot_id,
135            first_row_id,
136        )
137    }
138
139    fn new(
140        format_version: FormatVersion,
141        writer: Box<dyn FileWrite>,
142        metadata: HashMap<String, String>,
143        sequence_number: i64,
144        snapshot_id: i64,
145        first_row_id: Option<u64>,
146    ) -> Self {
147        let avro_schema = match format_version {
148            FormatVersion::V1 => &MANIFEST_LIST_AVRO_SCHEMA_V1,
149            FormatVersion::V2 => &MANIFEST_LIST_AVRO_SCHEMA_V2,
150            FormatVersion::V3 => &MANIFEST_LIST_AVRO_SCHEMA_V3,
151        };
152        let mut avro_writer = Writer::new(avro_schema, Vec::new());
153        for (key, value) in metadata {
154            avro_writer
155                .add_user_metadata(key, value)
156                .expect("Avro metadata should be added to the writer before the first record.");
157        }
158        Self {
159            format_version,
160            writer,
161            avro_writer,
162            sequence_number,
163            snapshot_id,
164            next_row_id: first_row_id,
165        }
166    }
167
168    /// Append manifests to be written.
169    ///
170    /// If V3 Manifests are added and the `first_row_id` of any data manifest is unassigned,
171    /// it will be assigned based on the `next_row_id` of the writer, and the `next_row_id` of the writer will be updated accordingly.
172    /// If `first_row_id` is already assigned, it will be validated against the `next_row_id` of the writer.
173    pub fn add_manifests(&mut self, manifests: impl Iterator<Item = ManifestFile>) -> Result<()> {
174        match self.format_version {
175            FormatVersion::V1 => {
176                for manifest in manifests {
177                    let manifests: ManifestFileV1 = manifest.try_into()?;
178                    self.avro_writer.append_ser(manifests)?;
179                }
180            }
181            FormatVersion::V2 | FormatVersion::V3 => {
182                for mut manifest in manifests {
183                    self.assign_sequence_numbers(&mut manifest)?;
184
185                    if self.format_version == FormatVersion::V2 {
186                        let manifest_entry: ManifestFileV2 = manifest.try_into()?;
187                        self.avro_writer.append_ser(manifest_entry)?;
188                    } else if self.format_version == FormatVersion::V3 {
189                        self.assign_first_row_id(&mut manifest)?;
190                        let manifest_entry: ManifestFileV3 = manifest.try_into()?;
191                        self.avro_writer.append_ser(manifest_entry)?;
192                    }
193                }
194            }
195        }
196        Ok(())
197    }
198
199    /// Write the manifest list to the output file.
200    pub async fn close(mut self) -> Result<()> {
201        let data = self.avro_writer.into_inner()?;
202        self.writer.write(Bytes::from(data)).await?;
203        self.writer.close().await?;
204        Ok(())
205    }
206
207    /// Assign sequence numbers to manifest if they are unassigned
208    fn assign_sequence_numbers(&self, manifest: &mut ManifestFile) -> Result<()> {
209        if manifest.sequence_number == UNASSIGNED_SEQUENCE_NUMBER {
210            if manifest.added_snapshot_id != self.snapshot_id {
211                return Err(Error::new(
212                    ErrorKind::DataInvalid,
213                    format!(
214                        "Found unassigned sequence number for a manifest from snapshot {}.",
215                        manifest.added_snapshot_id
216                    ),
217                ));
218            }
219            manifest.sequence_number = self.sequence_number;
220        }
221
222        if manifest.min_sequence_number == UNASSIGNED_SEQUENCE_NUMBER {
223            if manifest.added_snapshot_id != self.snapshot_id {
224                return Err(Error::new(
225                    ErrorKind::DataInvalid,
226                    format!(
227                        "Found unassigned sequence number for a manifest from snapshot {}.",
228                        manifest.added_snapshot_id
229                    ),
230                ));
231            }
232            manifest.min_sequence_number = self.sequence_number;
233        }
234
235        Ok(())
236    }
237
238    /// Returns number of newly assigned first-row-ids, if any.
239    fn assign_first_row_id(&mut self, manifest: &mut ManifestFile) -> Result<()> {
240        match manifest.content {
241            ManifestContentType::Data => {
242                match (self.next_row_id, manifest.first_row_id) {
243                    (Some(_), Some(_)) => {
244                        // Case: Manifest with already assigned first row ID.
245                        // No need to increase next_row_id, as this manifest is already assigned.
246                    }
247                    (None, Some(manifest_first_row_id)) => {
248                        // Case: Assigned first row ID for data manifest, but the writer does not have a next-row-id assigned.
249                        return Err(Error::new(
250                            ErrorKind::Unexpected,
251                            format!(
252                                "Found invalid first-row-id assignment for Manifest {}. Writer does not have a next-row-id assigned, but the manifest has first-row-id assigned to {}.",
253                                manifest.manifest_path, manifest_first_row_id,
254                            ),
255                        ));
256                    }
257                    (Some(writer_next_row_id), None) => {
258                        // Case: Unassigned first row ID for data manifest. This is either a new
259                        // manifest, or a manifest from a pre-v3 snapshot. We need to assign one.
260                        let (existing_rows_count, added_rows_count) =
261                            require_row_counts_in_manifest(manifest)?;
262                        manifest.first_row_id = Some(writer_next_row_id);
263
264                        self.next_row_id = writer_next_row_id
265                        .checked_add(existing_rows_count)
266                        .and_then(|sum| sum.checked_add(added_rows_count))
267                        .ok_or_else(|| {
268                            Error::new(
269                                ErrorKind::DataInvalid,
270                                format!(
271                                    "Row ID overflow when computing next row ID for Manifest {}. Next Row ID: {writer_next_row_id}, Existing Rows Count: {existing_rows_count}, Added Rows Count: {added_rows_count}",
272                                    manifest.manifest_path
273                                ),
274                            )
275                        }).map(Some)?;
276                    }
277                    (None, None) => {
278                        // Case: Table without row lineage. No action needed.
279                    }
280                }
281            }
282            ManifestContentType::Deletes => {
283                // Deletes never have a first-row-id assigned.
284                manifest.first_row_id = None;
285            }
286        };
287
288        Ok(())
289    }
290}
291
292fn require_row_counts_in_manifest(manifest: &ManifestFile) -> Result<(u64, u64)> {
293    let existing_rows_count = manifest.existing_rows_count.ok_or_else(|| {
294        Error::new(
295            ErrorKind::DataInvalid,
296            format!(
297                "Cannot include a Manifest without existing-rows-count to a table with row lineage enabled. Manifest path: {}",
298                manifest.manifest_path,
299            ),
300        )
301    })?;
302    let added_rows_count = manifest.added_rows_count.ok_or_else(|| {
303        Error::new(
304            ErrorKind::DataInvalid,
305            format!(
306                "Cannot include a Manifest without added-rows-count to a table with row lineage enabled. Manifest path: {}",
307                manifest.manifest_path,
308            ),
309        )
310    })?;
311    Ok((existing_rows_count, added_rows_count))
312}
313
314#[cfg(test)]
315mod test {
316    use std::fs;
317    use std::path::Path;
318    use std::sync::Arc;
319
320    use tempfile::TempDir;
321
322    use super::ManifestListWriter;
323    use crate::encryption::kms::{KeyManagementClient, MemoryKeyManagementClient};
324    use crate::encryption::{EncryptedInputFile, EncryptionManager};
325    use crate::io::{FileIO, FileWrite};
326    use crate::spec::{
327        Datum, FieldSummary, ManifestContentType, ManifestFile, ManifestList,
328        UNASSIGNED_SEQUENCE_NUMBER,
329    };
330
331    #[tokio::test]
332    async fn test_manifest_list_writer_v1() {
333        let expected_manifest_list = ManifestList {
334            entries: vec![ManifestFile {
335                manifest_path: "/opt/bitnami/spark/warehouse/db/table/metadata/10d28031-9739-484c-92db-cdf2975cead4-m0.avro".to_string(),
336                manifest_length: 5806,
337                partition_spec_id: 1,
338                content: ManifestContentType::Data,
339                sequence_number: 0,
340                min_sequence_number: 0,
341                added_snapshot_id: 1646658105718557341,
342                added_files_count: Some(3),
343                existing_files_count: Some(0),
344                deleted_files_count: Some(0),
345                added_rows_count: Some(3),
346                existing_rows_count: Some(0),
347                deleted_rows_count: Some(0),
348                partitions: Some(
349                    vec![FieldSummary { contains_null: false, contains_nan: Some(false), lower_bound: Some(Datum::long(1).to_bytes().unwrap()), upper_bound: Some(Datum::long(1).to_bytes().unwrap())}],
350                ),
351                key_metadata: None,
352                first_row_id: None,
353            }]
354        };
355
356        let temp_dir = TempDir::new().unwrap();
357        let path = temp_dir.path().join("manifest_list_v1.avro");
358        let io = FileIO::new_with_fs();
359        let file_writer = file_writer(&path, io).await;
360
361        let mut writer = ManifestListWriter::v1(file_writer, 1646658105718557341, Some(0));
362        writer
363            .add_manifests(expected_manifest_list.entries.clone().into_iter())
364            .unwrap();
365        writer.close().await.unwrap();
366
367        let bs = fs::read(path).unwrap();
368
369        let manifest_list =
370            ManifestList::parse_with_version(&bs, crate::spec::FormatVersion::V1).unwrap();
371        assert_eq!(manifest_list, expected_manifest_list);
372
373        temp_dir.close().unwrap();
374    }
375
376    #[tokio::test]
377    async fn test_manifest_list_writer_v2() {
378        let snapshot_id = 377075049360453639;
379        let seq_num = 1;
380        let mut expected_manifest_list = ManifestList {
381            entries: vec![ManifestFile {
382                manifest_path: "s3a://icebergdata/demo/s1/t1/metadata/05ffe08b-810f-49b3-a8f4-e88fc99b254a-m0.avro".to_string(),
383                manifest_length: 6926,
384                partition_spec_id: 1,
385                content: ManifestContentType::Data,
386                sequence_number: UNASSIGNED_SEQUENCE_NUMBER,
387                min_sequence_number: UNASSIGNED_SEQUENCE_NUMBER,
388                added_snapshot_id: snapshot_id,
389                added_files_count: Some(1),
390                existing_files_count: Some(0),
391                deleted_files_count: Some(0),
392                added_rows_count: Some(3),
393                existing_rows_count: Some(0),
394                deleted_rows_count: Some(0),
395                partitions: Some(
396                    vec![FieldSummary { contains_null: false, contains_nan: Some(false), lower_bound: Some(Datum::long(1).to_bytes().unwrap()), upper_bound: Some(Datum::long(1).to_bytes().unwrap())}]
397                ),
398                key_metadata: None,
399                first_row_id: None,
400            }]
401        };
402
403        let temp_dir = TempDir::new().unwrap();
404        let path = temp_dir.path().join("manifest_list_v2.avro");
405        let io = FileIO::new_with_fs();
406        let file_writer = file_writer(&path, io).await;
407
408        let mut writer = ManifestListWriter::v2(file_writer, snapshot_id, Some(0), seq_num);
409        writer
410            .add_manifests(expected_manifest_list.entries.clone().into_iter())
411            .unwrap();
412        writer.close().await.unwrap();
413
414        let bs = fs::read(path).unwrap();
415        let manifest_list =
416            ManifestList::parse_with_version(&bs, crate::spec::FormatVersion::V2).unwrap();
417        expected_manifest_list.entries[0].sequence_number = seq_num;
418        expected_manifest_list.entries[0].min_sequence_number = seq_num;
419        assert_eq!(manifest_list, expected_manifest_list);
420
421        temp_dir.close().unwrap();
422    }
423
424    #[tokio::test]
425    async fn test_manifest_list_writer_v3() {
426        let snapshot_id = 377075049360453639;
427        let seq_num = 1;
428        let mut expected_manifest_list = ManifestList {
429            entries: vec![ManifestFile {
430                manifest_path: "s3a://icebergdata/demo/s1/t1/metadata/05ffe08b-810f-49b3-a8f4-e88fc99b254a-m0.avro".to_string(),
431                manifest_length: 6926,
432                partition_spec_id: 1,
433                content: ManifestContentType::Data,
434                sequence_number: UNASSIGNED_SEQUENCE_NUMBER,
435                min_sequence_number: UNASSIGNED_SEQUENCE_NUMBER,
436                added_snapshot_id: snapshot_id,
437                added_files_count: Some(1),
438                existing_files_count: Some(0),
439                deleted_files_count: Some(0),
440                added_rows_count: Some(3),
441                existing_rows_count: Some(0),
442                deleted_rows_count: Some(0),
443                partitions: Some(
444                    vec![FieldSummary { contains_null: false, contains_nan: Some(false), lower_bound: Some(Datum::long(1).to_bytes().unwrap()), upper_bound: Some(Datum::long(1).to_bytes().unwrap())}]
445                ),
446                key_metadata: None,
447                first_row_id: Some(10),
448            }]
449        };
450
451        let temp_dir = TempDir::new().unwrap();
452        let path = temp_dir.path().join("manifest_list_v2.avro");
453        let io = FileIO::new_with_fs();
454        let file_writer = file_writer(&path, io).await;
455
456        let mut writer =
457            ManifestListWriter::v3(file_writer, snapshot_id, Some(0), seq_num, Some(10));
458        writer
459            .add_manifests(expected_manifest_list.entries.clone().into_iter())
460            .unwrap();
461        writer.close().await.unwrap();
462
463        let bs = fs::read(path).unwrap();
464        let manifest_list =
465            ManifestList::parse_with_version(&bs, crate::spec::FormatVersion::V3).unwrap();
466        expected_manifest_list.entries[0].sequence_number = seq_num;
467        expected_manifest_list.entries[0].min_sequence_number = seq_num;
468        expected_manifest_list.entries[0].first_row_id = Some(10);
469        assert_eq!(manifest_list, expected_manifest_list);
470
471        temp_dir.close().unwrap();
472    }
473
474    #[tokio::test]
475    async fn test_manifest_list_writer_v1_as_v2() {
476        let expected_manifest_list = ManifestList {
477            entries: vec![ManifestFile {
478                manifest_path: "/opt/bitnami/spark/warehouse/db/table/metadata/10d28031-9739-484c-92db-cdf2975cead4-m0.avro".to_string(),
479                manifest_length: 5806,
480                partition_spec_id: 1,
481                content: ManifestContentType::Data,
482                sequence_number: 0,
483                min_sequence_number: 0,
484                added_snapshot_id: 1646658105718557341,
485                added_files_count: Some(3),
486                existing_files_count: Some(0),
487                deleted_files_count: Some(0),
488                added_rows_count: Some(3),
489                existing_rows_count: Some(0),
490                deleted_rows_count: Some(0),
491                partitions: Some(
492                    vec![FieldSummary { contains_null: false, contains_nan: Some(false), lower_bound: Some(Datum::long(1).to_bytes().unwrap()), upper_bound: Some(Datum::long(1).to_bytes().unwrap())}]
493                ),
494                key_metadata: None,
495                first_row_id: None,
496            }]
497        };
498
499        let temp_dir = TempDir::new().unwrap();
500        let path = temp_dir.path().join("manifest_list_v1.avro");
501        let io = FileIO::new_with_fs();
502        let file_writer = file_writer(&path, io).await;
503
504        let mut writer = ManifestListWriter::v1(file_writer, 1646658105718557341, Some(0));
505        writer
506            .add_manifests(expected_manifest_list.entries.clone().into_iter())
507            .unwrap();
508        writer.close().await.unwrap();
509
510        let bs = fs::read(path).unwrap();
511
512        let manifest_list =
513            ManifestList::parse_with_version(&bs, crate::spec::FormatVersion::V2).unwrap();
514        assert_eq!(manifest_list, expected_manifest_list);
515
516        temp_dir.close().unwrap();
517    }
518
519    #[tokio::test]
520    async fn test_manifest_list_writer_v1_as_v3() {
521        let expected_manifest_list = ManifestList {
522            entries: vec![ManifestFile {
523                manifest_path: "/opt/bitnami/spark/warehouse/db/table/metadata/10d28031-9739-484c-92db-cdf2975cead4-m0.avro".to_string(),
524                manifest_length: 5806,
525                partition_spec_id: 1,
526                content: ManifestContentType::Data,
527                sequence_number: 0,
528                min_sequence_number: 0,
529                added_snapshot_id: 1646658105718557341,
530                added_files_count: Some(3),
531                existing_files_count: Some(0),
532                deleted_files_count: Some(0),
533                added_rows_count: Some(3),
534                existing_rows_count: Some(0),
535                deleted_rows_count: Some(0),
536                partitions: Some(
537                    vec![FieldSummary { contains_null: false, contains_nan: Some(false), lower_bound: Some(Datum::long(1).to_bytes().unwrap()), upper_bound: Some(Datum::long(1).to_bytes().unwrap())}]
538                ),
539                key_metadata: None,
540                first_row_id: None,
541            }]
542        };
543
544        let temp_dir = TempDir::new().unwrap();
545        let path = temp_dir.path().join("manifest_list_v1.avro");
546        let io = FileIO::new_with_fs();
547        let file_writer = file_writer(&path, io).await;
548
549        let mut writer = ManifestListWriter::v1(file_writer, 1646658105718557341, Some(0));
550        writer
551            .add_manifests(expected_manifest_list.entries.clone().into_iter())
552            .unwrap();
553        writer.close().await.unwrap();
554
555        let bs = fs::read(path).unwrap();
556
557        let manifest_list =
558            ManifestList::parse_with_version(&bs, crate::spec::FormatVersion::V3).unwrap();
559        assert_eq!(manifest_list, expected_manifest_list);
560
561        temp_dir.close().unwrap();
562    }
563
564    #[tokio::test]
565    async fn test_manifest_list_writer_v2_as_v3() {
566        let snapshot_id = 377075049360453639;
567        let seq_num = 1;
568        let mut expected_manifest_list = ManifestList {
569            entries: vec![ManifestFile {
570                manifest_path: "s3a://icebergdata/demo/s1/t1/metadata/05ffe08b-810f-49b3-a8f4-e88fc99b254a-m0.avro".to_string(),
571                manifest_length: 6926,
572                partition_spec_id: 1,
573                content: ManifestContentType::Data,
574                sequence_number: UNASSIGNED_SEQUENCE_NUMBER,
575                min_sequence_number: UNASSIGNED_SEQUENCE_NUMBER,
576                added_snapshot_id: snapshot_id,
577                added_files_count: Some(1),
578                existing_files_count: Some(0),
579                deleted_files_count: Some(0),
580                added_rows_count: Some(3),
581                existing_rows_count: Some(0),
582                deleted_rows_count: Some(0),
583                partitions: Some(
584                    vec![FieldSummary { contains_null: false, contains_nan: Some(false), lower_bound: Some(Datum::long(1).to_bytes().unwrap()), upper_bound: Some(Datum::long(1).to_bytes().unwrap())}]
585                ),
586                key_metadata: None,
587                first_row_id: None,
588            }]
589        };
590
591        let temp_dir = TempDir::new().unwrap();
592        let path = temp_dir.path().join("manifest_list_v2.avro");
593        let io = FileIO::new_with_fs();
594        let file_writer = file_writer(&path, io).await;
595
596        let mut writer = ManifestListWriter::v2(file_writer, snapshot_id, Some(0), seq_num);
597        writer
598            .add_manifests(expected_manifest_list.entries.clone().into_iter())
599            .unwrap();
600        writer.close().await.unwrap();
601
602        let bs = fs::read(path).unwrap();
603
604        let manifest_list =
605            ManifestList::parse_with_version(&bs, crate::spec::FormatVersion::V3).unwrap();
606        expected_manifest_list.entries[0].sequence_number = seq_num;
607        expected_manifest_list.entries[0].min_sequence_number = seq_num;
608        assert_eq!(manifest_list, expected_manifest_list);
609
610        temp_dir.close().unwrap();
611    }
612
613    #[tokio::test]
614    async fn test_manifest_list_writer_v3_encrypted_round_trip() {
615        let (mgr, file_io) = fresh_encryption_manager_and_io();
616        let path = "memory:///manifest_list_v3_encrypted.avro";
617
618        let encrypted_output = mgr.encrypt(file_io.new_output(path).unwrap());
619        let key_metadata = encrypted_output.key_metadata().clone();
620
621        let snapshot_id = 9_000_000_000_000_001i64;
622        let seq_num = 7i64;
623        let mut expected = ManifestList {
624            entries: vec![ManifestFile {
625                manifest_path: "memory:///encrypted/v3_m0.avro".to_string(),
626                manifest_length: 1234,
627                partition_spec_id: 0,
628                content: ManifestContentType::Data,
629                sequence_number: UNASSIGNED_SEQUENCE_NUMBER,
630                min_sequence_number: UNASSIGNED_SEQUENCE_NUMBER,
631                added_snapshot_id: snapshot_id,
632                added_files_count: Some(2),
633                existing_files_count: Some(0),
634                deleted_files_count: Some(0),
635                added_rows_count: Some(10),
636                existing_rows_count: Some(0),
637                deleted_rows_count: Some(0),
638                partitions: Some(vec![FieldSummary {
639                    contains_null: false,
640                    contains_nan: Some(false),
641                    lower_bound: Some(Datum::long(1).to_bytes().unwrap()),
642                    upper_bound: Some(Datum::long(1).to_bytes().unwrap()),
643                }]),
644                key_metadata: None,
645                first_row_id: None,
646            }],
647        };
648
649        let file_writer = encrypted_output.writer().await.unwrap();
650        let mut writer = ManifestListWriter::v3(file_writer, snapshot_id, Some(0), seq_num, None);
651        writer
652            .add_manifests(expected.entries.clone().into_iter())
653            .unwrap();
654        writer.close().await.unwrap();
655
656        let raw_bytes = file_io.new_input(path).unwrap().read().await.unwrap();
657        assert!(
658            ManifestList::parse_with_version(&raw_bytes, crate::spec::FormatVersion::V3).is_err(),
659            "raw bytes should be ciphertext, not parseable as Avro"
660        );
661
662        let plaintext = EncryptedInputFile::new(file_io.new_input(path).unwrap(), key_metadata)
663            .read()
664            .await
665            .unwrap();
666        let manifest_list =
667            ManifestList::parse_with_version(&plaintext, crate::spec::FormatVersion::V3).unwrap();
668
669        expected.entries[0].sequence_number = seq_num;
670        expected.entries[0].min_sequence_number = seq_num;
671        assert_eq!(manifest_list, expected);
672    }
673
674    fn fresh_encryption_manager_and_io() -> (EncryptionManager, FileIO) {
675        let kms = MemoryKeyManagementClient::new();
676        kms.add_master_key("master-1").unwrap();
677        let mgr = EncryptionManager::builder()
678            .kms_client(Arc::new(kms) as Arc<dyn KeyManagementClient>)
679            .table_key_id("master-1")
680            .build();
681        (mgr, FileIO::new_with_memory())
682    }
683
684    async fn file_writer(path: &Path, io: FileIO) -> Box<dyn FileWrite> {
685        io.new_output(path.to_str().unwrap())
686            .unwrap()
687            .writer()
688            .await
689            .unwrap()
690    }
691}
iceberg/spec/manifest_list/writer.rs

iceberg/spec/manifest_list/
writer.rs
