Add logic to commit metadata changes (no retry(yet))

CTTY · CTTY · commit d631b3c93649 · 2025-05-31T21:58:25.000-07:00
diff --git a/crates/catalog/memory/src/catalog.rs b/crates/catalog/memory/src/catalog.rs
@@ -29,6 +29,7 @@ use iceberg::{
     TableIdent, TableRequirement, TableUpdate,
 };
 use itertools::Itertools;
+use regex::Regex;
 use uuid::Uuid;
 
 use crate::namespace_state::NamespaceState;
@@ -278,32 +279,74 @@ impl Catalog for MemoryCatalog {
 
     /// Update a table to the catalog.
     async fn update_table(&self, commit: TableCommit) -> Result<Table> {
-        // TODO persist the updated metadata
-        // let mut root_namespace_state = self.root_namespace_state.lock().await;
-        // let current_table = self.load_table(commit.identifier()).await?;
-        // let updated_staged_table = update_and_stage_table(Some(&current_table), &commit)?;
-        //
-        // if current_table.metadata() == updated_staged_table.metadata() {
-        //     // no changes
-        //     return Ok(current_table);
-        // }
-        //
-        // // write metadata
-        // self.file_io
-        //     .new_output(&updated_staged_table.metadata_location())?
-        //     .write(serde_json::to_vec(updated_staged_table.metadata())?.into())
-        //     .await?;
-        //
-        // root_namespace_state.update_existing_table_location(
-        //     commit.identifier(),
-        //     updated_staged_table.metadata_location(),
-        // )?;
-        // Ok(updated_staged_table)
         Err(Error::new(
             ErrorKind::FeatureUnsupported,
             "MemoryCatalog does not currently support updating tables.",
         ))
     }
+
+    async fn commit_table(&self, base: &Table, current: Table) -> Result<Table> {
+        if base.metadata() == current.metadata() {
+            // no change
+            return Ok(current);
+        }
+
+        let mut root_namespace_state = self.root_namespace_state.lock().await;
+        // TODO: caller needs to retry on the error below
+        let _ = root_namespace_state
+            .check_metadata_location(base.identifier(), base.metadata_location())?;
+
+        let next_metadata_version = if let Some(base_metadata_location) = base.metadata_location() {
+            self.parse_metadata_version(base_metadata_location) + 1
+        } else {
+            0
+        };
+
+        // write metadata
+        let metadata_location = format!(
+            "{}/metadata/{}-{}.metadata.json",
+            current.metadata().location(),
+            next_metadata_version,
+            Uuid::new_v4()
+        );
+
+        // TODO instead of using current.metadata(), build a new metadata with some properties like last_updated_ms updated
+        self.file_io
+            .new_output(&metadata_location)?
+            .write(serde_json::to_vec(current.metadata())?.into())
+            .await?;
+
+        root_namespace_state
+            .update_existing_table_location(current.identifier(), current.metadata_location())?;
+
+        // TODO same here, need to update the metadata location
+        Ok(current)
+    }
+}
+
+// todo move this to metadata?
+fn parse_metadata_version(metadata_location: &str) -> Result<i32> {
+    let pattern = r"(\d+)-([\w-]{36})(?:\.\w+)?\.metadata\.json"; // todo make this constant
+
+    if let Some(metadata_file_name) = metadata_location.split('/').last() {
+        let re = Regex::new(pattern).expect("Failed to parse regex for metadata file!");
+        if let Some(caps) = re.captures(metadata_file_name) {
+            let metadata_version_str = &caps[1];
+            let uuid_str = &caps[2];
+
+            let metadata_version = metadata_version_str
+                .parse()
+                .expect(format!("Invalid metadata version: {metadata_version_str}").as_str());
+            let uuid = Uuid::parse_str(uuid_str)?;
+
+            return Ok(metadata_version);
+        }
+    }
+
+    Err(Error::new(
+        ErrorKind::Unexpected,
+        format!("Unrecognizable metadata location: {metadata_location}"),
+    ))
 }
 
 #[cfg(test)]
