Enforce max-ref-age-ms when expiring snapshots

mkdocs/docs/api.md

@@ -1527,6 +1527,35 @@ def cleanup_old_snapshots(table_name: str, snapshot_ids: list[int]):
 cleanup_old_snapshots("analytics.user_events", [12345, 67890, 11111])
 ```
 
+### Expiring Branches and Tags
+
+A branch or tag protects the snapshot it points at from expiration, along with that snapshot's
+ancestors. A ref left behind by a failed job therefore pins storage indefinitely unless it is
+removed.
+
+`remove_expired_refs()` drops refs whose retention period has elapsed, using the ref's own
+`max-ref-age-ms` when set and the `history.expire.max-ref-age-ms` table property otherwise.
+The `main` branch never expires.
+
+```python
+from datetime import datetime, timedelta, timezone
+
+# Give an audit branch a one-day lifetime
+table.manage_snapshots().create_branch(
+    snapshot_id=table.metadata.current_snapshot_id,
+    branch_name="audit-2024-01-15",
+    max_ref_age_ms=24 * 60 * 60 * 1000,
+).commit()
+
+# Later: drop expired refs, then reclaim the snapshots they were pinning
+table.maintenance.expire_snapshots().remove_expired_refs().older_than(
+    datetime.now(timezone.utc) - timedelta(days=3)
+).commit()
+```
+
+Call `remove_expired_refs()` and `older_than()` in either order. Snapshots released by the
+removed refs are eligible for expiration in the same commit.
+
 ## Views
 
 If PyIceberg is unable to automatically determine view support on your REST Catalog, you can manually specify, `"view-endpoints-supported": "true"`:

pyiceberg/table/__init__.py

@@ -212,6 +212,10 @@ class TableProperties:
     MIN_SNAPSHOTS_TO_KEEP = "history.expire.min-snapshots-to-keep"
     MIN_SNAPSHOTS_TO_KEEP_DEFAULT = 1
 
+    MAX_REF_AGE_MS = "history.expire.max-ref-age-ms"
+    # sys.maxsize would be 2**31-1 on a 32-bit build, silently expiring refs after ~25 days.
+    MAX_REF_AGE_MS_DEFAULT = 2**63 - 1
+
 
 class Transaction:
     _table: Table

pyiceberg/table/update/snapshot.py

@@ -21,7 +21,7 @@
 from abc import abstractmethod
 from collections import defaultdict
 from collections.abc import Callable
-from datetime import datetime
+from datetime import datetime, timezone
 from functools import cached_property
 from typing import TYPE_CHECKING, Generic
 
@@ -1040,12 +1040,16 @@ class ExpireSnapshots(UpdateTableMetadata["ExpireSnapshots"]):
     _updates: tuple[TableUpdate, ...]
     _requirements: tuple[TableRequirement, ...]
     _snapshot_ids_to_expire: set[int]
+    _ref_names_to_remove: set[str]
+    _expire_older_than_ms: int | None
 
     def __init__(self, transaction: Transaction) -> None:
         super().__init__(transaction)
         self._updates = ()
         self._requirements = ()
         self._snapshot_ids_to_expire = set()
+        self._ref_names_to_remove = set()
+        self._expire_older_than_ms = None
 
     def _commit(self) -> UpdatesAndRequirements:
         """
@@ -1057,9 +1061,18 @@ def _commit(self) -> UpdatesAndRequirements:
             Tuple of updates and requirements to be committed,
             as required by the calling parent apply functions.
         """
-        # Remove any protected snapshot IDs from the set to expire, just in case
+        # Resolved here rather than in older_than() so that snapshots released by
+        # remove_expired_refs() are expirable no matter which order the two were called in.
         protected_ids = self._get_protected_snapshot_ids()
+        if self._expire_older_than_ms is not None:
+            for snapshot in self._transaction.table_metadata.snapshots:
+                if snapshot.timestamp_ms < self._expire_older_than_ms and snapshot.snapshot_id not in protected_ids:
+                    self._snapshot_ids_to_expire.add(snapshot.snapshot_id)
+
+        # Remove any protected snapshot IDs from the set to expire, just in case
         self._snapshot_ids_to_expire -= protected_ids
+        for ref_name in self._ref_names_to_remove:
+            self._updates += (RemoveSnapshotRefUpdate(ref_name=ref_name),)
         update = RemoveSnapshotsUpdate(snapshot_ids=self._snapshot_ids_to_expire)
         self._updates += (update,)
         return self._updates, self._requirements
@@ -1069,16 +1082,58 @@ def _get_protected_snapshot_ids(self) -> set[int]:
         Get the IDs of protected snapshots.
 
         These are the HEAD snapshots of all branches and all tagged snapshots.  These ids are to be excluded from expiration.
+        Refs staged for removal by :meth:`remove_expired_refs` do not protect anything.
 
         Returns:
             Set of protected snapshot IDs to exclude from expiration.
         """
         return {
             ref.snapshot_id
-            for ref in self._transaction.table_metadata.refs.values()
+            for ref_name, ref in self._transaction.table_metadata.refs.items()
             if ref.snapshot_ref_type in [SnapshotRefType.TAG, SnapshotRefType.BRANCH]
+            and ref_name not in self._ref_names_to_remove
         }
 
+    def remove_expired_refs(self) -> ExpireSnapshots:
+        """
+        Remove branches and tags whose retention period has elapsed.
+
+        A ref's age is measured from the timestamp of the snapshot it points at, against its own
+        ``max-ref-age-ms`` when set, otherwise the ``history.expire.max-ref-age-ms`` table property.
+        The ``main`` branch never expires. A ref pointing at a snapshot that no longer exists is also
+        removed, since it can no longer be resolved.
+
+        This is step 2 of the snapshot retention policy in the Iceberg spec:
+        https://iceberg.apache.org/spec/#snapshot-retention-policy
+
+        Returns:
+            This for method chaining.
+        """
+        from pyiceberg.table import TableProperties
+
+        metadata = self._transaction.table_metadata
+        default_max_ref_age_ms = property_as_int(
+            metadata.properties,
+            TableProperties.MAX_REF_AGE_MS,
+            TableProperties.MAX_REF_AGE_MS_DEFAULT,
+        )
+        now_ms = datetime_to_millis(datetime.now(timezone.utc))
+
+        for ref_name, ref in metadata.refs.items():
+            if ref_name == MAIN_BRANCH:
+                continue
+
+            snapshot = metadata.snapshot_by_id(ref.snapshot_id)
+            if snapshot is None:
+                self._ref_names_to_remove.add(ref_name)
+                continue
+
+            max_ref_age_ms = ref.max_ref_age_ms if ref.max_ref_age_ms is not None else default_max_ref_age_ms
+            if max_ref_age_ms is not None and now_ms - snapshot.timestamp_ms > max_ref_age_ms:
+                self._ref_names_to_remove.add(ref_name)
+
+        return self
+
     def by_id(self, snapshot_id: int) -> ExpireSnapshots:
         """
         Expire a snapshot by its ID.
@@ -1125,9 +1180,7 @@ def older_than(self, dt: datetime) -> ExpireSnapshots:
         Returns:
             This for method chaining.
         """
-        protected_ids = self._get_protected_snapshot_ids()
         expire_from = datetime_to_millis(dt)
-        for snapshot in self._transaction.table_metadata.snapshots:
-            if snapshot.timestamp_ms < expire_from and snapshot.snapshot_id not in protected_ids:
-                self._snapshot_ids_to_expire.add(snapshot.snapshot_id)
+        if self._expire_older_than_ms is None or expire_from > self._expire_older_than_ms:
+            self._expire_older_than_ms = expire_from
         return self

tests/integration/test_snapshot_operations.py

@@ -14,8 +14,10 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import time
 import uuid
 from collections.abc import Generator
+from datetime import datetime, timedelta, timezone
 
 import pyarrow as pa
 import pytest
@@ -332,3 +334,66 @@ 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
+@pytest.mark.parametrize("catalog", [lf("session_catalog_hive"), lf("session_catalog")])
+def test_remove_expired_refs(catalog: Catalog) -> None:
+    """An expired branch is removed, and the snapshot it pinned becomes expirable."""
+    catalog.create_namespace_if_not_exists("default")
+    identifier = f"default.test_expire_refs_{uuid.uuid4().hex[:8]}"
+    arrow_schema = pa.schema([pa.field("id", pa.int64(), nullable=False)])
+    tbl = catalog.create_table(identifier=identifier, schema=arrow_schema)
+
+    for i in range(3):
+        tbl.append(pa.Table.from_pylist([{"id": i}], schema=arrow_schema))
+    tbl = catalog.load_table(identifier)
+
+    pinned_snapshot_id = tbl.metadata.snapshots[0].snapshot_id
+    tbl.manage_snapshots().create_branch(
+        snapshot_id=pinned_snapshot_id,
+        branch_name="audit",
+        max_ref_age_ms=1,
+    ).commit()
+
+    tbl = catalog.load_table(identifier)
+    assert "audit" in tbl.metadata.refs
+    time.sleep(0.05)
+
+    tbl.maintenance.expire_snapshots().remove_expired_refs().older_than(datetime.now(timezone.utc) + timedelta(days=1)).commit()
+
+    tbl = catalog.load_table(identifier)
+    assert "audit" not in tbl.metadata.refs
+    assert tbl.snapshot_by_id(pinned_snapshot_id) is None
+
+    catalog.drop_table(identifier)
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [lf("session_catalog_hive"), lf("session_catalog")])
+def test_remove_expired_refs_keeps_unexpired_branch(catalog: Catalog) -> None:
+    """A branch inside its retention window survives and keeps protecting its snapshot."""
+    catalog.create_namespace_if_not_exists("default")
+    identifier = f"default.test_expire_refs_keep_{uuid.uuid4().hex[:8]}"
+    arrow_schema = pa.schema([pa.field("id", pa.int64(), nullable=False)])
+    tbl = catalog.create_table(identifier=identifier, schema=arrow_schema)
+
+    for i in range(3):
+        tbl.append(pa.Table.from_pylist([{"id": i}], schema=arrow_schema))
+    tbl = catalog.load_table(identifier)
+
+    pinned_snapshot_id = tbl.metadata.snapshots[0].snapshot_id
+    tbl.manage_snapshots().create_branch(
+        snapshot_id=pinned_snapshot_id,
+        branch_name="audit",
+        max_ref_age_ms=7 * 24 * 60 * 60 * 1000,
+    ).commit()
+
+    tbl = catalog.load_table(identifier)
+    tbl.maintenance.expire_snapshots().remove_expired_refs().older_than(datetime.now(timezone.utc) + timedelta(days=1)).commit()
+
+    tbl = catalog.load_table(identifier)
+    assert "audit" in tbl.metadata.refs
+    assert tbl.snapshot_by_id(pinned_snapshot_id) is not None
+
+    catalog.drop_table(identifier)

tests/table/test_expire_snapshots.py

@@ -15,15 +15,21 @@
 # specific language governing permissions and limitations
 # under the License.
 import threading
-from datetime import datetime, timedelta
+import time
+from datetime import datetime, timedelta, timezone
 from unittest.mock import MagicMock, Mock
 from uuid import uuid4
 
+import pyarrow as pa
 import pytest
 
-from pyiceberg.table import CommitTableResponse, Table
+from pyiceberg.catalog import Catalog
+from pyiceberg.io.pyarrow import schema_to_pyarrow
+from pyiceberg.schema import Schema
+from pyiceberg.table import CommitTableResponse, Table, TableProperties
 from pyiceberg.table.update import RemoveSnapshotsUpdate, update_table_metadata
 from pyiceberg.table.update.snapshot import ExpireSnapshots
+from pyiceberg.types import LongType, NestedField
 
 
 def test_cannot_expire_protected_head_snapshot(table_v2: Table) -> None:
@@ -316,3 +322,124 @@ def test_update_remove_snapshots_with_statistics(table_v2_with_statistics: Table
     assert not any(stat.snapshot_id == REMOVE_SNAPSHOT for stat in new_metadata.statistics), (
         "Statistics for removed snapshot should be gone"
     )
+
+
+def _table_with_expired_branch(
+    catalog_with_warehouse: Catalog,
+    max_ref_age_ms: int,
+    namespace: str,
+) -> tuple[Table, int]:
+    """Create a table with three snapshots and a branch on the oldest, then wait out its TTL.
+
+    Returns the reloaded table and the snapshot id the branch pins.
+    """
+    catalog_with_warehouse.create_namespace(namespace)
+    schema = Schema(NestedField(field_id=1, name="id", field_type=LongType(), required=False))
+    table = catalog_with_warehouse.create_table(f"{namespace}.tbl", schema=schema)
+
+    arrow_schema = schema_to_pyarrow(schema, include_field_ids=False)
+    for i in range(3):
+        table.append(pa.table({"id": pa.array([i], type=pa.int64())}, schema=arrow_schema))
+    table = catalog_with_warehouse.load_table(f"{namespace}.tbl")
+
+    pinned_snapshot_id = table.metadata.snapshots[0].snapshot_id
+    table.manage_snapshots().create_branch(
+        snapshot_id=pinned_snapshot_id,
+        branch_name="audit",
+        max_ref_age_ms=max_ref_age_ms,
+    ).commit()
+
+    time.sleep(0.05)
+    return catalog_with_warehouse.load_table(f"{namespace}.tbl"), pinned_snapshot_id
+
+
+def test_remove_expired_refs_removes_expired_branch(catalog_with_warehouse: Catalog) -> None:
+    """An expired branch is dropped, and the snapshot it pinned becomes expirable."""
+    table, pinned_snapshot_id = _table_with_expired_branch(catalog_with_warehouse, namespace="ns_removes", max_ref_age_ms=1)
+    assert "audit" in table.metadata.refs
+
+    table.maintenance.expire_snapshots().remove_expired_refs().older_than(datetime.now(timezone.utc) + timedelta(days=1)).commit()
+
+    table = catalog_with_warehouse.load_table("ns_removes.tbl")
+    assert "audit" not in table.metadata.refs
+    assert table.snapshot_by_id(pinned_snapshot_id) is None
+
+
+@pytest.mark.parametrize("refs_first", [True, False])
+def test_remove_expired_refs_is_order_independent(catalog_with_warehouse: Catalog, refs_first: bool) -> None:
+    """Snapshots released by an expired ref are reclaimed whichever order the builder is chained in."""
+    table, pinned_snapshot_id = _table_with_expired_branch(catalog_with_warehouse, namespace="ns_order", max_ref_age_ms=1)
+    cutoff = datetime.now(timezone.utc) + timedelta(days=1)
+
+    expire = table.maintenance.expire_snapshots()
+    if refs_first:
+        expire.remove_expired_refs().older_than(cutoff).commit()
+    else:
+        expire.older_than(cutoff).remove_expired_refs().commit()
+
+    table = catalog_with_warehouse.load_table("ns_order.tbl")
+    assert "audit" not in table.metadata.refs
+    assert table.snapshot_by_id(pinned_snapshot_id) is None
+
+
+def test_remove_expired_refs_is_opt_in(catalog_with_warehouse: Catalog) -> None:
+    """Without remove_expired_refs(), an expired branch survives and keeps pinning its snapshot."""
+    table, pinned_snapshot_id = _table_with_expired_branch(catalog_with_warehouse, namespace="ns_optin", max_ref_age_ms=1)
+
+    table.maintenance.expire_snapshots().older_than(datetime.now(timezone.utc) + timedelta(days=1)).commit()
+
+    table = catalog_with_warehouse.load_table("ns_optin.tbl")
+    assert "audit" in table.metadata.refs
+    assert table.snapshot_by_id(pinned_snapshot_id) is not None
+
+
+def test_remove_expired_refs_keeps_unexpired_branch(catalog_with_warehouse: Catalog) -> None:
+    """A branch still inside its retention window is kept, and keeps protecting its snapshot."""
+    table, pinned_snapshot_id = _table_with_expired_branch(
+        catalog_with_warehouse, namespace="ns_keeps", max_ref_age_ms=7 * 24 * 60 * 60 * 1000
+    )
+
+    table.maintenance.expire_snapshots().remove_expired_refs().older_than(datetime.now(timezone.utc) + timedelta(days=1)).commit()
+
+    table = catalog_with_warehouse.load_table("ns_keeps.tbl")
+    assert "audit" in table.metadata.refs
+    assert table.snapshot_by_id(pinned_snapshot_id) is not None
+
+
+def test_remove_expired_refs_falls_back_to_table_property(catalog_with_warehouse: Catalog) -> None:
+    """A ref without its own TTL uses history.expire.max-ref-age-ms."""
+    catalog_with_warehouse.create_namespace("ns_prop")
+    schema = Schema(NestedField(field_id=1, name="id", field_type=LongType(), required=False))
+    table = catalog_with_warehouse.create_table("ns_prop.tbl", schema=schema, properties={TableProperties.MAX_REF_AGE_MS: "1"})
+    arrow_schema = schema_to_pyarrow(schema, include_field_ids=False)
+    table.append(pa.table({"id": pa.array([1], type=pa.int64())}, schema=arrow_schema))
+    table = catalog_with_warehouse.load_table("ns_prop.tbl")
+
+    snapshot_id = table.metadata.current_snapshot_id
+    assert snapshot_id is not None
+    table.manage_snapshots().create_tag(snapshot_id=snapshot_id, tag_name="release").commit()
+    table = catalog_with_warehouse.load_table("ns_prop.tbl")
+    assert table.metadata.refs["release"].max_ref_age_ms is None
+
+    time.sleep(0.05)
+    table.maintenance.expire_snapshots().remove_expired_refs().commit()
+
+    table = catalog_with_warehouse.load_table("ns_prop.tbl")
+    assert "release" not in table.metadata.refs
+
+
+def test_remove_expired_refs_never_removes_main(catalog_with_warehouse: Catalog) -> None:
+    """main is exempt from ref expiry even when the table-level TTL has elapsed."""
+    catalog_with_warehouse.create_namespace("ns_main")
+    schema = Schema(NestedField(field_id=1, name="id", field_type=LongType(), required=False))
+    table = catalog_with_warehouse.create_table("ns_main.tbl", schema=schema, properties={TableProperties.MAX_REF_AGE_MS: "1"})
+    arrow_schema = schema_to_pyarrow(schema, include_field_ids=False)
+    table.append(pa.table({"id": pa.array([1], type=pa.int64())}, schema=arrow_schema))
+    table = catalog_with_warehouse.load_table("ns_main.tbl")
+
+    time.sleep(0.05)
+    table.maintenance.expire_snapshots().remove_expired_refs().commit()
+
+    table = catalog_with_warehouse.load_table("ns_main.tbl")
+    assert "main" in table.metadata.refs
+    assert table.metadata.current_snapshot_id is not None