Commit learns how to rebase in case of conflict (#920)

`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

docs/docs/icechunk-python/version-control.md

@@ -72,7 +72,6 @@ print(session.commit(message="Update foo attribute on root group"))
 
 With a few snapshots committed, we can take a look at the ancestry of the `main` branch:
 
-
 ```python exec="on" session="version" source="material-block" result="code"
 for snapshot in repo.ancestry(branch="main"):
     print(snapshot)
@@ -89,7 +88,6 @@ gitGraph
 """.format(*[snap.id[:6] for snap in repo.ancestry(branch="main")]))
 ```
 
-
 ## Time Travel
 
 Now that we've created a new snapshot, we can time-travel back to the previous snapshot using the snapshot ID.
@@ -193,7 +191,6 @@ For example to tag the second commit in `main`'s history:
 repo.create_tag("v1.0.0", snapshot_id=list(repo.ancestry(branch="main"))[1].id)
 ```
 
-
 Because tags are immutable, we need to use a readonly `Session` to access the data referenced by a tag.
 
 ```python exec="on" session="version" source="material-block" result="code"
@@ -353,6 +350,7 @@ session = repo.readonly_session(branch="main")
 ```
 
 Lastly, if you make changes to non-conflicting chunks or attributes, you can rebase without having to resolve any conflicts.
+This time we will show how to use rebase automatically during the `commit` call:
 
 ```python
 session1 = repo.writable_session("main")
@@ -365,16 +363,9 @@ root1["data"][3,:] = 3
 root2["data"][4,:] = 4
 
 session1.commit(message="Update fourth row of data array")
+session2.commit(message="Update fifth row of data array", rebase_with=icechunk.ConflictDetector())
+print("Rebase+commit succeeded")
 
-try:
-    session2.rebase(icechunk.ConflictDetector())
-    print("Rebase succeeded")
-except icechunk.RebaseFailedError as e:
-    print(e.conflicts)
-
-session2.commit(message="Update fifth row of data array")
-
-# Rebase succeeded
 ```
 
 And now we can see the data in the `data` array to confirm that the changes were committed correctly.

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 = 1_000,
+    ) -> 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 = 1_000,
+    ) -> 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=1_000))]
     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(1_000),
+                            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