add UT and move out remove in separate PR

Author: mann-db

kernel/src/actions/mod.rs

@@ -877,9 +877,9 @@ pub(crate) struct CheckpointMetadata {
 #[derive(Debug, Clone, PartialEq, Eq, ToSchema, IntoEngineData)]
 #[internal_api]
 pub(crate) struct DomainMetadata {
-    pub(crate) domain: String,
-    pub(crate) configuration: String,
-    pub(crate) removed: bool,
+    domain: String,
+    configuration: String,
+    removed: bool,
 }
 
 impl DomainMetadata {
@@ -889,6 +889,18 @@ impl DomainMetadata {
     pub(crate) fn is_internal(&self) -> bool {
         self.domain.starts_with(INTERNAL_DOMAIN_PREFIX)
     }
+
+    pub fn domain(&self) -> &str {
+        &self.domain
+    }
+
+    pub(crate) fn new(domain: String, configuration: String, removed: bool) -> Self {
+        Self {
+            domain,
+            configuration,
+            removed,
+        }
+    }
 }
 
 #[cfg(test)]

kernel/src/transaction/mod.rs

@@ -1,10 +1,13 @@
-use std::collections::{HashSet, HashMap};
+use std::collections::{HashMap, HashSet};
 use std::iter;
 use std::sync::{Arc, LazyLock};
 use std::time::{SystemTime, UNIX_EPOCH};
 
 use crate::actions::domain_metadata::DomainMetadataMap;
-use crate::actions::{DomainMetadata, get_log_add_schema, get_log_commit_info_schema, get_log_domain_metadata_schema, get_log_txn_schema};
+use crate::actions::{
+    get_log_add_schema, get_log_commit_info_schema, get_log_domain_metadata_schema,
+    get_log_txn_schema, DomainMetadata,
+};
 use crate::actions::{CommitInfo, SetTransaction};
 use crate::error::Error;
 use crate::expressions::UnaryExpressionOp;
@@ -282,38 +285,34 @@ impl Transaction {
         self.add_files_metadata.push(add_metadata);
     }
 
-    pub fn add_domain_metadata(&mut self, domain: String, configuration: String) -> DeltaResult<()> {
-        let domain_metadata = DomainMetadata {
-            domain: domain.clone(),
-            configuration: configuration.to_string(),
-            removed: false,
-        };
+    /// Set domain metadata to be written to the Delta log. Each domain can only be modified once per
+    /// transaction. System-controlled domains (those starting with `delta.`) cannot be modified.
+    pub fn set_domain_metadata(
+        &mut self,
+        domain: String,
+        configuration: String,
+    ) -> DeltaResult<()> {
+        let domain_metadata = DomainMetadata::new(domain.clone(), configuration, false);
         self.validate_domain_metadata(&domain_metadata)?;
         self.domain_metadata.insert(domain, domain_metadata);
         Ok(())
     }
 
     fn validate_domain_metadata(&self, domain_metadata: &DomainMetadata) -> DeltaResult<()> {
         if domain_metadata.is_internal() {
-            return Err(Error::Generic("User metadata cannot be added to system-controlled 'delta.*' domain".to_string()));
+            return Err(Error::Generic(
+                "User metadata cannot be added to system-controlled 'delta.*' domain".to_string(),
+            ));
         }
 
-        if self.domain_metadata.contains_key(&domain_metadata.domain) {
-            return Err(Error::Generic(format!("Metadata for domain {} already specified in this transaction", domain_metadata.domain)));
+        if self.domain_metadata.contains_key(domain_metadata.domain()) {
+            return Err(Error::Generic(format!(
+                "Metadata for domain {} already specified in this transaction",
+                domain_metadata.domain()
+            )));
         }
         Ok(())
     }
-
-    pub fn remove_domain_metadata(&mut self, domain: String) -> DeltaResult<()> {
-        let domain_metadata = DomainMetadata {
-            domain: domain.clone(),
-            configuration: "".to_string(),
-            removed: true,
-        };
-        self.validate_domain_metadata(&domain_metadata)?;
-        self.domain_metadata.insert(domain, domain_metadata);
-        Ok(())
-    }
 }
 
 /// Convert file metadata provided by the engine into protocol-compliant add actions.

kernel/tests/write.rs

@@ -1152,3 +1152,90 @@ async fn test_shredded_variant_read_rejection() -> Result<(), Box<dyn std::error
 
     Ok(())
 }
+
+#[tokio::test]
+async fn test_set_domain_metadata_basic() -> Result<(), Box<dyn std::error::Error>> {
+    let _ = tracing_subscriber::fmt::try_init();
+
+    let schema = Arc::new(StructType::new(vec![StructField::nullable(
+        "number",
+        DataType::INTEGER,
+    )]));
+
+    for (table_url, engine, store, table_name) in
+        setup_test_tables(schema.clone(), &[], None, "test_table").await?
+    {
+        let snapshot = Arc::new(Snapshot::builder(table_url.clone()).build(&engine)?);
+        let mut txn = snapshot.transaction()?;
+
+        // Set multiple domain metadata
+        let domain1 = "app.config";
+        let config1 = r#"{"version": 1}"#;
+        let domain2 = "spark.settings";
+        let config2 = r#"{"cores": 4}"#;
+
+        txn.set_domain_metadata(domain1.into(), config1.into())?;
+        txn.set_domain_metadata(domain2.into(), config2.into())?;
+        txn.commit(&engine)?;
+
+        let commit_data = store
+            .get(&Path::from(format!(
+                "/{table_name}/_delta_log/00000000000000000001.json"
+            )))
+            .await?
+            .bytes()
+            .await?;
+        let actions: Vec<serde_json::Value> = Deserializer::from_slice(&commit_data)
+            .into_iter()
+            .try_collect()?;
+
+        let domain_actions: Vec<_> = actions
+            .iter()
+            .filter(|v| v.get("domainMetadata").is_some())
+            .collect();
+        assert_eq!(domain_actions.len(), 2);
+
+        // Check domains and their configurations
+        for action in &domain_actions {
+            let domain = action["domainMetadata"]["domain"].as_str().unwrap();
+            let config = action["domainMetadata"]["configuration"].as_str().unwrap();
+            assert!(!action["domainMetadata"]["removed"].as_bool().unwrap());
+
+            match domain {
+                d if d == domain1 => assert_eq!(config, config1),
+                d if d == domain2 => assert_eq!(config, config2),
+                _ => panic!("Unexpected domain: {}", domain),
+            }
+        }
+    }
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_set_domain_metadata_errors() -> Result<(), Box<dyn std::error::Error>> {
+    let _ = tracing_subscriber::fmt::try_init();
+
+    let schema = Arc::new(StructType::new(vec![StructField::nullable(
+        "number",
+        DataType::INTEGER,
+    )]));
+
+    for (table_url, engine, _, _) in
+        setup_test_tables(schema.clone(), &[], None, "test_table").await?
+    {
+        let snapshot = Arc::new(Snapshot::builder(table_url.clone()).build(&engine)?);
+        let mut txn = snapshot.transaction()?;
+
+        // System domain rejection
+        assert!(txn
+            .set_domain_metadata("delta.system".into(), "config".into())
+            .is_err());
+
+        // Duplicate domain rejection
+        txn.set_domain_metadata("app.config".into(), "v1".into())?;
+        assert!(txn
+            .set_domain_metadata("app.config".into(), "v2".into())
+            .is_err());
+    }
+    Ok(())
+}