Expire stale refs in ExpireSnapshots

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: alessandro-nori

pyiceberg/table/__init__.py

@@ -205,6 +205,8 @@ 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"
+
 
 class Transaction:
     _table: Table

pyiceberg/table/update/snapshot.py

@@ -1047,31 +1047,88 @@ def _commit(self) -> UpdatesAndRequirements:
         Commit the staged updates and requirements.
 
         This will remove the snapshots with the given IDs, but will always skip protected snapshots (branch/tag heads).
+        Refs (branches/tags) whose snapshot age exceeds their configured max-ref-age-ms are also removed.
 
         Returns:
             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
-        protected_ids = self._get_protected_snapshot_ids()
-        self._snapshot_ids_to_expire -= protected_ids
-        update = RemoveSnapshotsUpdate(snapshot_ids=self._snapshot_ids_to_expire)
-        self._updates += (update,)
+        now_ms = int(datetime.now().timestamp() * 1000)
+        expired_ref_names = self._compute_expired_refs(now_ms)
+        protected_ids = self._get_protected_snapshot_ids(expired_ref_names)
+
+        # Snapshots exclusively referenced by expired refs are also eligible for expiration
+        for ref_name in expired_ref_names:
+            ref = self._transaction.table_metadata.refs[ref_name]
+            if ref.snapshot_id not in protected_ids:
+                self._snapshot_ids_to_expire.add(ref.snapshot_id)
+
+        snapshot_ids_to_expire = self._snapshot_ids_to_expire - protected_ids
+
+        updates: list[TableUpdate] = list(self._updates)
+        for ref_name in expired_ref_names:
+            updates.append(RemoveSnapshotRefUpdate(ref_name=ref_name))
+        if snapshot_ids_to_expire:
+            updates.append(RemoveSnapshotsUpdate(snapshot_ids=snapshot_ids_to_expire))
+        self._updates = tuple(updates)
         return self._updates, self._requirements
 
