feat(catalog): Implement update_table for MemoryCatalog

crates/iceberg/src/catalog/memory/catalog.rs

@@ -20,17 +20,16 @@
 use std::collections::HashMap;
 
 use async_trait::async_trait;
-use futures::lock::Mutex;
+use futures::lock::{Mutex, MutexGuard};
 use itertools::Itertools;
-use uuid::Uuid;
 
 use super::namespace_state::NamespaceState;
 use crate::io::FileIO;
 use crate::spec::{TableMetadata, TableMetadataBuilder};
 use crate::table::Table;
 use crate::{
-    Catalog, Error, ErrorKind, Namespace, NamespaceIdent, Result, TableCommit, TableCreation,
-    TableIdent,
+    Catalog, Error, ErrorKind, MetadataLocation, Namespace, NamespaceIdent, Result, TableCommit,
+    TableCreation, TableIdent,
 };
 
 /// namespace `location` property
@@ -45,14 +44,31 @@ pub struct MemoryCatalog {
 }
 
 impl MemoryCatalog {
-    /// Creates an memory catalog.
+    /// Creates a memory catalog.
     pub fn new(file_io: FileIO, warehouse_location: Option<String>) -> Self {
         Self {
             root_namespace_state: Mutex::new(NamespaceState::default()),
             file_io,
             warehouse_location,
         }
     }
+
+    /// Loads a table from the locked namespace state.
+    async fn load_table_from_locked_state(
+        &self,
+        table_ident: &TableIdent,
+        root_namespace_state: &MutexGuard<'_, NamespaceState>,
+    ) -> Result<Table> {
+        let metadata_location = root_namespace_state.get_existing_table_location(table_ident)?;
+        let metadata = TableMetadata::read_from(&self.file_io, metadata_location).await?;
+
+        Table::builder()
+            .identifier(table_ident.clone())
+            .metadata(metadata)
+            .metadata_location(metadata_location.to_string())
+            .file_io(self.file_io.clone())
+            .build()
+    }
 }
 
 #[async_trait]