diff --git a/crates/catalog/memory/src/namespace_state.rs b/crates/catalog/memory/src/namespace_state.rs
@@ -262,7 +262,28 @@ impl NamespaceState {
         }
     }
 
-    /// TODO fix this
+    pub(crate) fn check_metadata_location(
+        &self,
+        table_ident: &TableIdent,
+        metadata_location: Option<&str>,
+    ) -> Result<()> {
+        let namespace = self.get_namespace(table_ident.namespace())?;
+
+        if namespace
+            .table_metadata_locations
+            .get(table_ident.name())
+            .map(|s| s.as_str())
+            != metadata_location
+        {
+            return Err(Error::new(
+                ErrorKind::DataInvalid,
+                format!("Metadata location does not match for table: {table_ident}!"),
+            ));
+        }
+
+        Ok(())
+    }
+
     pub(crate) fn update_existing_table_location(
         &mut self,
         table_ident: &TableIdent,
diff --git a/crates/iceberg/src/catalog/mod.rs b/crates/iceberg/src/catalog/mod.rs
@@ -94,6 +94,8 @@ pub trait Catalog: Debug + Sync + Send {
 
     /// Update a table to the catalog.
     async fn update_table(&self, commit: TableCommit) -> Result<Table>;
+
+    async fn commit_table(&self, base: &Table, current: Table) -> Result<Table>;
 }
 
 /// NamespaceIdent represents the identifier of a namespace in the catalog.
diff --git a/crates/iceberg/src/error.rs b/crates/iceberg/src/error.rs
@@ -39,7 +39,7 @@ pub enum ErrorKind {
     /// Iceberg data is invalid.
     ///
     /// This error is returned when we try to read a table from iceberg but
-    /// failed to parse it's metadata or data file correctly.
+    /// failed to parse its metadata or data file correctly.
     ///
     /// The table could be invalid or corrupted.
     DataInvalid,
diff --git a/crates/iceberg/src/transaction/action/mod.rs b/crates/iceberg/src/transaction/action/mod.rs
@@ -24,7 +24,7 @@ pub(crate) trait TransactionAction<'a>: Sync {
     /// Apply the pending changes and return the uncommitted changes
     /// TODO is this even needed?
     fn apply(&mut self) -> Result<Option<TableUpdate>>;
-    
+
     /// Commit the changes and apply the changes to the associated transaction
     fn commit(self) -> Result<Transaction<'a>>;
 }
@@ -36,12 +36,9 @@ pub struct SetLocation<'a> {
 
 impl<'a> SetLocation<'a> {
     pub fn new(tx: Transaction<'a>) -> Self {
-        SetLocation { 
-            tx,
-            location: None
-        }
+        SetLocation { tx, location: None }
     }
-    
+
     pub fn set_location(mut self, location: String) -> Self {
         self.location = Some(location);
         self
@@ -51,17 +48,19 @@ impl<'a> SetLocation<'a> {
 impl<'a> TransactionAction<'a> for SetLocation<'a> {
     fn apply(&mut self) -> Result<Option<TableUpdate>> {
         if self.location.is_none() {
-            return Ok(None)
+            return Ok(None);
         }
-        Ok(Some(TableUpdate::SetLocation { location: self.location.clone().unwrap() }))
+        Ok(Some(TableUpdate::SetLocation {
+            location: self.location.clone().unwrap(),
+        }))
     }
-    
+
     fn commit(mut self) -> Result<Transaction<'a>> {
         let location = &mut self.apply()?;
         if location.is_none() {
-            return Ok(self.tx)    
+            return Ok(self.tx);
         }
-        
+
         self.tx.apply(vec![location.clone().unwrap()], vec![])?;
         Ok(self.tx)
         // self.tx.actions.push(Box::new(self));
diff --git a/crates/iceberg/src/transaction/mod.rs b/crates/iceberg/src/transaction/mod.rs
@@ -71,7 +71,6 @@ impl<'a> Transaction<'a> {
         Ok(())
     }
 
-    // TODO deprecate this and move the logic to TransactionAction
     fn apply(
         &mut self,
         updates: Vec<TableUpdate>,
@@ -193,29 +192,52 @@ impl<'a> Transaction<'a> {
         Ok(SetLocation::new(self))
     }
 
+    fn refresh(&mut self, refreshed: Table) {
+        self.base_table = &refreshed;
+        self.current_table = refreshed.clone();
+    }
+
     /// Commit transaction.
-    pub async fn commit(self, catalog: &dyn Catalog) -> Result<Table> {
-        let table_commit = TableCommit::builder()
-            .ident(self.base_table.identifier().clone())
-            .updates(self.updates)
-            .requirements(self.requirements)
-            .build();
+    pub async fn commit(mut self, catalog: &dyn Catalog) -> Result<Table> {
+        // let table_commit = TableCommit::builder()
+        //     .ident(self.base_table.identifier().clone())
+        //     .updates(self.updates)
+        //     .requirements(self.requirements)
+        //     .build();
         if self.base_table.metadata() == self.current_table.metadata() {
             return Ok(self.current_table);
         }
 
         // TODO add refresh() in catalog?
-        let refreshed_table = catalog
-            .load_table(table_commit.identifier())
+        let refreshed = catalog
+            .load_table(self.base_table.identifier())
             .await
-            .expect(format!("Failed to refresh table {}", table_commit.identifier()).as_str());
+            .expect(format!("Failed to refresh table {}", self.base_table.identifier()).as_str());
 
-        if self.base_table.metadata() != refreshed_table.metadata() {
-            // TODO raise a real error and retry
-            panic!("Stale base table!")
+        if self.base_table.metadata() != refreshed.metadata()
+            || self.base_table.metadata_location() != refreshed.metadata_location()
+        {
+            self.refresh(refreshed);
+            self.apply(self.updates, self.requirements) // TODO need create new requirements based on the refreshed table
+                .expect("Failed to re-apply updates"); // re-apply updates
+            // TODO retry on this error
+            return Err(Error::new(
+                ErrorKind::DataInvalid,
+                "Cannot commit: stale base table metadata".to_string(),
+            ));
         }
 
-        catalog.update_table(table_commit).await
+        if self.base_table.metadata() == self.current_table.metadata()
+            && self.base_table.metadata_location() == self.current_table.metadata_location()
+        {
+            // nothing to commit, return current table
+            return Ok(self.current_table);
+        }
+
+        catalog
+            .commit_table(self.base_table, self.current_table)
+            .await
+        // catalog.update_table(table_commit).await
     }
 }
 
@@ -229,8 +251,8 @@ mod tests {
     use crate::spec::{FormatVersion, TableMetadata};
     use crate::table::Table;
     use crate::transaction::Transaction;
-    use crate::{TableIdent, TableUpdate};
     use crate::transaction::action::TransactionAction;
+    use crate::{TableIdent, TableUpdate};
 
     fn make_v1_table() -> Table {
         let file = File::open(format!(
@@ -368,7 +390,7 @@ mod tests {
             .set_location()
             .unwrap()
             .set_location(String::from("s3://bucket/prefix/new_table"));
-        
+
         let tx = set_location.commit().unwrap();
 
         assert_eq!(

Original file line number	Diff line number	Diff line change
`@@ -94,6 +94,8 @@ pub trait Catalog: Debug + Sync + Send {`
`94`	`94`
`95`	`95`	`/// Update a table to the catalog.`
`96`	`96`	`async fn update_table(&self, commit: TableCommit) -> Result<Table>;`
	`97`	`+`
	`98`	`+ async fn commit_table(&self, base: &Table, current: Table) -> Result<Table>;`
`97`	`99`	`}`
`98`	`100`
`99`	`101`	`/// NamespaceIdent represents the identifier of a namespace in the catalog.`