diff --git a/mkdocs/docs/api.md b/mkdocs/docs/api.md index 1e17e64f42..443df9234f 100644 --- a/mkdocs/docs/api.md +++ b/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. diff --git a/pyiceberg/table/update/__init__.py b/pyiceberg/table/update/__init__.py index 64838b0bd6..1cc559f5ab 100644 --- a/pyiceberg/table/update/__init__.py +++ b/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): diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 6bddd27905..682042d4e3 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/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. diff --git a/tests/integration/test_snapshot_operations.py b/tests/integration/test_snapshot_operations.py index 07fb77edbb..e829449849 100644 --- a/tests/integration/test_snapshot_operations.py +++ b/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() diff --git a/tests/table/test_manage_snapshots.py b/tests/table/test_manage_snapshots.py index 93301a01c7..152b541e50 100644 --- a/tests/table/test_manage_snapshots.py +++ b/tests/table/test_manage_snapshots.py @@ -18,12 +18,17 @@ from uuid import uuid4 import pytest +from pydantic import ValidationError as PydanticValidationError from pyiceberg.table import CommitTableResponse, Table -from pyiceberg.table.update import SetSnapshotRefUpdate, TableUpdate +from pyiceberg.table.refs import SnapshotRef, SnapshotRefType +from pyiceberg.table.update import AssertRefSnapshotId, SetSnapshotRefUpdate, TableRequirement, TableUpdate +PARENT_SNAPSHOT_ID = 3051729675574597004 +CHILD_SNAPSHOT_ID = 3055729675574597004 -def _mock_commit_response(table: Table) -> CommitTableResponse: + +def mock_commit_response(table: Table) -> CommitTableResponse: return CommitTableResponse( metadata=table.metadata, metadata_location="s3://bucket/tbl", @@ -31,23 +36,31 @@ def _mock_commit_response(table: Table) -> CommitTableResponse: ) -def _get_updates(mock_catalog: MagicMock) -> tuple[TableUpdate, ...]: +def get_updates(mock_catalog: MagicMock) -> tuple[TableUpdate, ...]: args, _ = mock_catalog.commit_table.call_args return args[2] +def get_requirements(mock_catalog: MagicMock) -> tuple[TableRequirement, ...]: + args, _ = mock_catalog.commit_table.call_args + return args[1] + + +def get_set_ref_updates(mock_catalog: MagicMock) -> list[SetSnapshotRefUpdate]: + return [update for update in get_updates(mock_catalog) if isinstance(update, SetSnapshotRefUpdate)] + + def test_set_current_snapshot_basic(table_v2: Table) -> None: snapshot_one = 3051729675574597004 table_v2.catalog = MagicMock() - table_v2.catalog.commit_table.return_value = _mock_commit_response(table_v2) + table_v2.catalog.commit_table.return_value = mock_commit_response(table_v2) table_v2.manage_snapshots().set_current_snapshot(snapshot_id=snapshot_one).commit() table_v2.catalog.commit_table.assert_called_once() - updates = _get_updates(table_v2.catalog) - set_ref_updates = [u for u in updates if isinstance(u, SetSnapshotRefUpdate)] + set_ref_updates = get_set_ref_updates(table_v2.catalog) assert len(set_ref_updates) == 1 update = set_ref_updates[0] @@ -71,7 +84,7 @@ def test_set_current_snapshot_to_current(table_v2: Table) -> None: assert current_snapshot is not None table_v2.catalog = MagicMock() - table_v2.catalog.commit_table.return_value = _mock_commit_response(table_v2) + table_v2.catalog.commit_table.return_value = mock_commit_response(table_v2) table_v2.manage_snapshots().set_current_snapshot(snapshot_id=current_snapshot.snapshot_id).commit() @@ -81,14 +94,13 @@ def test_set_current_snapshot_to_current(table_v2: Table) -> None: def test_set_current_snapshot_chained_with_tag(table_v2: Table) -> None: snapshot_one = 3051729675574597004 table_v2.catalog = MagicMock() - table_v2.catalog.commit_table.return_value = _mock_commit_response(table_v2) + table_v2.catalog.commit_table.return_value = mock_commit_response(table_v2) (table_v2.manage_snapshots().set_current_snapshot(snapshot_id=snapshot_one).create_tag(snapshot_one, "my-tag").commit()) table_v2.catalog.commit_table.assert_called_once() - updates = _get_updates(table_v2.catalog) - set_ref_updates = [u for u in updates if isinstance(u, SetSnapshotRefUpdate)] + set_ref_updates = get_set_ref_updates(table_v2.catalog) assert len(set_ref_updates) == 2 assert {u.ref_name for u in set_ref_updates} == {"main", "my-tag"} @@ -101,14 +113,13 @@ def test_set_current_snapshot_with_extensive_snapshots(table_v2_with_extensive_s target_snapshot = snapshots[50].snapshot_id table_v2_with_extensive_snapshots.catalog = MagicMock() - table_v2_with_extensive_snapshots.catalog.commit_table.return_value = _mock_commit_response(table_v2_with_extensive_snapshots) + table_v2_with_extensive_snapshots.catalog.commit_table.return_value = mock_commit_response(table_v2_with_extensive_snapshots) table_v2_with_extensive_snapshots.manage_snapshots().set_current_snapshot(snapshot_id=target_snapshot).commit() table_v2_with_extensive_snapshots.catalog.commit_table.assert_called_once() - updates = _get_updates(table_v2_with_extensive_snapshots.catalog) - set_ref_updates = [u for u in updates if isinstance(u, SetSnapshotRefUpdate)] + set_ref_updates = get_set_ref_updates(table_v2_with_extensive_snapshots.catalog) assert len(set_ref_updates) == 1 assert set_ref_updates[0].snapshot_id == target_snapshot @@ -119,12 +130,11 @@ def test_set_current_snapshot_by_ref_name(table_v2: Table) -> None: assert current_snapshot is not None table_v2.catalog = MagicMock() - table_v2.catalog.commit_table.return_value = _mock_commit_response(table_v2) + table_v2.catalog.commit_table.return_value = mock_commit_response(table_v2) table_v2.manage_snapshots().set_current_snapshot(ref_name="main").commit() - updates = _get_updates(table_v2.catalog) - set_ref_updates = [u for u in updates if isinstance(u, SetSnapshotRefUpdate)] + set_ref_updates = get_set_ref_updates(table_v2.catalog) assert len(set_ref_updates) == 1 assert set_ref_updates[0].snapshot_id == current_snapshot.snapshot_id @@ -155,7 +165,7 @@ def test_set_current_snapshot_requires_one_argument(table_v2: Table) -> None: def test_set_current_snapshot_chained_with_create_tag(table_v2: Table) -> None: snapshot_one = 3051729675574597004 table_v2.catalog = MagicMock() - table_v2.catalog.commit_table.return_value = _mock_commit_response(table_v2) + table_v2.catalog.commit_table.return_value = mock_commit_response(table_v2) # create a tag and immediately use it to set current snapshot ( @@ -167,8 +177,7 @@ def test_set_current_snapshot_chained_with_create_tag(table_v2: Table) -> None: table_v2.catalog.commit_table.assert_called_once() - updates = _get_updates(table_v2.catalog) - set_ref_updates = [u for u in updates if isinstance(u, SetSnapshotRefUpdate)] + set_ref_updates = get_set_ref_updates(table_v2.catalog) # should have the tag and the main branch update assert len(set_ref_updates) == 2 @@ -177,3 +186,219 @@ def test_set_current_snapshot_chained_with_create_tag(table_v2: Table) -> None: # The main branch should point to the same snapshot as the tag main_update = next(u for u in set_ref_updates if u.ref_name == "main") assert main_update.snapshot_id == snapshot_one + + +def test_fast_forward_branch_basic(table_v2: Table) -> None: + table_v2.metadata.refs["lagging"] = SnapshotRef( + snapshot_id=PARENT_SNAPSHOT_ID, + snapshot_ref_type=SnapshotRefType.BRANCH, + ) + table_v2.catalog = MagicMock() + table_v2.catalog.commit_table.return_value = mock_commit_response(table_v2) + + table_v2.manage_snapshots().fast_forward_branch(from_branch="lagging", to_ref="main").commit() + + set_ref_updates = get_set_ref_updates(table_v2.catalog) + assert len(set_ref_updates) == 1 + assert set_ref_updates[0].ref_name == "lagging" + assert set_ref_updates[0].snapshot_id == CHILD_SNAPSHOT_ID + assert set_ref_updates[0].type == SnapshotRefType.BRANCH + + +def test_fast_forward_branch_to_tag(table_v2: Table) -> None: + table_v2.metadata.refs["lagging"] = SnapshotRef( + snapshot_id=PARENT_SNAPSHOT_ID, + snapshot_ref_type=SnapshotRefType.BRANCH, + ) + table_v2.metadata.refs["release"] = SnapshotRef( + snapshot_id=CHILD_SNAPSHOT_ID, + snapshot_ref_type=SnapshotRefType.TAG, + ) + table_v2.catalog = MagicMock() + table_v2.catalog.commit_table.return_value = mock_commit_response(table_v2) + + table_v2.manage_snapshots().fast_forward_branch(from_branch="lagging", to_ref="release").commit() + + set_ref_updates = get_set_ref_updates(table_v2.catalog) + assert len(set_ref_updates) == 1 + assert set_ref_updates[0].ref_name == "lagging" + assert set_ref_updates[0].snapshot_id == CHILD_SNAPSHOT_ID + assert set_ref_updates[0].type == SnapshotRefType.BRANCH + + +def test_fast_forward_branch_creates_branch(table_v2: Table) -> None: + table_v2.catalog = MagicMock() + table_v2.catalog.commit_table.return_value = mock_commit_response(table_v2) + + table_v2.manage_snapshots().fast_forward_branch(from_branch="brand-new", to_ref="main").commit() + + set_ref_updates = get_set_ref_updates(table_v2.catalog) + assert len(set_ref_updates) == 1 + update = set_ref_updates[0] + assert update.ref_name == "brand-new" + assert update.snapshot_id == CHILD_SNAPSHOT_ID + assert update.type == SnapshotRefType.BRANCH + # a branch created by the fast-forward gets the default retention + assert update.max_ref_age_ms is None + assert update.max_snapshot_age_ms is None + assert update.min_snapshots_to_keep is None + + +def test_fast_forward_branch_noop(table_v2: Table) -> None: + table_v2.metadata.refs["peer"] = SnapshotRef( + snapshot_id=CHILD_SNAPSHOT_ID, + snapshot_ref_type=SnapshotRefType.BRANCH, + ) + table_v2.catalog = MagicMock() + table_v2.catalog.commit_table.return_value = mock_commit_response(table_v2) + + table_v2.manage_snapshots().fast_forward_branch(from_branch="peer", to_ref="main").commit() + + table_v2.catalog.commit_table.assert_not_called() + + +def test_fast_forward_branch_preserves_retention(table_v2: Table) -> None: + table_v2.metadata.refs["lagging"] = SnapshotRef( + snapshot_id=PARENT_SNAPSHOT_ID, + snapshot_ref_type=SnapshotRefType.BRANCH, + max_ref_age_ms=3_600_000, + max_snapshot_age_ms=7_200_000, + min_snapshots_to_keep=5, + ) + table_v2.catalog = MagicMock() + table_v2.catalog.commit_table.return_value = mock_commit_response(table_v2) + + table_v2.manage_snapshots().fast_forward_branch(from_branch="lagging", to_ref="main").commit() + + set_ref_updates = get_set_ref_updates(table_v2.catalog) + assert len(set_ref_updates) == 1 + update = set_ref_updates[0] + assert update.snapshot_id == CHILD_SNAPSHOT_ID + assert update.max_ref_age_ms == 3_600_000 + assert update.max_snapshot_age_ms == 7_200_000 + assert update.min_snapshots_to_keep == 5 + + +def test_fast_forward_branch_chained_with_create_branch(table_v2: Table) -> None: + table_v2.catalog = MagicMock() + table_v2.catalog.commit_table.return_value = mock_commit_response(table_v2) + + with table_v2.manage_snapshots() as ms: + ms.create_branch( + snapshot_id=PARENT_SNAPSHOT_ID, + branch_name="staging", + max_ref_age_ms=3_600_000, + max_snapshot_age_ms=7_200_000, + min_snapshots_to_keep=5, + ).fast_forward_branch(from_branch="staging", to_ref="main") + + # the fast-forward sees the branch created in the same chain and replaces its update + set_ref_updates = get_set_ref_updates(table_v2.catalog) + assert len(set_ref_updates) == 1 + update = set_ref_updates[0] + assert update.ref_name == "staging" + assert update.snapshot_id == CHILD_SNAPSHOT_ID + assert update.max_ref_age_ms == 3_600_000 + assert update.max_snapshot_age_ms == 7_200_000 + assert update.min_snapshots_to_keep == 5 + + +def test_fast_forward_branch_unknown_ref(table_v2: Table) -> None: + table_v2.catalog = MagicMock() + + with pytest.raises(ValueError, match="Ref does not exist: nonexistent"): + table_v2.manage_snapshots().fast_forward_branch(from_branch="main", to_ref="nonexistent") + + table_v2.catalog.commit_table.assert_not_called() + + +def test_fast_forward_branch_from_tag(table_v2: Table) -> None: + table_v2.catalog = MagicMock() + + with pytest.raises(ValueError, match="Ref test is a tag not a branch"): + table_v2.manage_snapshots().fast_forward_branch(from_branch="test", to_ref="main") + + table_v2.catalog.commit_table.assert_not_called() + + +def test_fast_forward_branch_not_ancestor(table_v2: Table) -> None: + # the tag points at the parent of the current snapshot + table_v2.catalog = MagicMock() + + with pytest.raises(ValueError, match="Cannot fast-forward: main is not an ancestor of test"): + table_v2.manage_snapshots().fast_forward_branch(from_branch="main", to_ref="test") + + table_v2.catalog.commit_table.assert_not_called() + + +def test_fast_forward_branch_chained_with_remove_branch(table_v2: Table) -> None: + table_v2.metadata.refs["stale"] = SnapshotRef( + snapshot_id=PARENT_SNAPSHOT_ID, + snapshot_ref_type=SnapshotRefType.BRANCH, + ) + table_v2.catalog = MagicMock() + table_v2.catalog.commit_table.return_value = mock_commit_response(table_v2) + + # the branch removed earlier in the chain is created again at main + with table_v2.manage_snapshots() as ms: + ms.remove_branch("stale").fast_forward_branch(from_branch="stale", to_ref="main") + + set_ref_updates = get_set_ref_updates(table_v2.catalog) + assert len(set_ref_updates) == 1 + assert set_ref_updates[0].ref_name == "stale" + assert set_ref_updates[0].snapshot_id == CHILD_SNAPSHOT_ID + + +def test_fast_forward_branch_requirement(table_v2: Table) -> None: + table_v2.metadata.refs["lagging"] = SnapshotRef( + snapshot_id=PARENT_SNAPSHOT_ID, + snapshot_ref_type=SnapshotRefType.BRANCH, + ) + table_v2.catalog = MagicMock() + table_v2.catalog.commit_table.return_value = mock_commit_response(table_v2) + + table_v2.manage_snapshots().fast_forward_branch(from_branch="lagging", to_ref="main").commit() + + # the requirement asserts the committed snapshot of the branch + ref_requirements = [r for r in get_requirements(table_v2.catalog) if isinstance(r, AssertRefSnapshotId)] + assert len(ref_requirements) == 1 + assert ref_requirements[0].ref == "lagging" + assert ref_requirements[0].snapshot_id == PARENT_SNAPSHOT_ID + + +def test_fast_forward_branch_chained_requirement(table_v2: Table) -> None: + table_v2.catalog = MagicMock() + table_v2.catalog.commit_table.return_value = mock_commit_response(table_v2) + + with table_v2.manage_snapshots() as ms: + ms.create_branch(snapshot_id=PARENT_SNAPSHOT_ID, branch_name="staging").fast_forward_branch( + from_branch="staging", to_ref="main" + ) + + ref_requirements = [r for r in get_requirements(table_v2.catalog) if isinstance(r, AssertRefSnapshotId)] + assert len(ref_requirements) == 1 + assert ref_requirements[0].ref == "staging" + assert ref_requirements[0].snapshot_id is None + + +def test_fast_forward_branch_unknown_snapshot(table_v2: Table) -> None: + table_v2.metadata.refs["dangling"] = SnapshotRef( + snapshot_id=1234567890000, + snapshot_ref_type=SnapshotRefType.BRANCH, + ) + table_v2.catalog = MagicMock() + + with pytest.raises(ValueError, match="Cannot fast-forward to unknown snapshot id: 1234567890000"): + table_v2.manage_snapshots().fast_forward_branch(from_branch="main", to_ref="dangling") + + table_v2.catalog.commit_table.assert_not_called() + + +def test_create_branch_invalid_retention(table_v2: Table) -> None: + table_v2.catalog = MagicMock() + + # rejected when passed, not when the ref is built + with pytest.raises(PydanticValidationError, match="min_snapshots_to_keep"): + table_v2.manage_snapshots().create_branch(snapshot_id=PARENT_SNAPSHOT_ID, branch_name="b", min_snapshots_to_keep=0) + + table_v2.catalog.commit_table.assert_not_called()