Have tx own base table, mark tx_commit_res as not used, need to fix retry

Author: CTTY

crates/iceberg/Cargo.toml

@@ -55,6 +55,7 @@ arrow-select = { workspace = true }
 arrow-string = { workspace = true }
 async-std = { workspace = true, optional = true, features = ["attributes"] }
 async-trait = { workspace = true }
+tokio-retry2 = { version = "0.5.7" }
 base64 = { workspace = true }
 bimap = { workspace = true }
 bytes = { workspace = true }

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

@@ -16,26 +16,29 @@
 // under the License.
 
 use std::mem::take;
-
+use std::sync::Arc;
+use async_trait::async_trait;
 use crate::{Error, ErrorKind, Result, TableRequirement, TableUpdate};
+use crate::transaction::Transaction;
 
 pub type BoxedTransactionAction = Arc<dyn TransactionAction>;
 
 #[async_trait]
 pub(crate) trait TransactionAction: Sync + Send {
     /// Commit the changes and apply the changes to the transaction,
     /// return the transaction with the updated current_table
-    fn commit(self: Arc<Self>, tx:  &Transaction) -> Result<TransactionActionCommit>;
+    fn commit(self: Arc<Self>, tx: &mut Transaction) -> Result<()>;
 }
 
 pub struct TransactionActionCommit {
-    action: Option<PendingAction>,
+    action: Option<BoxedTransactionAction>,
     updates: Vec<TableUpdate>,
     requirements: Vec<TableRequirement>,
 }
 
