add tx action-ish for set_location (does not work well :(

Author: CTTY

crates/catalog/memory/src/catalog.rs

@@ -26,7 +26,7 @@ use iceberg::spec::{TableMetadata, TableMetadataBuilder};
 use iceberg::table::Table;
 use iceberg::{
     Catalog, Error, ErrorKind, Namespace, NamespaceIdent, Result, TableCommit, TableCreation,
-    TableIdent,
+    TableIdent, TableRequirement, TableUpdate,
 };
 use itertools::Itertools;
 use uuid::Uuid;
@@ -277,7 +277,28 @@ impl Catalog for MemoryCatalog {
     }
 
     /// Update a table to the catalog.
-    async fn update_table(&self, _commit: TableCommit) -> Result<Table> {
+    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.",

crates/catalog/memory/src/namespace_state.rs

@@ -262,6 +262,24 @@ impl NamespaceState {
         }
     }
 
+    /// TODO fix this
+    pub(crate) fn update_existing_table_location(
+        &mut self,
+        table_ident: &TableIdent,
+        new_metadata_location: Option<&str>,
+    ) -> Result<()> {
+        if new_metadata_location.is_none() {
+            return Ok(());
+        }
+
+        let mut namespace = self.get_mut_namespace(table_ident.namespace())?;
+        namespace
+            .table_metadata_locations
+            .entry(table_ident.name().to_string())
+            .insert_entry(new_metadata_location.unwrap().into_string());
+        Ok(())
+    }
+
     // Inserts the given table or returns an error if it already exists
     pub(crate) fn insert_new_table(
         &mut self,

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

@@ -0,0 +1,69 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::transaction::Transaction;
+use crate::{Result, TableUpdate};
+
+pub type ActionElement<'a> = Box<dyn TransactionAction<'a>>;
+
+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 struct SetLocation<'a> {
+    pub tx: Transaction<'a>,
+    location: Option<String>,
+}
+
+impl<'a> SetLocation<'a> {
+    pub fn new(tx: Transaction<'a>) -> Self {
+        SetLocation { 
+            tx,
+            location: None
+        }
+    }
+    
+    pub fn set_location(mut self, location: String) -> Self {
+        self.location = Some(location);
+        self
+    }
+}
+
+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)    
+        }
+        
+        self.tx.apply(vec![location.clone().unwrap()], vec![])?;
+        Ok(self.tx)
+        // self.tx.actions.push(Box::new(self));
+    }
+}

crates/iceberg/src/transaction/mod.rs

@@ -17,6 +17,7 @@
 
 //! This module contains transaction api.
 
+mod action;
 mod append;
 mod snapshot;
 mod sort_order;
@@ -32,6 +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::append::FastAppendAction;
 use crate::transaction::sort_order::ReplaceSortOrderAction;
 use crate::{Catalog, Error, ErrorKind, TableCommit, TableRequirement, TableUpdate};
@@ -40,6 +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?
     updates: Vec<TableUpdate>,
     requirements: Vec<TableRequirement>,
 }
@@ -50,6 +53,7 @@ impl<'a> Transaction<'a> {
         Self {
             base_table: table,
             current_table: table.clone(),
+            _actions: vec![],
             updates: vec![],
             requirements: vec![],
         }
@@ -67,6 +71,7 @@ impl<'a> Transaction<'a> {
         Ok(())
     }
 
+    // TODO deprecate this and move the logic to TransactionAction
     fn apply(
         &mut self,
         updates: Vec<TableUpdate>,
@@ -184,9 +189,8 @@ impl<'a> Transaction<'a> {
     }
 
     /// Set the location of table
-    pub fn set_location(mut self, location: String) -> Result<Self> {
-        self.apply(vec![TableUpdate::SetLocation { location }], vec![])?;
-        Ok(self)
+    pub fn set_location(self) -> Result<SetLocation<'a>> {
+        Ok(SetLocation::new(self))
     }
 
     /// Commit transaction.
@@ -196,6 +200,20 @@ impl<'a> Transaction<'a> {
             .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())
+            .await
+            .expect(format!("Failed to refresh table {}", table_commit.identifier()).as_str());
+
+        if self.base_table.metadata() != refreshed_table.metadata() {
+            // TODO raise a real error and retry
+            panic!("Stale base table!")
+        }
 
         catalog.update_table(table_commit).await
     }
@@ -212,6 +230,7 @@ mod tests {
     use crate::table::Table;
     use crate::transaction::Transaction;
     use crate::{TableIdent, TableUpdate};
+    use crate::transaction::action::TransactionAction;
 
     fn make_v1_table() -> Table {
         let file = File::open(format!(
@@ -345,9 +364,12 @@ mod tests {
     fn test_set_location() {
         let table = make_v2_table();
         let tx = Transaction::new(&table);
-        let tx = tx
-            .set_location(String::from("s3://bucket/prefix/new_table"))
-            .unwrap();
+        let set_location = tx
+            .set_location()
+            .unwrap()
+            .set_location(String::from("s3://bucket/prefix/new_table"));
+        
+        let tx = set_location.commit().unwrap();
 
         assert_eq!(
             vec![TableUpdate::SetLocation {