From bd48508cb7275a08fa52bd2e26b3607a7e57d88f Mon Sep 17 00:00:00 2001 From: Jeroen Schmidt Date: Mon, 13 Jul 2026 10:30:00 +0200 Subject: [PATCH] feat: add ManageSnapshots.fast_forward_branch test: fast_forward_branch auto-creates missing from_branch test: fast_forward_branch is a no-op when snapshots already equal test: fast_forward_branch rejects a tag as the source ref test: fast_forward_branch rejects missing to_ref test: fast_forward_branch rejects non-ancestor target test: fast_forward_branch preserves retention fields test: fast_forward_branch composes in a manage_snapshots chain test(integration): fast_forward_branch on real catalogs docs: add fast_forward_branch section with WAP example chore: apply ruff-format and add None-checks for mypy --- mkdocs/docs/api.md | 66 +++++++ pyiceberg/exceptions.py | 12 ++ pyiceberg/table/update/snapshot.py | 70 +++++++ tests/integration/test_snapshot_operations.py | 33 ++++ tests/table/test_manage_snapshots.py | 173 ++++++++++++++++++ tests/test_exceptions.py | 38 ++++ 6 files changed, 392 insertions(+) create mode 100644 tests/test_exceptions.py diff --git a/mkdocs/docs/api.md b/mkdocs/docs/api.md index b23f7b976e..10608d76c2 100644 --- a/mkdocs/docs/api.md +++ b/mkdocs/docs/api.md @@ -1483,6 +1483,72 @@ Remove an existing branch: table.manage_snapshots().remove_branch("dev").commit() ``` +#### Fast-forwarding a branch + +Fast-forward `from_branch` to point at the snapshot referenced by +`to_ref`. `to_ref` may be a branch or tag; `from_branch` must be a +branch. If `from_branch` does not yet exist it is created pointing at +`to_ref`'s snapshot. If both already point at the same snapshot the +call is a no-op. Otherwise `from_branch`'s current snapshot must be an +ancestor of `to_ref`'s snapshot; if not, `NotAncestorError` is raised. + +```python +with table.manage_snapshots() as ms: + ms.fast_forward_branch("main", "audit-branch") +``` + +##### End-to-end: write-audit-publish + +The canonical use case for fast-forward is the write-audit-publish +(WAP) pattern: writes proceed on a side branch, validation runs +against that branch, and only after validation succeeds is the main +branch advanced to publish the new data. + +```python +import pyarrow as pa +import pyarrow.compute as pc +from pyiceberg.catalog import load_catalog + +catalog = load_catalog("prod") +table = catalog.load_table("sales.orders") + +# 1. WRITE — create a side branch off main and append to it. +main_snapshot_id = table.current_snapshot().snapshot_id +table.manage_snapshots().create_branch( + snapshot_id=main_snapshot_id, + branch_name="audit", +).commit() + +new_rows = pa.table({ + "order_id": [1001, 1002, 1003], + "amount": [ 49.99, 129.00, 12.50], +}) +table.append(new_rows, branch="audit") + +# 2. AUDIT — scan the audit branch and run whatever validation +# your data-quality contract requires. Nothing on `main` has +# changed yet, so readers of `main` still see the pre-write state. +audit_snapshot_id = table.refs()["audit"].snapshot_id +audit_data = table.scan(snapshot_id=audit_snapshot_id).to_arrow() + +assert audit_data.num_rows > 0, "audit branch is empty" +assert pc.all(pc.greater(audit_data["amount"], 0)).as_py(), \ + "found non-positive amounts" + +# 3. PUBLISH — validation passed; fast-forward main to audit. +# Because the audit branch was created from main and only appended +# to, main's current snapshot is an ancestor of audit's snapshot, +# so the fast-forward is valid. +with table.manage_snapshots() as ms: + ms.fast_forward_branch("main", "audit") + ms.remove_branch("audit") # optional: clean up the side branch +``` + +If validation fails, callers simply skip the fast-forward step. The +audit branch (and its data files) can then be inspected, rewritten, +or removed via `remove_branch` and subsequent snapshot expiration — +without ever having polluted `main`. + ## 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/exceptions.py b/pyiceberg/exceptions.py index 4a22f1bb21..dee43a4042 100644 --- a/pyiceberg/exceptions.py +++ b/pyiceberg/exceptions.py @@ -138,3 +138,15 @@ class WaitingForLockException(Exception): class ValidationException(Exception): """Raised when validation fails.""" + + +class NoSuchSnapshotRefError(ValueError): + """Raised when a named snapshot ref (branch or tag) does not exist.""" + + +class SnapshotRefTypeError(ValueError): + """Raised when an operation expects a branch and gets a tag (or vice versa).""" + + +class NotAncestorError(ValueError): + """Raised when an operation requires ancestry between two snapshots and it does not hold.""" diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 7931edacdd..d2f64794c8 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/pyiceberg/table/update/snapshot.py @@ -26,6 +26,11 @@ from typing import TYPE_CHECKING, Generic from pyiceberg.avro.codecs import AvroCompressionCodec +from pyiceberg.exceptions import ( + NoSuchSnapshotRefError, + NotAncestorError, + SnapshotRefTypeError, +) from pyiceberg.expressions import AlwaysFalse, BooleanExpression, Or from pyiceberg.expressions.visitors import ( ROWS_MIGHT_NOT_MATCH, @@ -56,6 +61,7 @@ SnapshotSummaryCollector, Summary, ancestors_of, + is_ancestor_of, latest_ancestor_before_timestamp, update_snapshot_summaries, ) @@ -1028,6 +1034,70 @@ def _current_ancestors(self) -> set[int]: ) } + def fast_forward_branch(self, from_branch: str, to_ref: str) -> ManageSnapshots: + """Fast-forward ``from_branch`` to the snapshot referenced by ``to_ref``. + + If ``from_branch`` does not exist, it is created pointing at ``to_ref``'s snapshot (Java/Spark parity). + If both refs already point to the same snapshot the call is a no-op. + Otherwise ``from_branch`` must be a branch (not a tag) and its current + snapshot must be an ancestor of ``to_ref``'s snapshot. + + Note: + Unlike Java's ``ManageSnapshots.fastForwardBranch``, this method does not + provide Java-parity for intra-chain semantics. + + 1) Java maintains a mutable ``updatedRefs`` map that reflects prior operations + in the same chain, so calling ``createBranch`` and then ``fastForwardBranch`` + on the same ref in one chain observes the freshly-created state. + + 2) pyiceberg's ``ManageSnapshots`` accumulates ``SetSnapshotRefUpdate`` values without + mutating ``table_metadata.refs`` between chained calls; + -> The no-op and ancestry checks below operate on committed metadata only. + -> Callers that need Java-parity behavior should commit between chain steps. + + Args: + from_branch: name of the branch to advance. + to_ref: name of the branch or tag whose snapshot ``from_branch`` will point to. + + Returns: + This for method chaining. + + Raises: + NoSuchSnapshotRefError: ``to_ref`` does not exist. + SnapshotRefTypeError: ``from_branch`` exists but is a tag. + NotAncestorError: ``from_branch``'s snapshot is not an ancestor of ``to_ref``'s snapshot. + """ + refs = self._transaction.table_metadata.refs + + if to_ref not in refs: + raise NoSuchSnapshotRefError(f"Ref does not exist: {to_ref}") + to_snapshot_id = refs[to_ref].snapshot_id + + if from_branch not in refs: + return self.create_branch(snapshot_id=to_snapshot_id, branch_name=from_branch) + + from_ref = refs[from_branch] + if from_ref.snapshot_ref_type != SnapshotRefType.BRANCH: + raise SnapshotRefTypeError(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, self._transaction.table_metadata): + raise NotAncestorError(f"Cannot fast-forward: {from_branch} is not an ancestor of {to_ref}") + + update, requirement = self._transaction._set_ref_snapshot( + snapshot_id=to_snapshot_id, + ref_name=from_branch, + 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, + ) + self._updates += update + self._requirements += requirement + 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..6dda20313b 100644 --- a/tests/integration/test_snapshot_operations.py +++ b/tests/integration/test_snapshot_operations.py @@ -282,6 +282,39 @@ def test_rollback_to_timestamp_no_valid_snapshot(table_with_snapshots: Table) -> table_with_snapshots.manage_snapshots().rollback_to_timestamp(timestamp_ms=oldest_timestamp).commit() +@pytest.mark.integration +@pytest.mark.parametrize("catalog", [lf("session_catalog_hive"), lf("session_catalog")]) +def test_fast_forward_branch(catalog: Catalog) -> None: + identifier = "default.test_table_snapshot_operations" + tbl = catalog.load_table(identifier) + assert len(tbl.history()) > 2 + + # Create a side branch off an older snapshot on main, then append to it + # so that the side branch is a strict descendant of main's older snapshot + # but the *current* main is not yet caught up. + current_snapshot = tbl.current_snapshot() + assert current_snapshot is not None + main_snapshot_id = current_snapshot.snapshot_id + side_branch = "audit_ff" + + tbl.manage_snapshots().create_branch(snapshot_id=main_snapshot_id, branch_name=side_branch).commit() + + arrow_schema = tbl.schema().as_arrow() + new_rows = pa.Table.from_pylist([{col.name: None for col in arrow_schema}], schema=arrow_schema) + tbl.append(new_rows, branch=side_branch) + + # Validate appending to the side branch has advanced its snapshot. + tbl = catalog.load_table(identifier) + audit_snapshot_id = tbl.refs()[side_branch].snapshot_id + assert audit_snapshot_id != main_snapshot_id, "side-branch append should have advanced it" + + # Fast-forward main to the side branch. + tbl.manage_snapshots().fast_forward_branch(from_branch="main", to_ref=side_branch).commit() + + tbl = catalog.load_table(identifier) + assert tbl.refs()["main"].snapshot_id == audit_snapshot_id + + @pytest.mark.integration def test_rollback_to_timestamp(table_with_snapshots: Table) -> None: current_snapshot = table_with_snapshots.current_snapshot() diff --git a/tests/table/test_manage_snapshots.py b/tests/table/test_manage_snapshots.py index 93301a01c7..b77a68373d 100644 --- a/tests/table/test_manage_snapshots.py +++ b/tests/table/test_manage_snapshots.py @@ -19,7 +19,13 @@ import pytest +from pyiceberg.exceptions import ( + NoSuchSnapshotRefError, + NotAncestorError, + SnapshotRefTypeError, +) from pyiceberg.table import CommitTableResponse, Table +from pyiceberg.table.refs import SnapshotRefType from pyiceberg.table.update import SetSnapshotRefUpdate, TableUpdate @@ -177,3 +183,170 @@ 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_advances_to_descendant(table_v2: Table) -> None: + parent_snapshot_id = 3051729675574597004 + child_snapshot_id = 3055729675574597004 + + # Create a lagging branch at the parent snapshot, then reset the mock so + # the next commit's updates are the fast-forward alone. + table_v2.catalog = MagicMock() + table_v2.catalog.commit_table.return_value = _mock_commit_response(table_v2) + table_v2.manage_snapshots().create_branch(snapshot_id=parent_snapshot_id, branch_name="lagging").commit() + table_v2.catalog.commit_table.reset_mock() + + table_v2.manage_snapshots().fast_forward_branch(from_branch="lagging", to_ref="main").commit() + + updates = _get_updates(table_v2.catalog) + set_ref_updates = [u for u in updates if isinstance(u, SetSnapshotRefUpdate)] + + assert len(set_ref_updates) == 1 + update = set_ref_updates[0] + assert update.ref_name == "lagging" + assert update.snapshot_id == child_snapshot_id + assert update.type == "branch" + + +def test_fast_forward_branch_creates_missing_from(table_v2: Table) -> None: + current_snapshot = table_v2.current_snapshot() + assert current_snapshot is not None + main_snapshot_id = current_snapshot.snapshot_id + + 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() + + updates = _get_updates(table_v2.catalog) + set_ref_updates = [u for u in updates if isinstance(u, SetSnapshotRefUpdate)] + + assert len(set_ref_updates) == 1 + assert set_ref_updates[0].ref_name == "brand-new" + assert set_ref_updates[0].snapshot_id == main_snapshot_id + assert set_ref_updates[0].type == "branch" + + +def test_fast_forward_branch_noop_when_already_equal(table_v2: Table) -> None: + # The no-op check operates on committed metadata only (see + # ``fast_forward_branch`` docstring — intra-chain Java-parity is not + # implemented). Inject a second branch pointing at main's snapshot + # directly into metadata, then confirm the fast-forward stages nothing. + from pyiceberg.table.refs import SnapshotRef + + current_snapshot = table_v2.current_snapshot() + assert current_snapshot is not None + main_snapshot_id = current_snapshot.snapshot_id + table_v2.metadata.refs["peer"] = SnapshotRef( + snapshot_id=main_snapshot_id, + snapshot_ref_type="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() + + # A pure no-op stages no updates; commit_transaction short-circuits. + table_v2.catalog.commit_table.assert_not_called() + + +def test_fast_forward_branch_rejects_tag_as_source(table_v2: Table) -> None: + # Precondition: the fixture provides a tag named "test". + assert table_v2.metadata.refs["test"].snapshot_ref_type == SnapshotRefType.TAG + table_v2.catalog = MagicMock() + + with pytest.raises(SnapshotRefTypeError, match="Ref test is a tag, not a branch"): + table_v2.manage_snapshots().fast_forward_branch(from_branch="test", to_ref="main").commit() + + table_v2.catalog.commit_table.assert_not_called() + + +def test_fast_forward_branch_rejects_missing_to_ref(table_v2: Table) -> None: + table_v2.catalog = MagicMock() + + with pytest.raises(NoSuchSnapshotRefError, match="Ref does not exist: nonexistent"): + table_v2.manage_snapshots().fast_forward_branch(from_branch="main", to_ref="nonexistent").commit() + + table_v2.catalog.commit_table.assert_not_called() + + +def test_fast_forward_branch_rejects_non_ancestor(table_v2: Table) -> None: + # Non-ancestor check operates on committed metadata only (see + # ``fast_forward_branch`` docstring — intra-chain Java-parity is not + # implemented). Inject "ahead" branch directly into metadata so the + # ancestry check can observe it. + from pyiceberg.table.refs import SnapshotRef + + newer_snapshot_id = 3055729675574597004 # main's current snapshot; "test" tag points at older snapshot + + table_v2.metadata.refs["ahead"] = SnapshotRef( + snapshot_id=newer_snapshot_id, + snapshot_ref_type="branch", + ) + + table_v2.catalog = MagicMock() + + # Try to fast-forward "ahead" (at newer) backwards to "test" (at older). + with pytest.raises(NotAncestorError, match="Cannot fast-forward: ahead is not an ancestor of test"): + table_v2.manage_snapshots().fast_forward_branch(from_branch="ahead", to_ref="test").commit() + + table_v2.catalog.commit_table.assert_not_called() + + +def test_fast_forward_branch_preserves_retention_fields(table_v2: Table) -> None: + from pyiceberg.table.refs import SnapshotRef + + parent_snapshot_id = 3051729675574597004 + child_snapshot_id = 3055729675574597004 + + # Inject a branch with all three retention fields set into the metadata. + table_v2.metadata.refs["retained"] = SnapshotRef( + snapshot_id=parent_snapshot_id, + snapshot_ref_type="branch", + max_ref_age_ms=1000, + max_snapshot_age_ms=2000, + min_snapshots_to_keep=3, + ) + + 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="retained", to_ref="main").commit() + + updates = _get_updates(table_v2.catalog) + set_ref_updates = [u for u in updates if isinstance(u, SetSnapshotRefUpdate)] + + assert len(set_ref_updates) == 1 + update = set_ref_updates[0] + assert update.ref_name == "retained" + assert update.snapshot_id == child_snapshot_id + assert update.max_ref_age_ms == 1000 + assert update.max_snapshot_age_ms == 2000 + assert update.min_snapshots_to_keep == 3 + + +def test_fast_forward_branch_chains(table_v2: Table) -> None: + parent_snapshot_id = 3051729675574597004 + + 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="stream").fast_forward_branch( + from_branch="stream", to_ref="main" + ).create_tag(snapshot_id=parent_snapshot_id, tag_name="stream-v1") + + updates = _get_updates(table_v2.catalog) + set_ref_updates = [u for u in updates if isinstance(u, SetSnapshotRefUpdate)] + + ref_names = {u.ref_name for u in set_ref_updates} + assert "stream" in ref_names # from create_branch AND fast_forward_branch + assert "stream-v1" in ref_names # from create_tag + + # There should be two updates for `stream` (create at parent, then fast-forward to child) + # and one for `stream-v1`. The commit protocol accepts multiple updates for the same ref. + stream_updates = [u for u in set_ref_updates if u.ref_name == "stream"] + assert len(stream_updates) == 2 + assert stream_updates[0].snapshot_id == parent_snapshot_id + assert stream_updates[1].snapshot_id == 3055729675574597004 diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py new file mode 100644 index 0000000000..d59a35779a --- /dev/null +++ b/tests/test_exceptions.py @@ -0,0 +1,38 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +import pytest + +from pyiceberg.exceptions import ( + NoSuchSnapshotRefError, + NotAncestorError, + SnapshotRefTypeError, +) + + +def test_no_such_snapshot_ref_error_is_value_error() -> None: + with pytest.raises(ValueError): + raise NoSuchSnapshotRefError("nope") + + +def test_snapshot_ref_type_error_is_value_error() -> None: + with pytest.raises(ValueError): + raise SnapshotRefTypeError("nope") + + +def test_not_ancestor_error_is_value_error() -> None: + with pytest.raises(ValueError): + raise NotAncestorError("nope")