+// TODO probably we don't need this?
 impl TransactionActionCommit {
-    pub fn take_action(&mut self) -> Option<PendingAction> {
+    pub fn take_action(&mut self) -> Option<BoxedTransactionAction> {
         take(&mut self.action)
     }
 
@@ -64,7 +67,7 @@ impl SetLocation {
 }
 
 impl TransactionAction for SetLocation {
-    fn commit(self: Box<Self>) -> Result<TransactionActionCommit> {
+    fn commit(self: Arc<Self>, tx: &mut Transaction) -> Result<()> {
         let updates: Vec<TableUpdate>;
         let requirements: Vec<TableRequirement>;
         if let Some(location) = self.location.clone() {
@@ -76,11 +79,9 @@ impl TransactionAction for SetLocation {
                 "Location is not set for SetLocation!",
             ));
         }
-
-        Ok(TransactionActionCommit {
-            action: Some(self),
-            updates,
-            requirements,
-        })
+        
+        tx.actions.push(self);
+        
+        tx.apply(updates, requirements)
     }
 }

crates/iceberg/src/transaction/append.rs

@@ -31,15 +31,15 @@ use crate::writer::file_writer::ParquetWriter;
 use crate::{Error, ErrorKind};
 
 /// FastAppendAction is a transaction action for fast append data files to the table.
-pub struct FastAppendAction<'a> {
-    snapshot_produce_action: SnapshotProduceAction<'a>,
+pub struct FastAppendAction {
+    snapshot_produce_action: SnapshotProduceAction,
     check_duplicate: bool,
 }
 
-impl<'a> FastAppendAction<'a> {
+impl FastAppendAction {
     #[allow(clippy::too_many_arguments)]
     pub(crate) fn new(
-        tx: Transaction<'a>,
+        tx: Transaction,
         snapshot_id: i64,
         commit_uuid: Uuid,
         key_metadata: Vec<u8>,
@@ -89,7 +89,7 @@ impl<'a> FastAppendAction<'a> {
     /// Specifically, schema compatibility checks and support for adding to partitioned tables
     /// have not yet been implemented.
     #[allow(dead_code)]
-    async fn add_parquet_files(mut self, file_path: Vec<String>) -> Result<Transaction<'a>> {
+    async fn add_parquet_files(mut self, file_path: Vec<String>) -> Result<Transaction> {
         if !self
             .snapshot_produce_action
             .tx
@@ -119,7 +119,7 @@ impl<'a> FastAppendAction<'a> {
     }
 
     /// Finished building the action and apply it to the transaction.
-    pub async fn apply(self) -> Result<Transaction<'a>> {
+    pub async fn apply(self) -> Result<Transaction> {
         // Checks duplicate files
         if self.check_duplicate {
             let new_files: HashSet<&str> = self
@@ -185,14 +185,14 @@ impl SnapshotProduceOperation for FastAppendOperation {
 
     async fn delete_entries(
         &self,
-        _snapshot_produce: &SnapshotProduceAction<'_>,
+        _snapshot_produce: &SnapshotProduceAction,
     ) -> Result<Vec<ManifestEntry>> {
         Ok(vec![])
     }
 
     async fn existing_manifest(
         &self,
-        snapshot_produce: &SnapshotProduceAction<'_>,
+        snapshot_produce: &SnapshotProduceAction,
     ) -> Result<Vec<ManifestFile>> {
         let Some(snapshot) = snapshot_produce
             .tx
@@ -234,7 +234,7 @@ mod tests {
     #[tokio::test]
     async fn test_empty_data_append_action() {
         let table = make_v2_minimal_table();
-        let tx = Transaction::new(&table);
+        let tx = Transaction::new(table);
         let mut action = tx.fast_append(None, vec![]).unwrap();
         action.add_data_files(vec![]).unwrap();
         assert!(action.apply().await.is_err());
@@ -243,7 +243,7 @@ mod tests {
     #[tokio::test]
     async fn test_set_snapshot_properties() {
         let table = make_v2_minimal_table();
-        let tx = Transaction::new(&table);
+        let tx = Transaction::new(table.clone());
         let mut action = tx.fast_append(None, vec![]).unwrap();
 
         let mut snapshot_properties = HashMap::new();
@@ -281,7 +281,7 @@ mod tests {
     #[tokio::test]
     async fn test_fast_append_action() {
         let table = make_v2_minimal_table();
-        let tx = Transaction::new(&table);
+        let tx = Transaction::new(table.clone());
         let mut action = tx.fast_append(None, vec![]).unwrap();
 
         // check add data file with incompatible partition value
@@ -367,7 +367,7 @@ mod tests {
     async fn test_add_existing_parquet_files_to_unpartitioned_table() {
         let mut fixture = TableTestFixture::new_unpartitioned();
         fixture.setup_unpartitioned_manifest_files().await;
-        let tx = crate::transaction::Transaction::new(&fixture.table);
+        let tx = crate::transaction::Transaction::new(fixture.table.clone());
 
         let file_paths = vec![
             format!("{}/1.parquet", &fixture.table_location),

crates/iceberg/src/transaction/mod.rs

@@ -35,26 +35,26 @@ use crate::error::Result;
 use crate::spec::FormatVersion;
 use crate::table::Table;
 use crate::transaction::action::{
-    PendingAction, SetLocation, TransactionAction, TransactionActionCommit,
+    BoxedTransactionAction, SetLocation, TransactionAction, TransactionActionCommit,
 };
 use crate::transaction::append::FastAppendAction;
 use crate::transaction::sort_order::ReplaceSortOrderAction;
 use crate::{Catalog, Error, ErrorKind, TableCommit, TableRequirement, TableUpdate};
 
 /// Table transaction.
-pub struct Transaction<'a> {
-    base_table: &'a Table,
+pub struct Transaction {
+    base_table: Table,
     current_table: Table,
-    actions: Vec<PendingAction>,
+    actions: Vec<BoxedTransactionAction>,
     updates: Vec<TableUpdate>,
     requirements: Vec<TableRequirement>,
 }
 
-impl<'a> Transaction<'a> {
+impl Transaction {
     /// Creates a new transaction.
-    pub fn new(table: &'a Table) -> Self {
+    pub fn new(table: Table) -> Self {
         Self {
-            base_table: table,
+            base_table: table.clone(),
             current_table: table.clone(),
             actions: vec![],
             updates: vec![],
@@ -174,7 +174,7 @@ impl<'a> Transaction<'a> {
         self,
         commit_uuid: Option<Uuid>,
         key_metadata: Vec<u8>,
-    ) -> Result<FastAppendAction<'a>> {
+    ) -> Result<FastAppendAction> {
         let snapshot_id = self.generate_unique_snapshot_id();
         FastAppendAction::new(
             self,
@@ -186,7 +186,7 @@ impl<'a> Transaction<'a> {
     }
 
     /// Creates replace sort order action.
-    pub fn replace_sort_order(self) -> ReplaceSortOrderAction<'a> {
+    pub fn replace_sort_order(self) -> ReplaceSortOrderAction {
         ReplaceSortOrderAction {
             tx: self,
             sort_fields: vec![],
@@ -205,8 +205,7 @@ impl<'a> Transaction<'a> {
     /// Set the location of table
     pub fn set_location(mut self, location: String) -> Result<Self> {
         let set_location = SetLocation::new().set_location(location);
-        self.apply_commit_result(Box::new(set_location).commit()?)
-            .expect("Some error msg");
+        Arc::new(set_location).commit(&mut self)?;
         Ok(self)
     }
 
@@ -243,14 +242,13 @@ impl<'a> Transaction<'a> {
             || self.base_table.metadata_location() != refreshed.metadata_location()
         {
             // current base is stale, use refreshed as base and re-apply transaction actions
-            self.base_table = &refreshed.clone();
+            self.base_table = refreshed.clone();
             self.current_table = refreshed.clone();
 
             let pending_actions = take(&mut self.actions);
 
             for action in pending_actions {
-                self.apply_commit_result(action.commit()?)
-                    .expect("Failed to apply updates!");
+                action.commit(self).expect("Failed to apply updates!");
             }
         }
 
@@ -347,7 +345,7 @@ mod tests {
     #[test]
     fn test_upgrade_table_version_v1_to_v2() {
         let table = make_v1_table();
-        let tx = Transaction::new(&table);
+        let tx = Transaction::new(table);
         let tx = tx.upgrade_table_version(FormatVersion::V2).unwrap();
 
         assert_eq!(
@@ -361,7 +359,7 @@ mod tests {
     #[test]
     fn test_upgrade_table_version_v2_to_v2() {
         let table = make_v2_table();
-        let tx = Transaction::new(&table);
+        let tx = Transaction::new(table);
         let tx = tx.upgrade_table_version(FormatVersion::V2).unwrap();
 
         assert!(
@@ -377,7 +375,7 @@ mod tests {
     #[test]
     fn test_downgrade_table_version() {
         let table = make_v2_table();
-        let tx = Transaction::new(&table);
+        let tx = Transaction::new(table);
         let tx = tx.upgrade_table_version(FormatVersion::V1);
 
         assert!(tx.is_err(), "Downgrade table version should fail!");
@@ -386,7 +384,7 @@ mod tests {
     #[test]
     fn test_set_table_property() {
         let table = make_v2_table();
-        let tx = Transaction::new(&table);
+        let tx = Transaction::new(table);
         let tx = tx
             .set_properties(HashMap::from([("a".to_string(), "b".to_string())]))
             .unwrap();
@@ -402,7 +400,7 @@ mod tests {
     #[test]
     fn test_remove_property() {
         let table = make_v2_table();
-        let tx = Transaction::new(&table);
+        let tx = Transaction::new(table);
         let tx = tx
             .remove_properties(vec!["a".to_string(), "b".to_string()])
             .unwrap();
@@ -418,7 +416,7 @@ mod tests {
     #[test]
     fn test_set_location() {
         let table = make_v2_table();
-        let tx = Transaction::new(&table);
+        let tx = Transaction::new(table);
         let tx = tx
             .set_location(String::from("s3://bucket/prefix/new_table"))
             .unwrap();
@@ -434,7 +432,7 @@ mod tests {
     #[tokio::test]
     async fn test_transaction_apply_upgrade() {
         let table = make_v1_table();
-        let tx = Transaction::new(&table);
+        let tx = Transaction::new(table);
         // Upgrade v1 to v1, do nothing.
         let tx = tx.upgrade_table_version(FormatVersion::V1).unwrap();
         // Upgrade v1 to v2, success.

crates/iceberg/src/transaction/snapshot.rs

@@ -59,8 +59,8 @@ pub(crate) trait ManifestProcess: Send + Sync {
     fn process_manifests(&self, manifests: Vec<ManifestFile>) -> Vec<ManifestFile>;
 }
 
-pub(crate) struct SnapshotProduceAction<'a> {
-    pub tx: Transaction<'a>,
+pub(crate) struct SnapshotProduceAction {
+    pub tx: Transaction,
     snapshot_id: i64,
     key_metadata: Vec<u8>,
     commit_uuid: Uuid,
@@ -72,9 +72,9 @@ pub(crate) struct SnapshotProduceAction<'a> {
     manifest_counter: RangeFrom<u64>,
 }
 
-impl<'a> SnapshotProduceAction<'a> {
+impl SnapshotProduceAction {
     pub(crate) fn new(
-        tx: Transaction<'a>,
+        tx: Transaction,
         snapshot_id: i64,
         key_metadata: Vec<u8>,
         commit_uuid: Uuid,
@@ -310,7 +310,7 @@ impl<'a> SnapshotProduceAction<'a> {
         mut self,
         snapshot_produce_operation: OP,
         process: MP,
-    ) -> Result<Transaction<'a>> {
+    ) -> Result<Transaction> {
         let new_manifests = self
             .manifest_file(&snapshot_produce_operation, &process)
             .await?;

crates/iceberg/src/transaction/sort_order.rs

@@ -21,12 +21,12 @@ use crate::transaction::Transaction;
 use crate::{Error, ErrorKind, TableRequirement, TableUpdate};
 
 /// Transaction action for replacing sort order.
-pub struct ReplaceSortOrderAction<'a> {
-    pub tx: Transaction<'a>,
+pub struct ReplaceSortOrderAction {
+    pub tx: Transaction,
     pub sort_fields: Vec<SortField>,
 }
 
-impl<'a> ReplaceSortOrderAction<'a> {
+impl ReplaceSortOrderAction {
     /// Adds a field for sorting in ascending order.
     pub fn asc(self, name: &str, null_order: NullOrder) -> Result<Self> {
         self.add_sort_field(name, SortDirection::Ascending, null_order)
@@ -38,7 +38,7 @@ impl<'a> ReplaceSortOrderAction<'a> {
     }
 
     /// Finished building the action and apply it to the transaction.
-    pub fn apply(mut self) -> Result<Transaction<'a>> {
+    pub fn apply(mut self) -> Result<Transaction> {
         let unbound_sort_order = SortOrder::builder()
             .with_fields(self.sort_fields)
             .build_unbound()?;
@@ -114,7 +114,7 @@ mod tests {
     #[test]
     fn test_replace_sort_order() {
         let table = make_v2_table();
-        let tx = Transaction::new(&table);
+        let tx = Transaction::new(table);
         let tx = tx.replace_sort_order().apply().unwrap();
 
         assert_eq!(