Add optional default commit metadata to Repository (#860)

Author: mpiannucci

icechunk-python/python/icechunk/_icechunk_python.pyi

@@ -964,6 +964,10 @@ class PyRepository:
     def save_config(self) -> None: ...
     def config(self) -> RepositoryConfig: ...
     def storage(self) -> Storage: ...
+    def set_default_commit_metadata(
+        self, metadata: dict[str, Any] | None = None
+    ) -> None: ...
+    def default_commit_metadata(self) -> dict[str, Any] | None: ...
     def async_ancestry(
         self,
         *,

icechunk-python/python/icechunk/repository.py

@@ -1,6 +1,6 @@
 import datetime
 from collections.abc import AsyncIterator, Iterator
-from typing import Self, cast
+from typing import Any, Self, cast
 
 from icechunk._icechunk_python import (
     Diff,
@@ -216,6 +216,36 @@ def storage(self) -> Storage:
         """
         return self._repository.storage()
 
+    def set_default_commit_metadata(self, metadata: dict[str, Any] | None = None) -> None:
+        """
+        Set the default commit metadata for the repository. This is useful for providing
+        addition static system conexted metadata to all commits.
+
+        When a commit is made, the metadata will be merged with the metadata provided, with any
+        duplicate keys being overwritten by the metadata provided in the commit.
+
+        !!! warning
+            This metadata is only applied to sessions that are created after this call. Any open
+            writable sessions will not be affected and will not use the new default metadata.
+
+        Parameters
+        ----------
+        metadata : dict[str, Any], optional
+            The default commit metadata.
+        """
+        return self._repository.set_default_commit_metadata(metadata)
+
+    def default_commit_metadata(self) -> dict[str, Any] | None:
+        """
+        Get the current configured default commit metadata for the repository.
+
+        Returns
+        -------
+        dict[str, Any] | None
+            The default commit metadata.
+        """
+        return self._repository.default_commit_metadata()
+
     def ancestry(
         self,
         *,

icechunk-python/src/repository.rs

@@ -550,6 +550,25 @@ impl PyRepository {
         PyStorage(Arc::clone(self.0.storage()))
     }
 
+    #[pyo3(signature = (metadata))]
+    pub fn set_default_commit_metadata(
+        &self,
+        metadata: Option<PySnapshotProperties>,
+    ) -> PyResult<()> {
+        let metadata = metadata.map(|m| m.into());
+        pyo3_async_runtimes::tokio::get_runtime().block_on(async move {
+            self.0.set_default_commit_metadata(metadata).await;
+            Ok(())
+        })
+    }
+
+    pub fn default_commit_metadata(&self) -> PyResult<Option<PySnapshotProperties>> {
+        pyo3_async_runtimes::tokio::get_runtime().block_on(async move {
+            let metadata = self.0.default_commit_metadata().await;
+            Ok(metadata.map(|m| m.into()))
+        })
+    }
+
     /// Returns an object that is both a sync and an async iterator
     #[pyo3(signature = (*, branch = None, tag = None, snapshot_id = None))]
     pub fn async_ancestry(

icechunk-python/tests/test_timetravel.py

@@ -293,3 +293,17 @@ async def test_session_with_as_of() -> None:
         expected_children = {f"child {j}" for j in range(i)}
         actual_children = {g[0] for g in group.members()}
         assert expected_children == actual_children
+
+
+async def test_default_commit_metadata() -> None:
+    repo = ic.Repository.create(
+        storage=ic.in_memory_storage(),
+    )
+
+    repo.set_default_commit_metadata({"user": "test"})
+    session = repo.writable_session("main")
+    root = zarr.group(store=session.store, overwrite=True)
+    root.create_group("child")
+    sid = session.commit("root")
+    snap = next(repo.ancestry(snapshot_id=sid))
+    assert snap.metadata == {"user": "test"}

icechunk/src/asset_manager.rs

@@ -323,6 +323,16 @@ impl AssetManager {
             .get_snapshot_last_modified(&self.storage_settings, snapshot_id)
             .await?)
     }
+
+    #[instrument(skip(self))]
+    pub async fn fetch_snapshot_info(
+        &self,
+        snapshot_id: &SnapshotId,
+    ) -> RepositoryResult<SnapshotInfo> {
+        let snapshot = self.fetch_snapshot(snapshot_id).await?;
+        let info = snapshot.as_ref().try_into()?;
+        Ok(info)
+    }
 }
 
 fn binary_file_header(

icechunk/src/repository.rs

@@ -16,6 +16,7 @@ use futures::{
 use regex::bytes::Regex;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
+use tokio::sync::RwLock;
 use tokio::task::JoinError;
 use tracing::{Instrument, debug, error, instrument, trace};
 
@@ -27,7 +28,9 @@ use crate::{
     format::{
         IcechunkFormatError, IcechunkFormatErrorKind, ManifestId, NodeId, Path,
         SnapshotId,
-        snapshot::{ManifestFileInfo, NodeData, Snapshot, SnapshotInfo},
+        snapshot::{
+            ManifestFileInfo, NodeData, Snapshot, SnapshotInfo, SnapshotProperties,
+        },
         transaction_log::{Diff, DiffBuilder},
     },
     refs::{
@@ -136,6 +139,8 @@ pub struct Repository {
     asset_manager: Arc<AssetManager>,
     virtual_resolver: Arc<VirtualChunkResolver>,
     virtual_chunk_credentials: HashMap<ContainerName, Credentials>,
+    #[serde(skip)]
+    default_commit_metadata: Arc<RwLock<Option<SnapshotProperties>>>,
 }
 
 impl Repository {
@@ -307,6 +312,7 @@ impl Repository {
             virtual_resolver,
             asset_manager,
             virtual_chunk_credentials,
+            default_commit_metadata: Arc::new(RwLock::new(None)),
         })
     }
 
@@ -375,6 +381,20 @@ impl Repository {
         .await
     }
 
+    #[instrument(skip_all)]
+    pub async fn set_default_commit_metadata(
+        &self,
+        metadata: Option<SnapshotProperties>,
+    ) {
+        let mut guard = self.default_commit_metadata.write().await;
+        *guard = metadata;
+    }
+
+    #[instrument(skip_all)]
+    pub async fn default_commit_metadata(&self) -> Option<SnapshotProperties> {
+        self.default_commit_metadata.read().await.clone()
+    }
+
     #[instrument(skip(storage, config))]
     pub(crate) async fn store_config(
         storage: &(dyn Storage + Send + Sync),
@@ -754,6 +774,7 @@ impl Repository {
             self.virtual_resolver.clone(),
             branch.to_string(),
             ref_data.snapshot.clone(),
+            self.default_commit_metadata.read().await.clone(),
         );
 
         self.preload_manifests(ref_data.snapshot);

icechunk/src/session.rs

@@ -166,6 +166,7 @@ pub struct Session {
     branch_name: Option<String>,
     snapshot_id: SnapshotId,
     change_set: ChangeSet,
+    default_commit_metadata: Option<SnapshotProperties>,
 }
 
 impl Session {
@@ -186,9 +187,11 @@ impl Session {
             branch_name: None,
             snapshot_id,
             change_set: ChangeSet::default(),
+            default_commit_metadata: None,
         }
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn create_writable_session(
         config: RepositoryConfig,
         storage_settings: storage::Settings,
@@ -197,6 +200,7 @@ impl Session {
         virtual_resolver: Arc<VirtualChunkResolver>,
         branch_name: String,
         snapshot_id: SnapshotId,
+        default_commit_metadata: Option<SnapshotProperties>,
     ) -> Self {
         Self {
             config,
@@ -207,6 +211,7 @@ impl Session {
             branch_name: Some(branch_name),
             snapshot_id,
             change_set: ChangeSet::default(),
+            default_commit_metadata,
         }
     }
 
@@ -806,6 +811,17 @@ impl Session {
             return Err(SessionErrorKind::ReadOnlySession.into());
         };
 
+        let properties = match (properties, self.default_commit_metadata.as_ref()) {
+            (Some(p), None) => Some(p),
+            (None, Some(d)) => Some(d.clone()),
+            (Some(p), Some(d)) => {
+                let mut merged = d.clone();
+                merged.extend(p.into_iter());
+                Some(merged)
+            }
+            (None, None) => None,
+        };
+
         let current = fetch_branch_tip(
             self.storage.as_ref(),
             self.storage_settings.as_ref(),
@@ -1950,6 +1966,52 @@ mod tests {
         prop_assert_eq!(to, expected_to);
     }
 
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_repository_with_default_commit_metadata() -> Result<(), Box<dyn Error>>
+    {
+        let repo = create_memory_store_repository().await;
+        let mut ds = repo.writable_session("main").await?;
+        ds.add_group(Path::root(), Bytes::new()).await?;
+        let snapshot = ds.commit("commit", None).await?;
+
+        // Verify that the first commit has no metadata
+        let ancestry = repo.snapshot_ancestry(&snapshot).await?;
+        let snapshot_infos = ancestry.try_collect::<Vec<_>>().await?;
+        assert!(snapshot_infos[0].metadata.is_empty());
+
+        // Set some default metadata
+        let mut default_metadata = SnapshotProperties::default();
+        default_metadata.insert("author".to_string(), "John Doe".to_string().into());
+        default_metadata.insert("project".to_string(), "My Project".to_string().into());
+        repo.set_default_commit_metadata(Some(default_metadata.clone())).await;
+
+        let mut ds = repo.writable_session("main").await?;
+        ds.add_group("/group".try_into().unwrap(), Bytes::new()).await?;
+        let snapshot = ds.commit("commit", None).await?;
+
+        let snapshot_info = repo.snapshot_ancestry(&snapshot).await?;
+        let snapshot_infos = snapshot_info.try_collect::<Vec<_>>().await?;
+        assert_eq!(snapshot_infos[0].metadata, default_metadata);
+
+        // Check that metadata is merged with users provided metadata taking precedence
+        let mut metadata = SnapshotProperties::default();
+        metadata.insert("author".to_string(), "Jane Doe".to_string().into());
+        metadata.insert("id".to_string(), "ideded".to_string().into());
+        let mut ds = repo.writable_session("main").await?;
+        ds.add_group("/group2".try_into().unwrap(), Bytes::new()).await?;
+        let snapshot = ds.commit("commit", Some(metadata.clone())).await?;
+
+        let snapshot_info = repo.snapshot_ancestry(&snapshot).await?;
+        let snapshot_infos = snapshot_info.try_collect::<Vec<_>>().await?;
+        let mut expected_result = SnapshotProperties::default();
+        expected_result.insert("author".to_string(), "Jane Doe".to_string().into());
+        expected_result.insert("project".to_string(), "My Project".to_string().into());
+        expected_result.insert("id".to_string(), "ideded".to_string().into());
+        assert_eq!(snapshot_infos[0].metadata, expected_result);
+
+        Ok(())
+    }
+
     #[tokio::test(flavor = "multi_thread")]
     async fn test_repository_with_updates() -> Result<(), Box<dyn Error>> {
         let storage: Arc<dyn Storage + Send + Sync> = new_in_memory_storage().await?;