refresh-ish action

Author: CTTY

crates/iceberg/src/transaction/action/mod.rs

@@ -18,25 +18,20 @@
 use crate::transaction::Transaction;
 use crate::{Result, TableUpdate};
 
-pub type ActionElement<'a> = Box<dyn TransactionAction<'a>>;
+pub type PendingAction = Box<dyn TransactionAction>;
 
-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>>;
+pub(crate) trait TransactionAction: Sync {
+    /// Commit the changes and apply the changes to the transaction, return the updated transaction
+    fn commit(self, tx: Transaction) -> Result<Transaction>;
 }
 
-pub struct SetLocation<'a> {
-    pub tx: Transaction<'a>,
-    location: Option<String>,
+pub struct SetLocation {
+    pub location: Option<String>,
 }
 
-impl<'a> SetLocation<'a> {
-    pub fn new(tx: Transaction<'a>) -> Self {
-        SetLocation { tx, location: None }
+impl SetLocation {
+    pub fn new() -> Self {
+        SetLocation { location: None }
     }
 
     pub fn set_location(mut self, location: String) -> Self {
@@ -45,24 +40,13 @@ 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);
-        }
-        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);
+impl TransactionAction for SetLocation {
+    fn commit(self, mut tx: Transaction) -> Result<Transaction> {
+        if let Some(location) = self.location.clone() {
+            tx.apply(vec![TableUpdate::SetLocation { location }], vec![])?;
         }
+        tx.actions.push(Box::new(self));
 
-        self.tx.apply(vec![location.clone().unwrap()], vec![])?;
-        Ok(self.tx)
-        // self.tx.actions.push(Box::new(self));
+        Ok(tx)
     }
 }

crates/iceberg/src/transaction/mod.rs

@@ -33,7 +33,7 @@ use crate::TableUpdate::UpgradeFormatVersion;
 use crate::error::Result;
 use crate::spec::FormatVersion;
 use crate::table::Table;
-use crate::transaction::action::{ActionElement, SetLocation};
+use crate::transaction::action::{SetLocation, TransactionAction, PendingAction};
 use crate::transaction::append::FastAppendAction;
 use crate::transaction::sort_order::ReplaceSortOrderAction;
 use crate::{Catalog, Error, ErrorKind, TableCommit, TableRequirement, TableUpdate};
@@ -42,7 +42,7 @@ use crate::{Catalog, Error, ErrorKind, TableCommit, TableRequirement, TableUpdat
 pub struct Transaction<'a> {
     base_table: &'a Table,
     current_table: Table,
-    _actions: Vec<ActionElement<'a>>, // TODO unused for now, should we use this to reapply actions?
+    actions: Vec<PendingAction>,
     updates: Vec<TableUpdate>,
     requirements: Vec<TableRequirement>,
 }
@@ -53,12 +53,21 @@ impl<'a> Transaction<'a> {
         Self {
             base_table: table,
             current_table: table.clone(),
-            _actions: vec![],
+            actions: vec![],
             updates: vec![],
             requirements: vec![],
         }
     }
 
+    pub fn refresh(old_tx: Transaction<'a>, refreshed: Table) -> Result<Self> {
+        let mut new_tx = Transaction::new(&refreshed.clone());
+        for action in &old_tx.actions {
+            new_tx = action.commit(new_tx)?
+        }
+        
+        Ok(new_tx)
+    }
+
     fn update_table_metadata(&mut self, updates: &[TableUpdate]) -> Result<()> {
         let mut metadata_builder = self.current_table.metadata().clone().into_builder(None);
         for update in updates {
@@ -188,17 +197,12 @@ impl<'a> Transaction<'a> {
     }
 
     /// Set the location of table
-    pub fn set_location(self) -> Result<SetLocation<'a>> {
-        Ok(SetLocation::new(self))
-    }
-
-    fn refresh(&mut self, refreshed: Table) {
-        self.base_table = &refreshed;
-        self.current_table = refreshed.clone();
+    pub fn set_location(self, location: String) -> Result<Transaction<'a>> {
+        Ok(SetLocation::new().set_location(location).commit(self)?)
     }
 
     /// Commit transaction.
-    pub async fn commit(mut self, catalog: &dyn Catalog) -> Result<Table> {
+    pub async fn commit(mut self: Transaction<'a>, catalog: &dyn Catalog) -> Result<Table> {
         // let table_commit = TableCommit::builder()
         //     .ident(self.base_table.identifier().clone())
         //     .updates(self.updates)
@@ -217,14 +221,14 @@ impl<'a> Transaction<'a> {
         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(),
-            ));
+            // refresh table
+            let new_tx = Transaction::refresh(self, refreshed)?;
+            return new_tx.commit(catalog).await
+            // TODO instead of refreshing directly, retry on this error
+            // return Err(Error::new(
+            //     ErrorKind::DataInvalid,
+            //     "Cannot commit: stale base table metadata".to_string(),
+            // ));
         }
 
         if self.base_table.metadata() == self.current_table.metadata()
@@ -386,12 +390,9 @@ mod tests {
     fn test_set_location() {
         let table = make_v2_table();
         let tx = Transaction::new(&table);
-        let set_location = tx
-            .set_location()
-            .unwrap()
-            .set_location(String::from("s3://bucket/prefix/new_table"));
-
-        let tx = set_location.commit().unwrap();
+        let tx = tx
+            .set_location(String::from("s3://bucket/prefix/new_table"))
+            .unwrap();
 
         assert_eq!(
             vec![TableUpdate::SetLocation {