Commit learns how to rebase in case of conflict

`Session.commit` accepts new arguments:

```
rebase_with: ConflictSolver | None = None,
rebase_tries: int = 5
```

It uses the `ConflictSolver` to call `Session.rebase` in a loop, up to
`rebase_tries` times, and try to commit after each rebase.

Author: paraseba

icechunk-python/examples/bank_accounts.py

@@ -138,15 +138,6 @@ def slow_transfer_task(repo: icechunk.Repository) -> WireResult:
     return res
 
 
-def rebase_loop(session: icechunk.Session, message: str) -> None:
-    while True:
-        try:
-            session.commit(message)
-            return
-        except icechunk.ConflictError:
-            session.rebase(icechunk.ConflictDetector())
-
-
 def rebase_transfer_task(repo: icechunk.Repository) -> WireResult:
     """Safe and fast approach to concurrent transfers.
 
@@ -169,7 +160,10 @@ def rebase_transfer_task(repo: icechunk.Repository) -> WireResult:
             return res
         if res == WireResult.DONE:
             try:
-                rebase_loop(session, f"wired ${amount}: {from_account} -> {to_account}")
+                session.commit(
+                    f"wired ${amount}: {from_account} -> {to_account}",
+                    rebase_with=icechunk.ConflictDetector(),
+                )
                 return WireResult.DONE
             except icechunk.RebaseFailedError:
                 pass

icechunk-python/python/icechunk/_icechunk_python.pyi

@@ -1066,7 +1066,13 @@ class PySession:
     @property
     def store(self) -> PyStore: ...
     def merge(self, other: PySession) -> None: ...
-    def commit(self, message: str, metadata: dict[str, Any] | None = None) -> str: ...
+    def commit(
+        self,
+        message: str,
+        metadata: dict[str, Any] | None = None,
+        rebase_with: ConflictSolver | None = None,
+        rebase_tries: int = 5,
+    ) -> str: ...
     def rebase(self, solver: ConflictSolver) -> None: ...
 
 class PyStore:

icechunk-python/python/icechunk/session.py

@@ -250,7 +250,13 @@ def merge(self, other: Self) -> None:
         """
         self._session.merge(other._session)
 
-    def commit(self, message: str, metadata: dict[str, Any] | None = None) -> str:
+    def commit(
+        self,
+        message: str,
+        metadata: dict[str, Any] | None = None,
+        rebase_with: ConflictSolver | None = None,
+        rebase_tries: int = 5,
+    ) -> str:
         """
         Commit the changes in the session to the repository.
 
@@ -264,6 +270,10 @@ def commit(self, message: str, metadata: dict[str, Any] | None = None) -> str:
             The message to write with the commit.
         metadata : dict[str, Any] | None, optional
             Additional metadata to store with the commit snapshot.
+        rebase_with : ConflictSolver | None, optional
+            If other session committed while the current session was writing, use Session.rebase with this solver.
+        rebase_tries : int, optional
+            If other session committed while the current session was writing, use Session.rebase up to this many times in a loop.
 
         Returns
         -------
@@ -276,9 +286,13 @@ def commit(self, message: str, metadata: dict[str, Any] | None = None) -> str:
             If the session is out of date and a conflict occurs.
         """
         try:
-            return self._session.commit(message, metadata)
+            return self._session.commit(
+                message, metadata, rebase_with=rebase_with, rebase_tries=rebase_tries
+            )
         except PyConflictError as e:
             raise ConflictError(e) from None
