Safer reset_branch (#1259)

An optional third argument was added to `reset_branch`:

```python
def reset_branch(self, branch: str, snapshot_id: str, *, from_snapshot_id: str | None = None) -> None:
```

If `from_snapshot_id` is passed, the operation raises an exception and
leaves the branch untouched if it's not currently pointing to that
snapshot id.

Author: paraseba

docs/docs/icechunk-for-git-users.md

@@ -71,10 +71,22 @@ In git, you can reset a branch to previous commit. Similarly, in Icechunk you ca
 repo.reset_branch("my-new-branch", "198273178639187")
 ```
 
+At this point, the tip of the branch is now the snapshot `198273178639187` and any changes made to the branch will be based on this snapshot. This also means the history of the branch is now same as the ancestry of this snapshot.
+
 !!! warning
     This is a destructive operation. It will overwrite the branch reference with the snapshot immediately. It can only be undone by resetting the branch again.
 
-At this point, the tip of the branch is now the snapshot `198273178639187` and any changes made to the branch will be based on this snapshot. This also means the history of the branch is now same as the ancestry of this snapshot.
+To make `reset_branch` less dangerous, you can pass an optional third argument, the snapshot id currently pointed
+by the branch:
+
+```python
+# reset the branch to the initial commit
+repo.reset_branch("my-new-branch", "1CECHNKREP0F1RSTCMT0", from_snapshot_id="198273178639187")
+```
+
+If the branch is not currently pointing to snapshot `198273178639187`, the operation will be rejected
+with an exception. Using this approach, you can make sure to only reset a branch if no ether commits
+were done to it since you last checked.
 
 ### Branch History
 

icechunk-python/python/icechunk/_icechunk_python.pyi

@@ -1395,8 +1395,12 @@ class PyRepository:
     async def lookup_branch_async(self, branch: str) -> str: ...
     def lookup_snapshot(self, snapshot_id: str) -> SnapshotInfo: ...
     async def lookup_snapshot_async(self, snapshot_id: str) -> SnapshotInfo: ...
-    def reset_branch(self, branch: str, snapshot_id: str) -> None: ...
-    async def reset_branch_async(self, branch: str, snapshot_id: str) -> None: ...
+    def reset_branch(
+        self, branch: str, to_snapshot_id: str, from_snapshot_id: str | None
+    ) -> None: ...
+    async def reset_branch_async(
+        self, branch: str, to_snapshot_id: str, from_snapshot_id: str | None
+    ) -> None: ...
     def delete_branch(self, branch: str) -> None: ...
     async def delete_branch_async(self, branch: str) -> None: ...
     def delete_tag(self, tag: str) -> None: ...

icechunk-python/python/icechunk/repository.py

@@ -681,7 +681,9 @@ async def lookup_snapshot_async(self, snapshot_id: str) -> SnapshotInfo:
         """
         return await self._repository.lookup_snapshot_async(snapshot_id)
 
-    def reset_branch(self, branch: str, snapshot_id: str) -> None:
+    def reset_branch(
+        self, branch: str, snapshot_id: str, *, from_snapshot_id: str | None = None
+    ) -> None:
         """
         Reset a branch to a specific snapshot.
 
@@ -694,14 +696,19 @@ def reset_branch(self, branch: str, snapshot_id: str) -> None:
             The branch to reset.
         snapshot_id : str
             The snapshot ID to reset the branch to.
+        from_snapshot_id : str | None
+            If passed, the reset will only be executed if the branch currently
+            points to from_snapshot_id.
 
         Returns
         -------
         None
         """
-        self._repository.reset_branch(branch, snapshot_id)
+        self._repository.reset_branch(branch, snapshot_id, from_snapshot_id)
 
-    async def reset_branch_async(self, branch: str, snapshot_id: str) -> None:
+    async def reset_branch_async(
+        self, branch: str, snapshot_id: str, *, from_snapshot_id: str | None = None
+    ) -> None:
         """
         Reset a branch to a specific snapshot (async version).
 
@@ -714,12 +721,15 @@ async def reset_branch_async(self, branch: str, snapshot_id: str) -> None:
             The branch to reset.
         snapshot_id : str
             The snapshot ID to reset the branch to.
+        from_snapshot_id : str | None
+            If passed, the reset will only be executed if the branch currently
+            points to from_snapshot_id.
 
         Returns
         -------
         None
         """
-        await self._repository.reset_branch_async(branch, snapshot_id)
+        await self._repository.reset_branch_async(branch, snapshot_id, from_snapshot_id)
 
     def delete_branch(self, branch: str) -> None:
         """

icechunk-python/src/repository.rs

@@ -973,21 +973,33 @@ impl PyRepository {
         &self,
         py: Python<'_>,
         branch_name: &str,
-        snapshot_id: &str,
+        to_snapshot_id: &str,
+        from_snapshot_id: Option<&str>,
     ) -> PyResult<()> {
         // This function calls block_on, so we need to allow other thread python to make progress
         py.allow_threads(move || {
-            let snapshot_id = SnapshotId::try_from(snapshot_id).map_err(|_| {
+            let to_snapshot_id = SnapshotId::try_from(to_snapshot_id).map_err(|_| {
                 PyIcechunkStoreError::RepositoryError(
-                    RepositoryErrorKind::InvalidSnapshotId(snapshot_id.to_owned()).into(),
+                    RepositoryErrorKind::InvalidSnapshotId(to_snapshot_id.to_owned())
+                        .into(),
                 )
             })?;
 
+            let from_snapshot_id = from_snapshot_id
+                .map(|sid| {
+                    SnapshotId::try_from(sid).map_err(|_| {
+                        PyIcechunkStoreError::RepositoryError(
+                            RepositoryErrorKind::InvalidSnapshotId(sid.to_owned()).into(),
+                        )
+                    })
+                })
+                .transpose()?;
+
             pyo3_async_runtimes::tokio::get_runtime().block_on(async move {
                 self.0
                     .read()
                     .await
-                    .reset_branch(branch_name, &snapshot_id)
+                    .reset_branch(branch_name, &to_snapshot_id, from_snapshot_id.as_ref())
                     .await
                     .map_err(PyIcechunkStoreError::RepositoryError)?;
                 Ok(())
@@ -999,20 +1011,31 @@ impl PyRepository {
         &'py self,
         py: Python<'py>,
         branch_name: &str,
-        snapshot_id: &str,
+        to_snapshot_id: &str,
+        from_snapshot_id: Option<&str>,
     ) -> PyResult<Bound<'py, PyAny>> {
         let repository = self.0.clone();
         let branch_name = branch_name.to_owned();
-        let snapshot_id = SnapshotId::try_from(snapshot_id).map_err(|_| {
+        let to_snapshot_id = SnapshotId::try_from(to_snapshot_id).map_err(|_| {
             PyIcechunkStoreError::RepositoryError(
-                RepositoryErrorKind::InvalidSnapshotId(snapshot_id.to_owned()).into(),
+                RepositoryErrorKind::InvalidSnapshotId(to_snapshot_id.to_owned()).into(),
             )
         })?;
 
+        let from_snapshot_id = from_snapshot_id
+            .map(|sid| {
+                SnapshotId::try_from(sid).map_err(|_| {
+                    PyIcechunkStoreError::RepositoryError(
+                        RepositoryErrorKind::InvalidSnapshotId(sid.to_owned()).into(),
+                    )
+                })
+            })
+            .transpose()?;
+
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
             let repository = repository.read().await;
             repository
-                .reset_branch(&branch_name, &snapshot_id)
+                .reset_branch(&branch_name, &to_snapshot_id, from_snapshot_id.as_ref())
                 .await
                 .map_err(PyIcechunkStoreError::RepositoryError)?;
             Ok(())

icechunk-python/tests/test_timetravel.py

@@ -232,13 +232,21 @@ async def test_branch_reset() -> None:
 
     group = zarr.open_group(store=store)
     group.create_group("b")
-    session.commit("group b")
+    last_commit = session.commit("group b")
 
     keys = {k async for k in store.list()}
     assert "a/zarr.json" in keys
     assert "b/zarr.json" in keys
 
-    repo.reset_branch("main", prev_snapshot_id)
+    with pytest.raises(ic.IcechunkError, match="branch update conflict"):
+        repo.reset_branch(
+            "main", prev_snapshot_id, from_snapshot_id="1CECHNKREP0F1RSTCMT0"
+        )
+
+    assert last_commit == repo.lookup_branch("main")
+
+    repo.reset_branch("main", prev_snapshot_id, from_snapshot_id=last_commit)
+    assert prev_snapshot_id == repo.lookup_branch("main")
 
     session = repo.readonly_session("main")
     store = session.store

icechunk/src/repository.rs

@@ -574,7 +574,8 @@ impl Repository {
     pub async fn reset_branch(
         &self,
         branch: &str,
-        snapshot_id: &SnapshotId,
+        to_snapshot_id: &SnapshotId,
+        from_snapshot_id: Option<&SnapshotId>,
     ) -> RepositoryResult<()> {
         if !self.storage.can_write() {
             return Err(RepositoryErrorKind::ReadonlyStorage(
@@ -585,16 +586,19 @@ impl Repository {
         raise_if_invalid_snapshot_id(
             self.storage.as_ref(),
             &self.storage_settings,
-            snapshot_id,
+            to_snapshot_id,
         )
         .await?;
-        let branch_tip = self.lookup_branch(branch).await?;
+        let branch_tip = match from_snapshot_id {
+            Some(snap) => snap,
+            None => &self.lookup_branch(branch).await?,
+        };
         update_branch(
             self.storage.as_ref(),
             &self.storage_settings,
             branch,
-            snapshot_id.clone(),
-            Some(&branch_tip),
+            to_snapshot_id.clone(),
+            Some(branch_tip),
         )
         .await
         .err_into()

icechunk/src/store.rs

@@ -2348,7 +2348,7 @@ mod tests {
         assert!(store.exists("a/zarr.json").await?);
         assert!(store.exists("b/zarr.json").await?);
 
-        repo.reset_branch("main", &prev_snap).await?;
+        repo.reset_branch("main", &prev_snap, None).await?;
         let ds = Arc::new(RwLock::new(
             repo.readonly_session(&VersionInfo::BranchTipRef("main".to_string())).await?,
         ));