-    def _get_protected_snapshot_ids(self) -> set[int]:
+    def _compute_expired_refs(self, now_ms: int) -> set[str]:
+        """
+        Compute the set of ref names (branches/tags) that should be expired.
+
+        A ref is expired when the age of its snapshot exceeds:
+        - the ref's own max_ref_age_ms, or
+        - the table property history.expire.max-ref-age-ms, if the ref has no per-ref setting.
+        The main branch is never expired. Refs with no effective max-ref-age configuration are skipped.
+
+        Args:
+            now_ms: Current time in milliseconds.
+
+        Returns:
+            Set of ref names to remove.
+        """
+        from pyiceberg.table import TableProperties
+
+        props = self._transaction.table_metadata.properties
+        table_max_ref_age_ms: int | None = (
+            int(props[TableProperties.MAX_REF_AGE_MS]) if TableProperties.MAX_REF_AGE_MS in props else None
+        )
+
+        expired: set[str] = set()
+        for name, ref in self._transaction.table_metadata.refs.items():
+            if name == MAIN_BRANCH:
+                continue
+            effective_max_ref_age_ms = ref.max_ref_age_ms if ref.max_ref_age_ms is not None else table_max_ref_age_ms
+            if effective_max_ref_age_ms is None:
+                continue
+            snapshot = self._transaction.table_metadata.snapshot_by_id(ref.snapshot_id)
+            if snapshot is None:
+                expired.add(name)
+                continue
+            if (now_ms - snapshot.timestamp_ms) > effective_max_ref_age_ms:
+                expired.add(name)
+        return expired
+
+    def _get_protected_snapshot_ids(self, expired_ref_names: set[str] | None = None) -> 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.
+        These are the HEAD snapshots of all branches and all tagged snapshots that are NOT being expired.
+        Snapshots exclusively pointed to by expired refs are not protected.
+
+        Args:
+            expired_ref_names: Set of ref names that are being expired. Defaults to empty set.
 
         Returns:
             Set of protected snapshot IDs to exclude from expiration.
         """
+        if expired_ref_names is None:
+            expired_ref_names = set()
         return {
             ref.snapshot_id
-            for ref in self._transaction.table_metadata.refs.values()
-            if ref.snapshot_ref_type in [SnapshotRefType.TAG, SnapshotRefType.BRANCH]
+            for name, ref in self._transaction.table_metadata.refs.items()
+            if name not in expired_ref_names
         }
 
     def by_id(self, snapshot_id: int) -> ExpireSnapshots:

tests/table/test_expire_snapshots.py

@@ -22,7 +22,8 @@
 import pytest
 
 from pyiceberg.table import CommitTableResponse, Table
-from pyiceberg.table.update import RemoveSnapshotsUpdate, update_table_metadata
+from pyiceberg.table.refs import SnapshotRef, SnapshotRefType
+from pyiceberg.table.update import RemoveSnapshotRefUpdate, RemoveSnapshotsUpdate, update_table_metadata
 from pyiceberg.table.update.snapshot import ExpireSnapshots
 
 
@@ -92,8 +93,8 @@ def test_expire_unprotected_snapshot(table_v2: Table) -> None:
     table_v2.metadata = table_v2.metadata.model_copy(
         update={
             "refs": {
-                "main": MagicMock(snapshot_id=KEEP_SNAPSHOT, snapshot_ref_type="branch"),
-                "tag1": MagicMock(snapshot_id=KEEP_SNAPSHOT, snapshot_ref_type="tag"),
+                "main": SnapshotRef(**{"snapshot-id": KEEP_SNAPSHOT, "type": SnapshotRefType.BRANCH}),
+                "tag1": SnapshotRef(**{"snapshot-id": KEEP_SNAPSHOT, "type": SnapshotRefType.TAG}),
             }
         }
     )
@@ -134,8 +135,8 @@ def test_expire_snapshots_by_timestamp_skips_protected(table_v2: Table) -> None:
     table_v2.metadata = table_v2.metadata.model_copy(
         update={
             "refs": {
-                "main": MagicMock(snapshot_id=HEAD_SNAPSHOT, snapshot_ref_type="branch"),
-                "mytag": MagicMock(snapshot_id=TAGGED_SNAPSHOT, snapshot_ref_type="tag"),
+                "main": SnapshotRef(**{"snapshot-id": HEAD_SNAPSHOT, "type": SnapshotRefType.BRANCH}),
+                "mytag": SnapshotRef(**{"snapshot-id": TAGGED_SNAPSHOT, "type": SnapshotRefType.TAG}),
             },
             "snapshots": [
                 SimpleNamespace(snapshot_id=HEAD_SNAPSHOT, timestamp_ms=1, parent_snapshot_id=None),
@@ -165,13 +166,8 @@ def test_expire_snapshots_by_timestamp_skips_protected(table_v2: Table) -> None:
     assert HEAD_SNAPSHOT in remaining_ids
     assert TAGGED_SNAPSHOT in remaining_ids
 
-    # No snapshots should have been expired (commit_table called, but with empty snapshot_ids)
-    args, kwargs = table_v2.catalog.commit_table.call_args
-    updates = args[2] if len(args) > 2 else ()
-    # Find RemoveSnapshotsUpdate in updates
-    remove_update = next((u for u in updates if getattr(u, "action", None) == "remove-snapshots"), None)
-    assert remove_update is not None
-    assert remove_update.snapshot_ids == []
+    # No snapshots expired and no refs expired — commit_table should not be called at all
+    table_v2.catalog.commit_table.assert_not_called()
 
 
 def test_expire_snapshots_by_ids(table_v2: Table) -> None:
@@ -188,24 +184,14 @@ def test_expire_snapshots_by_ids(table_v2: Table) -> None:
     table_v2.catalog = MagicMock()
     table_v2.catalog.commit_table.return_value = mock_response
 
-    # Remove any refs that protect the snapshots to be expired
-    table_v2.metadata = table_v2.metadata.model_copy(
-        update={
-            "refs": {
-                "main": MagicMock(snapshot_id=KEEP_SNAPSHOT, snapshot_ref_type="branch"),
-                "tag1": MagicMock(snapshot_id=KEEP_SNAPSHOT, snapshot_ref_type="tag"),
-            }
-        }
-    )
-
     # Add snapshots to metadata for multi-id test
     from types import SimpleNamespace
 
     table_v2.metadata = table_v2.metadata.model_copy(
         update={
             "refs": {
-                "main": MagicMock(snapshot_id=KEEP_SNAPSHOT, snapshot_ref_type="branch"),
-                "tag1": MagicMock(snapshot_id=KEEP_SNAPSHOT, snapshot_ref_type="tag"),
+                "main": SnapshotRef(**{"snapshot-id": KEEP_SNAPSHOT, "type": SnapshotRefType.BRANCH}),
+                "tag1": SnapshotRef(**{"snapshot-id": KEEP_SNAPSHOT, "type": SnapshotRefType.TAG}),
             },
             "snapshots": [
                 SimpleNamespace(snapshot_id=EXPIRE_SNAPSHOT_1, timestamp_ms=1, parent_snapshot_id=None),
@@ -316,3 +302,158 @@ 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"
     )
+
+
+# --- Ref expiration tests ---
+# The table_v2 fixture has two snapshots:
+#   3051729675574597004  (timestamp_ms=1515100955770, ~Jan 2018)
+#   3055729675574597004  (timestamp_ms=1555100955770, ~Apr 2019, current/main)
+# And a "test" tag pointing to 3051729675574597004 with max-ref-age-ms=10000000 (~2.7 h).
+
+OLD_SNAPSHOT = 3051729675574597004
+CURRENT_SNAPSHOT = 3055729675574597004
+
+
+def _make_commit_response(table: Table) -> CommitTableResponse:
+    return CommitTableResponse(
+        metadata=table.metadata,
+        metadata_location="mock://metadata/location",
+        uuid=uuid4(),
+    )
+
+
+def test_ref_expiration_removes_old_tag_and_snapshot(table_v2: Table) -> None:
+    """A tag whose snapshot age exceeds max_ref_age_ms is removed along with its snapshot."""
+    table_v2.catalog = MagicMock()
+    table_v2.catalog.commit_table.return_value = _make_commit_response(table_v2)
+
+    # "test" tag (fixture) points to OLD_SNAPSHOT with max-ref-age-ms=10000000 (~2.7 h).
+    # OLD_SNAPSHOT timestamp is from 2018 — definitely older than 2.7 h.
+    assert "test" in table_v2.metadata.refs
+    assert table_v2.metadata.refs["test"].snapshot_id == OLD_SNAPSHOT
+
+    table_v2.maintenance.expire_snapshots().commit()
+
+    args, _ = table_v2.catalog.commit_table.call_args
+    updates = args[2]
+
+    ref_updates = [u for u in updates if isinstance(u, RemoveSnapshotRefUpdate)]
+    snap_updates = [u for u in updates if isinstance(u, RemoveSnapshotsUpdate)]
+
+    assert any(u.ref_name == "test" for u in ref_updates), "Expected 'test' tag to be removed"
+    assert any(OLD_SNAPSHOT in u.snapshot_ids for u in snap_updates), (
+        "Expected OLD_SNAPSHOT to be removed since it is no longer referenced"
+    )
+
+
+def test_ref_expiration_removes_old_branch(table_v2: Table) -> None:
+    """A non-main branch whose snapshot age exceeds max_ref_age_ms is removed."""
+    table_v2.catalog = MagicMock()
+    table_v2.catalog.commit_table.return_value = _make_commit_response(table_v2)
+
+    table_v2.metadata = table_v2.metadata.model_copy(
+        update={
+            "refs": {
+                "main": SnapshotRef(
+                    **{"snapshot-id": CURRENT_SNAPSHOT, "type": SnapshotRefType.BRANCH}
+                ),
+                "stale-branch": SnapshotRef(
+                    **{"snapshot-id": OLD_SNAPSHOT, "type": SnapshotRefType.BRANCH, "max-ref-age-ms": 1}
+                ),
+            }
+        }
+    )
+
+    table_v2.maintenance.expire_snapshots().commit()
+
+    args, _ = table_v2.catalog.commit_table.call_args
+    updates = args[2]
+    ref_updates = [u for u in updates if isinstance(u, RemoveSnapshotRefUpdate)]
+    assert any(u.ref_name == "stale-branch" for u in ref_updates)
+    assert not any(u.ref_name == "main" for u in ref_updates)
+
+
+def test_main_branch_never_expires(table_v2: Table) -> None:
+    """main branch is never removed regardless of age or max_ref_age_ms."""
+    table_v2.catalog = MagicMock()
+    table_v2.catalog.commit_table.return_value = _make_commit_response(table_v2)
+
+    # Only main exists, with max_ref_age_ms=1 (would expire anything else immediately)
+    table_v2.metadata = table_v2.metadata.model_copy(
+        update={
+            "refs": {
+                "main": SnapshotRef(
+                    **{"snapshot-id": CURRENT_SNAPSHOT, "type": SnapshotRefType.BRANCH, "max-ref-age-ms": 1}
+                ),
+            }
+        }
+    )
+
+    table_v2.maintenance.expire_snapshots().commit()
+
+    # commit_table should NOT be called — main is never expired, so nothing changes
+    table_v2.catalog.commit_table.assert_not_called()
+
+
+def test_table_property_max_ref_age_ms_used_as_default(table_v2: Table) -> None:
+    """A ref without per-ref max_ref_age_ms uses the table property as fallback."""
+    table_v2.catalog = MagicMock()
+    table_v2.catalog.commit_table.return_value = _make_commit_response(table_v2)
+
+    # Ref has no per-ref max_ref_age_ms
+    table_v2.metadata = table_v2.metadata.model_copy(
+        update={
+            "refs": {
+                "main": SnapshotRef(
+                    **{"snapshot-id": CURRENT_SNAPSHOT, "type": SnapshotRefType.BRANCH}
+                ),
+                "old-tag": SnapshotRef(
+                    **{"snapshot-id": OLD_SNAPSHOT, "type": SnapshotRefType.TAG}
+                ),
+            },
+            # Table-level default: 1 ms — everything is expired
+            "properties": {"history.expire.max-ref-age-ms": "1"},
+        }
+    )
+
+    table_v2.maintenance.expire_snapshots().commit()
+
+    args, _ = table_v2.catalog.commit_table.call_args
+    updates = args[2]
+    ref_updates = [u for u in updates if isinstance(u, RemoveSnapshotRefUpdate)]
+    assert any(u.ref_name == "old-tag" for u in ref_updates), (
+        "Expected old-tag to be expired via table property default"
+    )
+    assert not any(u.ref_name == "main" for u in ref_updates)
+
+
+def test_young_ref_is_retained(table_v2: Table) -> None:
+    """A ref whose snapshot is within max_ref_age_ms is not removed."""
+    table_v2.catalog = MagicMock()
+    table_v2.catalog.commit_table.return_value = _make_commit_response(table_v2)
+
+    # fresh-tag has a huge max_ref_age_ms — it should never expire
+    # stale-tag has max_ref_age_ms=1 — it will be expired (triggers a commit)
+    table_v2.metadata = table_v2.metadata.model_copy(
+        update={
+            "refs": {
+                "main": SnapshotRef(**{"snapshot-id": CURRENT_SNAPSHOT, "type": SnapshotRefType.BRANCH}),
+                "fresh-tag": SnapshotRef(
+                    **{"snapshot-id": OLD_SNAPSHOT, "type": SnapshotRefType.TAG, "max-ref-age-ms": 9999999999999}
+                ),
+                "stale-tag": SnapshotRef(
+                    **{"snapshot-id": OLD_SNAPSHOT, "type": SnapshotRefType.TAG, "max-ref-age-ms": 1}
+                ),
+            }
+        }
+    )
+
+    table_v2.maintenance.expire_snapshots().commit()
+
+    # stale-tag causes commit_table to be called; fresh-tag must not be in the expired list
+    table_v2.catalog.commit_table.assert_called_once()
+    args, _ = table_v2.catalog.commit_table.call_args
+    updates = args[2]
+    ref_updates = [u for u in updates if isinstance(u, RemoveSnapshotRefUpdate)]
+    assert any(u.ref_name == "stale-tag" for u in ref_updates), "stale-tag should be expired"
+    assert not any(u.ref_name == "fresh-tag" for u in ref_updates), "fresh-tag must not be expired"