+        except PyRebaseFailedError as e:
+            raise RebaseFailedError(e) from None
 
     def rebase(self, solver: ConflictSolver) -> None:
         """

icechunk-python/src/session.rs

@@ -180,23 +180,35 @@ impl PySession {
         })
     }
 
-    #[pyo3(signature = (message, metadata=None))]
+    #[pyo3(signature = (message, metadata=None, rebase_with=None, rebase_tries=5))]
     pub fn commit(
         &self,
         py: Python<'_>,
         message: &str,
         metadata: Option<PySnapshotProperties>,
+        rebase_with: Option<PyConflictSolver>,
+        rebase_tries: Option<u16>,
     ) -> PyResult<String> {
+        let metadata = metadata.map(|m| m.into());
         // This is blocking function, we need to release the Gil
         py.allow_threads(move || {
             pyo3_async_runtimes::tokio::get_runtime().block_on(async {
-                let snapshot_id = self
-                    .0
-                    .write()
-                    .await
-                    .commit(message, metadata.map(|m| m.into()))
-                    .await
-                    .map_err(PyIcechunkStoreError::SessionError)?;
+                let mut session = self.0.write().await;
+                let snapshot_id = if let Some(solver) = rebase_with {
+                    session
+                        .commit_rebasing(
+                            solver.as_ref(),
+                            rebase_tries.unwrap_or(5),
+                            message,
+                            metadata,
+                            |_| async {},
+                            |_| async {},
+                        )
+                        .await
+                } else {
+                    session.commit(message, metadata).await
+                }
+                .map_err(PyIcechunkStoreError::SessionError)?;
                 Ok(snapshot_id.to_string())
             })
         })

icechunk-python/tests/test_conflicts.py

@@ -83,8 +83,7 @@ def test_rebase_no_conflicts(repo: icechunk.Repository) -> None:
     array_b = cast(zarr.Array, root_b["foo/bar/some-array"])
     array_b.attrs["repo"] = 2
 
-    session_b.rebase(icechunk.ConflictDetector())
-    session_b.commit("update array")
+    session_b.commit("update array", rebase_with=icechunk.ConflictDetector())
 
     session_c = repo.readonly_session(branch="main")
     store_c = session_c.store
@@ -94,13 +93,7 @@ def test_rebase_no_conflicts(repo: icechunk.Repository) -> None:
     assert array_c.attrs["repo"] == 2
 
 
-@pytest.mark.parametrize(
-    "on_chunk_conflict",
-    [icechunk.VersionSelection.UseOurs, icechunk.VersionSelection.UseTheirs],
-)
-def test_rebase_fails_on_user_atts_double_edit(
-    repo: icechunk.Repository, on_chunk_conflict: icechunk.VersionSelection
-) -> None:
+def test_rebase_fails_on_user_atts_double_edit(repo: icechunk.Repository) -> None:
     session_a = repo.writable_session("main")
     session_b = repo.writable_session("main")
     store_a = session_a.store
@@ -120,7 +113,7 @@ def test_rebase_fails_on_user_atts_double_edit(
 
     # Make sure it fails if the resolver is not set
     with pytest.raises(icechunk.RebaseFailedError):
-        session_b.rebase(icechunk.BasicConflictSolver())
+        session_b.commit("update array", rebase_with=icechunk.BasicConflictSolver())
 
 
 @pytest.mark.parametrize(
@@ -150,10 +143,11 @@ def test_rebase_chunks_with_ours(
     # Make sure it fails if the resolver is not set
     with pytest.raises(icechunk.RebaseFailedError):
         try:
-            session_b.rebase(
-                icechunk.BasicConflictSolver(
+            session_b.commit(
+                "update first column of array",
+                rebase_with=icechunk.BasicConflictSolver(
                     on_chunk_conflict=icechunk.VersionSelection.Fail
-                )
+                ),
             )
         except icechunk.RebaseFailedError as e:
             assert e.conflicts[0].path == "/foo/bar/some-array"
@@ -184,8 +178,7 @@ def test_rebase_chunks_with_ours(
         on_chunk_conflict=on_chunk_conflict,
     )
 
-    session_b.rebase(solver)
-    session_b.commit("after conflict")
+    session_b.commit("after conflict", rebase_with=solver)
 
     session_c = repo.readonly_session(branch="main")
     store_c = session_c.store

icechunk/src/session.rs

@@ -877,6 +877,38 @@ impl Session {
         Ok(id)
     }
 
+    pub async fn commit_rebasing<F1, F2, Fut1, Fut2>(
+        &mut self,
+        solver: &dyn ConflictSolver,
+        rebase_attempts: u16,
+        message: &str,
+        properties: Option<SnapshotProperties>,
+        // We would prefer to make this argument optional, but passing None
+        // for this argument is so hard. Callers should just pass noop closure like
+        // |_| async {},
+        before_rebase: F1,
+        after_rebase: F2,
+    ) -> SessionResult<SnapshotId>
+    where
+        F1: Fn(u16) -> Fut1,
+        F2: Fn(u16) -> Fut2,
+        Fut1: Future<Output = ()>,
+        Fut2: Future<Output = ()>,
+    {
+        for attempt in 0..rebase_attempts {
+            match self.commit(message, properties.clone()).await {
+                Ok(snap) => return Ok(snap),
+                Err(SessionError { kind: SessionErrorKind::Conflict { .. }, .. }) => {
+                    before_rebase(attempt + 1).await;
+                    self.rebase(solver).await?;
+                    after_rebase(attempt + 1).await;
+                }
+                Err(other_err) => return Err(other_err),
+            }
+        }
+        self.commit(message, properties).await
+    }
+
     /// Detect and optionally fix conflicts between the current [`ChangeSet`] (or session) and
     /// the tip of the branch.
     ///
@@ -1754,7 +1786,11 @@ fn aggregate_extents<'a, T: std::fmt::Debug, E>(
 #[cfg(test)]
 #[allow(clippy::panic, clippy::unwrap_used, clippy::expect_used)]
 mod tests {
-    use std::{collections::HashMap, error::Error};
+    use std::{
+        collections::HashMap,
+        error::Error,
+        sync::atomic::{AtomicU16, Ordering},
+    };
 
     use crate::{
         ObjectStorage, Repository,
@@ -3752,6 +3788,122 @@ mod tests {
         Ok(())
     }
 
+    #[tokio_test]
+    /// Tests `commit_rebasing` retries the proper number of times when there are conflicts
+    async fn test_commit_rebasing_attempts() -> Result<(), Box<dyn Error>> {
+        let repo = Arc::new(create_memory_store_repository().await);
+        let mut session = repo.writable_session("main").await?;
+        session
+            .add_array("/array".try_into().unwrap(), basic_shape(), None, Bytes::new())
+            .await?;
+        session.commit("create array", None).await?;
+
+        // This is the main session we'll be trying to commit (and rebase)
+        let mut session = repo.writable_session("main").await?;
+        let path: Path = "/array".try_into().unwrap();
+        session
+            .set_chunk_ref(
+                path.clone(),
+                ChunkIndices(vec![1]),
+                Some(ChunkPayload::Inline("repo 1".into())),
+            )
+            .await?;
+
+        // we create an initial conflict for commit
+        let mut session2 = repo.writable_session("main").await.unwrap();
+        let path: Path = "/array".try_into().unwrap();
+        session2
+            .set_chunk_ref(
+                path.clone(),
+                ChunkIndices(vec![2]),
+                Some(ChunkPayload::Inline("repo 1".into())),
+            )
+            .await
+            .unwrap();
+        session2.commit("conflicting", None).await.unwrap();
+
+        let repo_ref = &repo;
+        let attempts = AtomicU16::new(0);
+        let attempts_ref = &attempts;
+
+        // after each rebase attempt we'll run this closure that creates a new conflict
+        // the result should be that it can never commit, failing after the indicated number of
+        // attempts
+        let conflicting = |attempt| async move {
+            attempts_ref.fetch_add(1, Ordering::SeqCst); //*attempts_ref = *attempts_ref + 1;;
+            assert_eq!(attempt, attempts_ref.load(Ordering::SeqCst));
+
+            let repo_c = Arc::clone(repo_ref);
+            let mut s = repo_c.writable_session("main").await.unwrap();
+            s.set_chunk_ref(
+                "/array".try_into().unwrap(),
+                ChunkIndices(vec![2]),
+                Some(ChunkPayload::Inline("repo 1".into())),
+            )
+            .await
+            .unwrap();
+            s.commit("conflicting", None).await.unwrap();
+        };
+
+        let res = session
+            .commit_rebasing(
+                &ConflictDetector,
+                3,
+                "updated non-conflict chunk",
+                None,
+                |_| async {},
+                conflicting,
+            )
+            .await;
+
+        // It has to give up eventually
+        assert!(matches!(
+            res,
+            Err(SessionError { kind: SessionErrorKind::Conflict { .. }, .. })
+        ));
+
+        // It has to rebase 3 times
+        assert_eq!(attempts.into_inner(), 3);
+
+        let attempts = AtomicU16::new(0);
+        let attempts_ref = &attempts;
+
+        // now we'll create a new conflict twice, and finally do nothing so the commit can succeed
+        let conflicting_twice = |attempt| async move {
+            attempts_ref.fetch_add(1, Ordering::SeqCst); //*attempts_ref = *attempts_ref + 1;;
+            assert_eq!(attempt, attempts_ref.load(Ordering::SeqCst));
+            if attempt <= 2 {
+                let repo_c = Arc::clone(repo_ref);
+
+                let mut s = repo_c.writable_session("main").await.unwrap();
+                s.set_chunk_ref(
+                    "/array".try_into().unwrap(),
+                    ChunkIndices(vec![2]),
+                    Some(ChunkPayload::Inline("repo 1".into())),
+                )
+                .await
+                .unwrap();
+                s.commit("conflicting", None).await.unwrap();
+            }
+        };
+
+        let res = session
+            .commit_rebasing(
+                &ConflictDetector,
+                42,
+                "updated non-conflict chunk",
+                None,
+                |_| async {},
+                conflicting_twice,
+            )
+            .await;
+
+        // The commit has to work after 3 rebase attempts
+        assert!(res.is_ok());
+        assert_eq!(attempts.into_inner(), 3);
+        Ok(())
+    }
+
     #[cfg(test)]
     mod state_machine_test {
         use crate::format::Path;