From 4f48ea49f5fae00ac61fefdbb7a46a95cacc355d Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Fri, 7 Aug 2026 21:23:23 -0700 Subject: [PATCH] Preserve a ref's retention policy when its snapshot moves create_branch() accepts max_ref_age_ms, max_snapshot_age_ms and min_snapshots_to_keep, but any later operation that moves the ref drops all three. Both paths that move a ref rebuild it from a SetSnapshotRefUpdate that omits the retention fields, so the applied ref falls back to the defaults: - writing to a branch, via _SnapshotProducer._commit() - set_current_snapshot(), and the rollbacks that delegate to it A branch created with a retention policy therefore loses it on its first write, which is the point at which the policy starts to matter. Carry the existing ref's retention fields through on both paths, mirroring Java's TableMetadata.Builder.setBranchSnapshotInternal, which rebuilds the moved ref with SnapshotRef.builderFrom(ref, replacementSnapshotId). create_tag() and create_branch() still set the fields explicitly, so a ref with no policy does not acquire one. Signed-off-by: 1fanwang <1fannnw@gmail.com> --- pyiceberg/table/update/snapshot.py | 8 ++ tests/integration/test_writes/test_writes.py | 27 +++++ tests/table/test_manage_snapshots.py | 104 +++++++++++++++++++ 3 files changed, 139 insertions(+) diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 7931edacdd..b5d942073e 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/pyiceberg/table/update/snapshot.py @@ -309,6 +309,7 @@ def _commit(self) -> UpdatesAndRequirements: (), ) else: + existing_ref = self._transaction.table_metadata.refs.get(self._target_branch) return ( ( add_snapshot_update, @@ -317,6 +318,9 @@ def _commit(self) -> UpdatesAndRequirements: parent_snapshot_id=self._parent_snapshot_id, ref_name=self._target_branch, type=SnapshotRefType.BRANCH, + max_ref_age_ms=existing_ref.max_ref_age_ms if existing_ref else None, + max_snapshot_age_ms=existing_ref.max_snapshot_age_ms if existing_ref else None, + min_snapshots_to_keep=existing_ref.min_snapshots_to_keep if existing_ref else None, ), ), ( @@ -966,10 +970,14 @@ def set_current_snapshot(self, snapshot_id: int | None = None, ref_name: str | N if self._transaction.table_metadata.snapshot_by_id(target_snapshot_id) is None: raise ValueError(f"Cannot set current snapshot to unknown snapshot id: {target_snapshot_id}") + existing_ref = self._transaction.table_metadata.refs.get(MAIN_BRANCH) update, requirement = self._transaction._set_ref_snapshot( snapshot_id=target_snapshot_id, ref_name=MAIN_BRANCH, type=SnapshotRefType.BRANCH, + max_ref_age_ms=existing_ref.max_ref_age_ms if existing_ref else None, + max_snapshot_age_ms=existing_ref.max_snapshot_age_ms if existing_ref else None, + min_snapshots_to_keep=existing_ref.min_snapshots_to_keep if existing_ref else None, ) self._transaction._stage(update, requirement) return self diff --git a/tests/integration/test_writes/test_writes.py b/tests/integration/test_writes/test_writes.py index 30fdd76ab7..51e12cebfd 100644 --- a/tests/integration/test_writes/test_writes.py +++ b/tests/integration/test_writes/test_writes.py @@ -2235,6 +2235,33 @@ def test_append_to_existing_branch(session_catalog: Catalog, arrow_table_with_nu assert branch_snapshot.parent_snapshot_id == main_snapshot.snapshot_id +@pytest.mark.integration +def test_write_to_branch_preserves_retention(session_catalog: Catalog, arrow_table_with_null: pa.Table) -> None: + """Writing to a branch must not reset the retention policy it was created with.""" + identifier = "default.test_branch_write_retention" + branch = "retention_branch" + tbl = _create_table(session_catalog, identifier, {"format-version": "2"}, [arrow_table_with_null]) + assert tbl.metadata.current_snapshot_id is not None + + tbl.manage_snapshots().create_branch( + snapshot_id=tbl.metadata.current_snapshot_id, + branch_name=branch, + max_ref_age_ms=86400000, + max_snapshot_age_ms=3600000, + min_snapshots_to_keep=5, + ).commit() + + tbl.append(arrow_table_with_null, branch=branch) + tbl.append(arrow_table_with_null, branch=branch) + + ref = tbl.metadata.refs[branch] + assert ref.max_ref_age_ms == 86400000 + assert ref.max_snapshot_age_ms == 3600000 + assert ref.min_snapshots_to_keep == 5 + # a branch with no retention configured stays that way + assert tbl.metadata.refs["main"].max_ref_age_ms is None + + @pytest.mark.integration def test_delete_to_existing_branch(session_catalog: Catalog, arrow_table_with_null: pa.Table) -> None: identifier = "default.test_existing_branch_delete" diff --git a/tests/table/test_manage_snapshots.py b/tests/table/test_manage_snapshots.py index 93301a01c7..773943e00b 100644 --- a/tests/table/test_manage_snapshots.py +++ b/tests/table/test_manage_snapshots.py @@ -17,8 +17,10 @@ from unittest.mock import MagicMock from uuid import uuid4 +import pyarrow as pa import pytest +from pyiceberg.catalog import Catalog from pyiceberg.table import CommitTableResponse, Table from pyiceberg.table.update import SetSnapshotRefUpdate, TableUpdate @@ -177,3 +179,105 @@ 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_branch_write_preserves_retention(catalog_with_warehouse: Catalog) -> None: + """Writing to a branch keeps the retention policy it was created with.""" + catalog_with_warehouse.create_namespace("branch_retention") + schema = pa.schema([pa.field("id", pa.int64())]) + tbl = catalog_with_warehouse.create_table("branch_retention.tbl", schema=schema) + tbl.append(pa.table({"id": [1]}, schema=schema)) + tbl = catalog_with_warehouse.load_table("branch_retention.tbl") + + snapshot_id = tbl.metadata.current_snapshot_id + assert snapshot_id is not None + tbl.manage_snapshots().create_branch( + snapshot_id=snapshot_id, + branch_name="audit", + max_ref_age_ms=86400000, + max_snapshot_age_ms=3600000, + min_snapshots_to_keep=5, + ).commit() + + for i in range(3): + tbl = catalog_with_warehouse.load_table("branch_retention.tbl") + tbl.append(pa.table({"id": [i + 2]}, schema=schema), branch="audit") + + ref = catalog_with_warehouse.load_table("branch_retention.tbl").metadata.refs["audit"] + assert ref.max_ref_age_ms == 86400000 + assert ref.max_snapshot_age_ms == 3600000 + assert ref.min_snapshots_to_keep == 5 + + +def test_branch_write_without_retention_stays_unset(catalog_with_warehouse: Catalog) -> None: + """A branch created without a retention policy does not gain one from a write.""" + catalog_with_warehouse.create_namespace("branch_no_retention") + schema = pa.schema([pa.field("id", pa.int64())]) + tbl = catalog_with_warehouse.create_table("branch_no_retention.tbl", schema=schema) + tbl.append(pa.table({"id": [1]}, schema=schema)) + tbl = catalog_with_warehouse.load_table("branch_no_retention.tbl") + + snapshot_id = tbl.metadata.current_snapshot_id + assert snapshot_id is not None + tbl.manage_snapshots().create_branch(snapshot_id=snapshot_id, branch_name="plain").commit() + + tbl = catalog_with_warehouse.load_table("branch_no_retention.tbl") + tbl.append(pa.table({"id": [2]}, schema=schema), branch="plain") + + ref = catalog_with_warehouse.load_table("branch_no_retention.tbl").metadata.refs["plain"] + assert ref.max_ref_age_ms is None + assert ref.max_snapshot_age_ms is None + assert ref.min_snapshots_to_keep is None + + +def test_rollback_preserves_retention(catalog_with_warehouse: Catalog) -> None: + """Moving a ref keeps its retention policy; rollback and set_current_snapshot move main.""" + catalog_with_warehouse.create_namespace("rollback_retention") + schema = pa.schema([pa.field("id", pa.int64())]) + tbl = catalog_with_warehouse.create_table("rollback_retention.tbl", schema=schema) + tbl.append(pa.table({"id": [1]}, schema=schema)) + tbl = catalog_with_warehouse.load_table("rollback_retention.tbl") + first = tbl.metadata.current_snapshot_id + assert first is not None + + tbl.append(pa.table({"id": [2]}, schema=schema)) + tbl = catalog_with_warehouse.load_table("rollback_retention.tbl") + head = tbl.metadata.current_snapshot_id + assert head is not None + + tbl.manage_snapshots().create_branch( + snapshot_id=head, + branch_name="main", + max_ref_age_ms=99999, + max_snapshot_age_ms=1234, + min_snapshots_to_keep=7, + ).commit() + + tbl = catalog_with_warehouse.load_table("rollback_retention.tbl") + tbl.manage_snapshots().rollback_to_snapshot(first).commit() + + ref = catalog_with_warehouse.load_table("rollback_retention.tbl").metadata.refs["main"] + assert ref.snapshot_id == first + assert ref.max_ref_age_ms == 99999 + assert ref.max_snapshot_age_ms == 1234 + assert ref.min_snapshots_to_keep == 7 + + +def test_rollback_without_retention_stays_unset(catalog_with_warehouse: Catalog) -> None: + """A ref with no retention policy does not gain one from being moved.""" + catalog_with_warehouse.create_namespace("rollback_plain") + schema = pa.schema([pa.field("id", pa.int64())]) + tbl = catalog_with_warehouse.create_table("rollback_plain.tbl", schema=schema) + tbl.append(pa.table({"id": [1]}, schema=schema)) + tbl = catalog_with_warehouse.load_table("rollback_plain.tbl") + first = tbl.metadata.current_snapshot_id + assert first is not None + + tbl.append(pa.table({"id": [2]}, schema=schema)) + tbl = catalog_with_warehouse.load_table("rollback_plain.tbl") + tbl.manage_snapshots().rollback_to_snapshot(first).commit() + + ref = catalog_with_warehouse.load_table("rollback_plain.tbl").metadata.refs["main"] + assert ref.max_ref_age_ms is None + assert ref.max_snapshot_age_ms is None + assert ref.min_snapshots_to_keep is None