Add ManageSnapshots.fast_forward_branch

mkdocs/docs/api.md

@@ -1483,6 +1483,31 @@ Remove an existing branch:
 table.manage_snapshots().remove_branch("dev").commit()
 ```
 
+#### Fast-forwarding a branch
+
+Advance a branch to the snapshot another ref points to. `to_ref` is left untouched and may be a branch or a tag, while `from_branch` must be a branch. The fast-forward only succeeds when the snapshot of `from_branch` is an ancestor of the snapshot of `to_ref`, otherwise a `ValueError` is raised. Fast-forwarding a branch onto the snapshot it already points to does nothing.
+
+A `from_branch` that does not exist yet is created at the snapshot of `to_ref` with the default retention properties, so a misspelled branch name creates a new branch instead of raising. An existing `from_branch` keeps its own retention properties.
+
+This is the publish step of the write-audit-publish pattern, where data is written to a side branch, validated there, and only then made visible on `main`:
+
+```python
+# Write: branch off main and append to the branch
+main_snapshot_id = table.refs()["main"].snapshot_id
+table.manage_snapshots().create_branch(snapshot_id=main_snapshot_id, branch_name="audit").commit()
+table.append(df, branch="audit")
+
+# Audit: main is untouched while the new data is validated on the branch
+assert table.refs()["main"].snapshot_id == main_snapshot_id
+audit_data = table.scan(snapshot_id=table.refs()["audit"].snapshot_id).to_arrow()
+assert audit_data.num_rows > 0
+
+# Publish: move main up to the audited snapshot
+with table.manage_snapshots() as ms:
+    ms.fast_forward_branch("main", "audit")
+    ms.remove_branch("audit")
+```
+
 ## Table Maintenance
 
 PyIceberg provides table maintenance operations through the `table.maintenance` API. This provides a clean interface for performing maintenance tasks like snapshot expiration.

pyiceberg/table/update/__init__.py

@@ -138,9 +138,9 @@ class SetSnapshotRefUpdate(IcebergBaseModel):
     ref_name: str = Field(alias="ref-name")
     type: Literal[SnapshotRefType.TAG, SnapshotRefType.BRANCH]
     snapshot_id: int = Field(alias="snapshot-id")
-    max_ref_age_ms: Annotated[int | None, Field(alias="max-ref-age-ms", default=None)]
-    max_snapshot_age_ms: Annotated[int | None, Field(alias="max-snapshot-age-ms", default=None)]
-    min_snapshots_to_keep: Annotated[int | None, Field(alias="min-snapshots-to-keep", default=None)]
+    max_ref_age_ms: Annotated[int | None, Field(alias="max-ref-age-ms", default=None, gt=0)]
+    max_snapshot_age_ms: Annotated[int | None, Field(alias="max-snapshot-age-ms", default=None, gt=0)]
+    min_snapshots_to_keep: Annotated[int | None, Field(alias="min-snapshots-to-keep", default=None, gt=0)]
 
 
 class RemoveSnapshotsUpdate(IcebergBaseModel):

pyiceberg/table/update/snapshot.py

@@ -59,6 +59,7 @@
     SnapshotSummaryCollector,
     Summary,
     ancestors_of,
+    is_ancestor_of,
     latest_ancestor_before_timestamp,
     update_snapshot_summaries,
 )
@@ -73,6 +74,7 @@
     U,
     UpdatesAndRequirements,
     UpdateTableMetadata,
+    update_table_metadata,
 )
 from pyiceberg.typedef import EMPTY_DICT, KeyDefaultDict, Record
 from pyiceberg.utils.bin_packing import ListPacker
@@ -1043,6 +1045,14 @@ def _commit_if_ref_updates_exist(self) -> None:
             self._updates = ()
             self._requirements = ()
 
+    def _pending_table_metadata(self) -> TableMetadata:
+        """Return the table metadata with the ref updates staged on this instance applied.
+
+        Only used to look at the result of operations chained earlier, requirements are still built
+        from the transaction metadata.
+        """
+        return update_table_metadata(self._transaction.table_metadata, self._updates)
+
     def _remove_ref_snapshot(self, ref_name: str) -> ManageSnapshots:
         """Remove a snapshot ref.
 
@@ -1192,6 +1202,8 @@ def rollback_to_snapshot(self, snapshot_id: int) -> ManageSnapshots:
         Raises:
             ValueError: If the snapshot does not exist or is not an ancestor of the current table state.
         """
+        self._commit_if_ref_updates_exist()
+
         if not self._transaction.table_metadata.snapshot_by_id(snapshot_id):
             raise ValueError(f"Cannot roll back to unknown snapshot id: {snapshot_id}")
 
@@ -1214,6 +1226,8 @@ def rollback_to_timestamp(self, timestamp_ms: int) -> ManageSnapshots:
         Raises:
             ValueError: If no valid snapshot exists older than the given timestamp.
         """
+        self._commit_if_ref_updates_exist()
+
         snapshot = latest_ancestor_before_timestamp(self._transaction.table_metadata, timestamp_ms)
         if snapshot is None:
             raise ValueError(f"Cannot roll back, no valid snapshot older than: {timestamp_ms}")
@@ -1232,6 +1246,79 @@ def _current_ancestors(self) -> set[int]:
             )
         }
 
+    def _stage_ref_snapshot(
+        self,
+        ref_name: str,
+        snapshot_id: int,
+        type: str,
+        max_ref_age_ms: int | None = None,
+        max_snapshot_age_ms: int | None = None,
+        min_snapshots_to_keep: int | None = None,
+    ) -> None:
+        """Stage a set-snapshot-ref, replacing any update already pending for the same ref."""
+        update, requirement = self._transaction._set_ref_snapshot(
+            snapshot_id=snapshot_id,
+            ref_name=ref_name,
+            type=type,
+            max_ref_age_ms=max_ref_age_ms,
+            max_snapshot_age_ms=max_snapshot_age_ms,
+            min_snapshots_to_keep=min_snapshots_to_keep,
+        )
+        self._updates = (
+            tuple(u for u in self._updates if not (isinstance(u, SetSnapshotRefUpdate) and u.ref_name == ref_name)) + update
+        )
+        if not any(isinstance(r, AssertRefSnapshotId) and r.ref == ref_name for r in self._requirements):
+            self._requirements += requirement
+
+    def fast_forward_branch(self, from_branch: str, to_ref: str) -> ManageSnapshots:
+        """
+        Fast-forward a branch to the snapshot that another ref points to.
+
+        The ref to fast-forward to can be a branch or a tag and is left untouched. A branch that does
+        not exist yet is created at that snapshot with default retention properties, otherwise the
+        branch keeps its retention properties. Fast-forwarding a branch to the snapshot it already
+        points to is a no-op.
+
+        Args:
+            from_branch (str): name of the branch to fast-forward
+            to_ref (str): name of the branch or tag to fast-forward to
+        Returns:
+            This for method chaining
+        Raises:
+            ValueError: If to_ref does not exist, from_branch is a tag, or from_branch is not an ancestor of to_ref.
+        """
+        table_metadata = self._pending_table_metadata()
+        refs = table_metadata.refs
+
+        if to_ref not in refs:
+            raise ValueError(f"Ref does not exist: {to_ref}")
+
+        to_snapshot_id = refs[to_ref].snapshot_id
+        if table_metadata.snapshot_by_id(to_snapshot_id) is None:
+            raise ValueError(f"Cannot fast-forward to unknown snapshot id: {to_snapshot_id}")
+
+        if (from_ref := refs.get(from_branch)) is None:
+            return self.create_branch(snapshot_id=to_snapshot_id, branch_name=from_branch)
+
+        if from_ref.snapshot_ref_type != SnapshotRefType.BRANCH:
+            raise ValueError(f"Ref {from_branch} is a tag not a branch")
+
+        if from_ref.snapshot_id == to_snapshot_id:
+            return self
+
+        if not is_ancestor_of(to_snapshot_id, from_ref.snapshot_id, table_metadata):
+            raise ValueError(f"Cannot fast-forward: {from_branch} is not an ancestor of {to_ref}")
+
+        self._stage_ref_snapshot(
+            ref_name=from_branch,
+            snapshot_id=to_snapshot_id,
+            type=SnapshotRefType.BRANCH,
+            max_ref_age_ms=from_ref.max_ref_age_ms,
+            max_snapshot_age_ms=from_ref.max_snapshot_age_ms,
+            min_snapshots_to_keep=from_ref.min_snapshots_to_keep,
+        )
+        return self
+
 
 class ExpireSnapshots(UpdateTableMetadata["ExpireSnapshots"]):
     """Expire snapshots by ID.

tests/integration/test_snapshot_operations.py

@@ -332,3 +332,69 @@ def test_rollback_to_timestamp_chained_with_tag(table_with_snapshots: Table) ->
     assert table_with_snapshots.metadata.refs[tag_name] == SnapshotRef(
         snapshot_id=current_snapshot.snapshot_id, snapshot_ref_type="tag"
     )
+
+
+@pytest.mark.integration
+def test_fast_forward_branch(table_with_snapshots: Table) -> None:
+    main_snapshot_id = table_with_snapshots.refs()["main"].snapshot_id
+    history_length = len(table_with_snapshots.history())
+
+    table_with_snapshots.manage_snapshots().create_branch(snapshot_id=main_snapshot_id, branch_name="audit").commit()
+
+    arrow_schema = table_with_snapshots.schema().as_arrow()
+    table_with_snapshots.append(
+        pa.Table.from_pylist([{"id": 5, "data": "e"}], schema=arrow_schema),
+        branch="audit",
+    )
+    audit_snapshot_id = table_with_snapshots.refs()["audit"].snapshot_id
+    assert audit_snapshot_id != main_snapshot_id
+    assert table_with_snapshots.refs()["main"].snapshot_id == main_snapshot_id
+
+    table_with_snapshots.manage_snapshots().fast_forward_branch(from_branch="main", to_ref="audit").commit()
+
+    assert table_with_snapshots.refs()["main"].snapshot_id == audit_snapshot_id
+    assert table_with_snapshots.refs()["audit"].snapshot_id == audit_snapshot_id
+    current_snapshot = table_with_snapshots.current_snapshot()
+    assert current_snapshot is not None
+    assert current_snapshot.snapshot_id == audit_snapshot_id
+    assert len(table_with_snapshots.history()) == history_length + 1
+
+
+@pytest.mark.integration
+def test_fast_forward_branch_preserves_retention(table_with_snapshots: Table) -> None:
+    main_snapshot_id = table_with_snapshots.refs()["main"].snapshot_id
+    older_snapshot_id = table_with_snapshots.history()[-2].snapshot_id
+    assert older_snapshot_id != main_snapshot_id
+
+    # distinct values so a swapped field is caught
+    table_with_snapshots.manage_snapshots().create_branch(
+        snapshot_id=older_snapshot_id,
+        branch_name="retained",
+        max_ref_age_ms=3_600_000,
+        max_snapshot_age_ms=7_200_000,
+        min_snapshots_to_keep=5,
+    ).commit()
+
+    table_with_snapshots.manage_snapshots().fast_forward_branch(from_branch="retained", to_ref="main").commit()
+
+    ref = table_with_snapshots.refs()["retained"]
+    assert ref.snapshot_id == main_snapshot_id
+    assert ref.max_ref_age_ms == 3_600_000
+    assert ref.max_snapshot_age_ms == 7_200_000
+    assert ref.min_snapshots_to_keep == 5
+
+
+@pytest.mark.integration
+def test_fast_forward_branch_not_an_ancestor(table_with_snapshots: Table) -> None:
+    older_snapshot_id = table_with_snapshots.history()[-2].snapshot_id
+
+    table_with_snapshots.manage_snapshots().create_branch(snapshot_id=older_snapshot_id, branch_name="diverged").commit()
+
+    arrow_schema = table_with_snapshots.schema().as_arrow()
+    table_with_snapshots.append(
+        pa.Table.from_pylist([{"id": 6, "data": "f"}], schema=arrow_schema),
+        branch="diverged",
+    )
+
+    with pytest.raises(ValueError, match="Cannot fast-forward: main is not an ancestor of diverged"):
+        table_with_snapshots.manage_snapshots().fast_forward_branch(from_branch="main", to_ref="diverged").commit()