(iceberg-rust) branch main updated: feat: add `construct_ref` for `table_metadata` (#1043)

[fokko]

This is an automated email from the ASF dual-hosted git repository.

fokko pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-rust.git


The following commit(s) were added to refs/heads/main by this push:
     new 82da8260 feat: add `construct_ref` for `table_metadata` (#1043)
82da8260 is described below

commit 82da82605d41b5f6d0bdc4bfaf74a7e60920d88e
Author: ZENOTME <43447882+zeno...@users.noreply.github.com>
AuthorDate: Fri Mar 7 20:54:12 2025 +0800

    feat: add `construct_ref` for `table_metadata` (#1043)
    
    ## Which issue does this PR close?
    
    <!--
    We generally require a GitHub issue to be filed for all bug fixes and
    enhancements and this helps us generate change logs for our releases.
    You can link an issue to this PR using the GitHub syntax. For example
    `Closes #123` indicates that this PR will close issue #123.
    -->
    
    
    ## What changes are included in this PR?
    
    <!--
    Provide a summary of the modifications in this PR. List the main changes
    such as new features, bug fixes, refactoring, or any other updates.
    -->
    
    This PR add a construct_ref to insert the main branch if it's not found
    in refs. ref from:
    
https://github.com/apache/iceberg-python/blob/f45966208dac7c0a2fbe5b16d643a816db2bacb3/pyiceberg/table/metadata.py#L117
    
    ## Are these changes tested?
    
    <!--
    Specify what test covers (unit test, integration test, etc.).
    
    If tests are not included in your PR, please explain why (for example,
    are they covered by existing tests)?
    -->
    
    ---------
    
    Co-authored-by: ZENOTME <st810918...@gmail.com>
---
 crates/iceberg/src/spec/table_metadata.rs         | 17 +++++++++++++
 crates/iceberg/src/spec/table_metadata_builder.rs | 31 +++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/crates/iceberg/src/spec/table_metadata.rs 
b/crates/iceberg/src/spec/table_metadata.rs
index 38204fc1..94f1191b 100644
--- a/crates/iceberg/src/spec/table_metadata.rs
+++ b/crates/iceberg/src/spec/table_metadata.rs
@@ -425,6 +425,22 @@ impl TableMetadata {
             .insert(snapshot.snapshot_id(), Arc::new(snapshot));
     }
 
+    fn construct_refs(&mut self) {
+        if let Some(current_snapshot_id) = self.current_snapshot_id {
+            if !self.refs.contains_key(MAIN_BRANCH) {
+                self.refs
+                    .insert(MAIN_BRANCH.to_string(), SnapshotReference {
+                        snapshot_id: current_snapshot_id,
+                        retention: SnapshotRetention::Branch {
+                            min_snapshots_to_keep: None,
+                            max_snapshot_age_ms: None,
+                            max_ref_age_ms: None,
+                        },
+                    });
+            }
+        }
+    }
+
     /// Normalize this partition spec.
     ///
     /// This is an internal method
@@ -435,6 +451,7 @@ impl TableMetadata {
     pub(super) fn try_normalize(&mut self) -> Result<&mut Self> {
         self.validate_current_schema()?;
         self.normalize_current_snapshot()?;
+        self.construct_refs();
         self.validate_refs()?;
         self.validate_chronological_snapshot_logs()?;
         self.validate_chronological_metadata_logs()?;
diff --git a/crates/iceberg/src/spec/table_metadata_builder.rs 
b/crates/iceberg/src/spec/table_metadata_builder.rs
index 26ee52b5..68ed7698 100644
--- a/crates/iceberg/src/spec/table_metadata_builder.rs
+++ b/crates/iceberg/src/spec/table_metadata_builder.rs
@@ -1220,14 +1220,19 @@ impl From<TableMetadataBuildResult> for TableMetadata {
 
 #[cfg(test)]
 mod tests {
+    use std::fs::File;
+    use std::io::BufReader;
     use std::thread::sleep;
 
     use super::*;
+    use crate::io::FileIOBuilder;
     use crate::spec::{
         BlobMetadata, NestedField, NullOrder, Operation, PartitionSpec, 
PrimitiveType, Schema,
         SnapshotRetention, SortDirection, SortField, StructType, Summary, 
Transform, Type,
         UnboundPartitionField,
     };
+    use crate::table::Table;
+    use crate::TableIdent;
 
     const TEST_LOCATION: &str = "s3://bucket/test/location";
     const LAST_ASSIGNED_COLUMN_ID: i32 = 3;
@@ -2381,4 +2386,30 @@ mod tests {
             last_updated_ms
         );
     }
+
+    #[test]
+    fn test_construct_default_main_branch() {
+        // Load the table without ref
+        let file = File::open(format!(
+            "{}/testdata/table_metadata/{}",
+            env!("CARGO_MANIFEST_DIR"),
+            "TableMetadataV2Valid.json"
+        ))
+        .unwrap();
+        let reader = BufReader::new(file);
+        let resp = serde_json::from_reader::<_, 
TableMetadata>(reader).unwrap();
+
+        let table = Table::builder()
+            .metadata(resp)
+            
.metadata_location("s3://bucket/test/location/metadata/v1.json".to_string())
+            .identifier(TableIdent::from_strs(["ns1", "test1"]).unwrap())
+            .file_io(FileIOBuilder::new("memory").build().unwrap())
+            .build()
+            .unwrap();
+
+        assert_eq!(
+            table.metadata().refs.get(MAIN_BRANCH).unwrap().snapshot_id,
+            table.metadata().current_snapshot_id().unwrap()
+        );
+    }
 }