@@ -203,12 +219,7 @@ impl Catalog for MemoryCatalog {
         let metadata = TableMetadataBuilder::from_table_creation(table_creation)?
             .build()?
             .metadata;
-        let metadata_location = format!(
-            "{}/metadata/{}-{}.metadata.json",
-            &location,
-            0,
-            Uuid::new_v4()
-        );
+        let metadata_location = MetadataLocation::new_with_location(location).to_string();
 
         metadata.write_to(&self.file_io, &metadata_location).await?;
 
@@ -226,15 +237,8 @@ impl Catalog for MemoryCatalog {
     async fn load_table(&self, table_ident: &TableIdent) -> Result<Table> {
         let root_namespace_state = self.root_namespace_state.lock().await;
 
-        let metadata_location = root_namespace_state.get_existing_table_location(table_ident)?;
-        let metadata = TableMetadata::read_from(&self.file_io, metadata_location).await?;
-
-        Table::builder()
-            .file_io(self.file_io.clone())
-            .metadata_location(metadata_location.clone())
-            .metadata(metadata)
-            .identifier(table_ident.clone())
-            .build()
+        self.load_table_from_locked_state(table_ident, &root_namespace_state)
+            .await
     }
 
     /// Drop a table from the catalog.
@@ -289,12 +293,31 @@ impl Catalog for MemoryCatalog {
             .build()
     }
 
-    /// Update a table to the catalog.
-    async fn update_table(&self, _commit: TableCommit) -> Result<Table> {
-        Err(Error::new(
-            ErrorKind::FeatureUnsupported,
-            "MemoryCatalog does not currently support updating tables.",
-        ))
+    /// Update a table in the catalog.
+    async fn update_table(&self, commit: TableCommit) -> Result<Table> {
+        let mut root_namespace_state = self.root_namespace_state.lock().await;
+
+        // Updates the current table version and writes a new metadata file.
+        let current_table = self
+            .load_table_from_locked_state(commit.identifier(), &root_namespace_state)
+            .await?;
+
+        // Apply TableCommit to get staged table
+        let staged_table = commit.apply(current_table)?;
+
+        // Write table metadata to the new location
+        staged_table
+            .metadata()
+            .write_to(
+                staged_table.file_io(),
+                staged_table.metadata_location_result()?,
+            )
+            .await?;
+
+        // Flip the pointer to reference the new metadata file.
+        let updated_table = root_namespace_state.commit_table_update(staged_table)?;
+
+        Ok(updated_table)
     }
 }
 
@@ -303,13 +326,15 @@ mod tests {
     use std::collections::HashSet;
     use std::hash::Hash;
     use std::iter::FromIterator;
+    use std::vec;
 
     use regex::Regex;
     use tempfile::TempDir;
 
     use super::*;
     use crate::io::FileIOBuilder;
     use crate::spec::{NestedField, PartitionSpec, PrimitiveType, Schema, SortOrder, Type};
+    use crate::transaction::{ApplyTransactionAction, Transaction};
 
     fn temp_path() -> String {
         let temp_dir = TempDir::new().unwrap();
@@ -335,7 +360,7 @@ mod tests {
         }
     }
 
-    fn to_set<T: std::cmp::Eq + Hash>(vec: Vec<T>) -> HashSet<T> {
+    fn to_set<T: Eq + Hash>(vec: Vec<T>) -> HashSet<T> {
         HashSet::from_iter(vec)
     }
 
@@ -348,8 +373,8 @@ mod tests {
             .unwrap()
     }
 
-    async fn create_table<C: Catalog>(catalog: &C, table_ident: &TableIdent) {
-        let _ = catalog
+    async fn create_table<C: Catalog>(catalog: &C, table_ident: &TableIdent) -> Table {
+        catalog
             .create_table(
                 &table_ident.namespace,
                 TableCreation::builder()
@@ -358,7 +383,7 @@ mod tests {
                     .build(),
             )
             .await
-            .unwrap();
+            .unwrap()
     }
 
     async fn create_tables<C: Catalog>(catalog: &C, table_idents: Vec<&TableIdent>) {
@@ -367,6 +392,14 @@ mod tests {
         }
     }
 
+    async fn create_table_with_namespace<C: Catalog>(catalog: &C) -> Table {
+        let namespace_ident = NamespaceIdent::new("abc".into());
+        create_namespace(catalog, &namespace_ident).await;
+
+        let table_ident = TableIdent::new(namespace_ident, "test".to_string());
+        create_table(catalog, &table_ident).await
+    }
+
     fn assert_table_eq(table: &Table, expected_table_ident: &TableIdent, expected_schema: &Schema) {
         assert_eq!(table.identifier(), expected_table_ident);
 
@@ -411,7 +444,12 @@ mod tests {
     fn assert_table_metadata_location_matches(table: &Table, regex_str: &str) {
         let actual = table.metadata_location().unwrap().to_string();
         let regex = Regex::new(regex_str).unwrap();
-        assert!(regex.is_match(&actual))
+        assert!(
+            regex.is_match(&actual),
+            "Expected metadata location to match regex, but got location: {} and regex: {}",
+            actual,
+            regex
+        )
     }
 
     #[tokio::test]
@@ -1063,7 +1101,7 @@ mod tests {
         let table_name = "tbl1";
         let expected_table_ident = TableIdent::new(namespace_ident.clone(), table_name.into());
         let expected_table_metadata_location_regex = format!(
-            "^{}/tbl1/metadata/0-{}.metadata.json$",
+            "^{}/tbl1/metadata/00000-{}.metadata.json$",
             namespace_location, UUID_REGEX_STR,
         );
 
@@ -1116,7 +1154,7 @@ mod tests {
         let expected_table_ident =
             TableIdent::new(nested_namespace_ident.clone(), table_name.into());
         let expected_table_metadata_location_regex = format!(
-            "^{}/tbl1/metadata/0-{}.metadata.json$",
+            "^{}/tbl1/metadata/00000-{}.metadata.json$",
             nested_namespace_location, UUID_REGEX_STR,
         );
 
@@ -1157,7 +1195,7 @@ mod tests {
         let table_name = "tbl1";
         let expected_table_ident = TableIdent::new(namespace_ident.clone(), table_name.into());
         let expected_table_metadata_location_regex = format!(
-            "^{}/a/tbl1/metadata/0-{}.metadata.json$",
+            "^{}/a/tbl1/metadata/00000-{}.metadata.json$",
             warehouse_location, UUID_REGEX_STR
         );
 
@@ -1205,7 +1243,7 @@ mod tests {
         let expected_table_ident =
             TableIdent::new(nested_namespace_ident.clone(), table_name.into());
         let expected_table_metadata_location_regex = format!(
-            "^{}/a/b/tbl1/metadata/0-{}.metadata.json$",
+            "^{}/a/b/tbl1/metadata/00000-{}.metadata.json$",
             warehouse_location, UUID_REGEX_STR
         );
 
@@ -1705,7 +1743,7 @@ mod tests {
                 .unwrap_err()
                 .to_string(),
             format!(
-                "TableAlreadyExists => Cannot create table {:? }. Table already exists.",
+                "TableAlreadyExists => Cannot create table {:?}. Table already exists.",
                 &dst_table_ident
             ),
         );
@@ -1754,4 +1792,89 @@ mod tests {
             metadata_location
         );
     }
+
+    #[tokio::test]
+    async fn test_update_table() {
+        let catalog = new_memory_catalog();
+
+        let table = create_table_with_namespace(&catalog).await;
+
+        // Assert the table doesn't contain the update yet
+        assert!(!table.metadata().properties().contains_key("key"));
+
+        // Update table metadata
+        let tx = Transaction::new(&table);
+        let updated_table = tx
+            .update_table_properties()
+            .set("key".to_string(), "value".to_string())
+            .apply(tx)
+            .unwrap()
+            .commit(&catalog)
+            .await
+            .unwrap();
+
+        assert_eq!(
+            updated_table.metadata().properties().get("key").unwrap(),
+            "value"
+        );
+
+        assert_eq!(table.identifier(), updated_table.identifier());
+        assert_eq!(table.metadata().uuid(), updated_table.metadata().uuid());
+        assert!(table.metadata().last_updated_ms() < updated_table.metadata().last_updated_ms());
+        assert_ne!(table.metadata_location(), updated_table.metadata_location());
+
+        println!("left: {:?}", table.metadata().metadata_log());
+        println!("right: {:?}", updated_table.metadata().metadata_log());
+        assert!(
+            table.metadata().metadata_log().len() < updated_table.metadata().metadata_log().len()
+        );
+    }
+
+    #[tokio::test]
+    async fn test_update_table_fails_if_table_doesnt_exist() {
+        let catalog = new_memory_catalog();
+
+        let namespace_ident = NamespaceIdent::new("a".into());
+        create_namespace(&catalog, &namespace_ident).await;
+
+        // This table is not known to the catalog.
+        let table_ident = TableIdent::new(namespace_ident, "test".to_string());
+        let table = build_table(table_ident);
+
+        let tx = Transaction::new(&table);
+        let err = tx
+            .update_table_properties()
+            .set("key".to_string(), "value".to_string())
+            .apply(tx)
+            .unwrap()
+            .commit(&catalog)
+            .await
+            .unwrap_err();
+        assert_eq!(err.kind(), ErrorKind::TableNotFound);
+    }
+
+    fn build_table(ident: TableIdent) -> Table {
+        let file_io = FileIOBuilder::new_fs_io().build().unwrap();
+
+        let temp_dir = TempDir::new().unwrap();
+        let location = temp_dir.path().to_str().unwrap().to_string();
+
+        let table_creation = TableCreation::builder()
+            .name(ident.name().to_string())
+            .schema(simple_table_schema())
+            .location(location)
+            .build();
+        let metadata = TableMetadataBuilder::from_table_creation(table_creation)
+            .unwrap()
+            .build()
+            .unwrap()
+            .metadata;
+
+        Table::builder()
+            .identifier(ident)
+            .metadata(metadata)
+            .file_io(file_io)
+            .build()
+            .unwrap()
+    }
 }

crates/iceberg/src/catalog/memory/namespace_state.rs

@@ -19,6 +19,7 @@ use std::collections::{HashMap, hash_map};
 
 use itertools::Itertools;
 
+use crate::table::Table;
 use crate::{Error, ErrorKind, NamespaceIdent, Result, TableIdent};
 
 // Represents the state of a namespace
@@ -259,7 +260,7 @@ impl NamespaceState {
 
         match namespace.table_metadata_locations.get(table_ident.name()) {
             None => no_such_table_err(table_ident),
-            Some(table_metadadata_location) => Ok(table_metadadata_location),
+            Some(table_metadata_location) => Ok(table_metadata_location),
         }
     }
 
@@ -296,4 +297,22 @@ impl NamespaceState {
             Some(metadata_location) => Ok(metadata_location),
         }
     }
+
+    // Updates the metadata location of the given table or returns an error if it doesn't exist
+    pub(crate) fn commit_table_update(&mut self, staged_table: Table) -> Result<Table> {
+        let namespace = self.get_mut_namespace(staged_table.identifier().namespace())?;
+
+        let _ = namespace
+            .table_metadata_locations
+            .insert(
+                staged_table.identifier().name().to_string(),
+                staged_table.metadata_location_result()?.to_string(),
+            )
+            .ok_or(Error::new(
+                ErrorKind::TableNotFound,
+                format!("No such table: {:?}", staged_table.identifier()),
+            ))?;
+
+        Ok(staged_table)
+    }
 }