Refactor: add remove_expired_refs method

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

Author: alessandro-nori

pyiceberg/table/update/snapshot.py

@@ -1025,110 +1025,71 @@ def _current_ancestors(self) -> set[int]:
 
 
 class ExpireSnapshots(UpdateTableMetadata["ExpireSnapshots"]):
-    """Expire snapshots by ID.
+    """Expire snapshots and refs.
 
     Use table.expire_snapshots().<operation>().commit() to run a specific operation.
     Use table.expire_snapshots().<operation-one>().<operation-two>().commit() to run multiple operations.
-    Pending changes are applied on commit.
+    Pending changes are applied on commit. Call order does not affect the result.
     """
 
     _updates: tuple[TableUpdate, ...]
     _requirements: tuple[TableRequirement, ...]
     _snapshot_ids_to_expire: set[int]
+    _ref_names_to_expire: 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_expire = set()
+        self._expire_older_than_ms = None
 
     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.
+        Applies all pending expirations: explicit snapshot IDs, age-based snapshot expiry,
+        and ref removals. Protected snapshots (branch/tag heads not being expired) are always
+        excluded. The age threshold from older_than() is evaluated here so that call order
+        with remove_expired_refs() does not affect the result.
 
         Returns:
             Tuple of updates and requirements to be committed,
             as required by the calling parent apply functions.
         """
-        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)
+        protected_ids = self._get_protected_snapshot_ids()
 
-        # 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)
+        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)
 
         snapshot_ids_to_expire = self._snapshot_ids_to_expire - protected_ids
 
         updates: list[TableUpdate] = list(self._updates)
-        for ref_name in expired_ref_names:
+        for ref_name in self._ref_names_to_expire:
             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 _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]:
+    def _get_protected_snapshot_ids(self) -> set[int]:
         """
-        Get the IDs of protected snapshots.
+        Get the IDs of snapshots that must not be expired.
 
-        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.
+        These are the HEAD snapshots of all branches and tags that are not
+        already marked for removal via remove_expired_refs().
 
         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 name, ref in self._transaction.table_metadata.refs.items()
-            if name not in expired_ref_names
+            if name not in self._ref_names_to_expire
         }
 
     def by_id(self, snapshot_id: int) -> ExpireSnapshots:
@@ -1169,17 +1130,46 @@ def by_ids(self, snapshot_ids: list[int]) -> ExpireSnapshots:
 
     def older_than(self, dt: datetime) -> ExpireSnapshots:
         """
-        Expire all unprotected snapshots with a timestamp older than a given value.
+        Expire all unprotected snapshots with a timestamp older than the given value.
+
+        The filter is evaluated at commit time so that snapshots left without a ref
+        by remove_expired_refs() are also considered, regardless of call order.
 
         Args:
-            dt (datetime): Only snapshots with datetime < this value will be expired.
+            dt (datetime): Only snapshots with timestamp < this value will be expired.
 
         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)
+        self._expire_older_than_ms = datetime_to_millis(dt)
+        return self
+
+    def remove_expired_refs(self, default_max_ref_age_ms: int | None = None) -> ExpireSnapshots:
+        """
+        Mark stale branches and tags for removal.
+
+        A ref is expired when the age of its snapshot exceeds its own max_ref_age_ms.
+        If a ref has no per-ref max_ref_age_ms set, default_max_ref_age_ms is used as fallback.
+        The main branch is never removed.
+
+        Snapshots left without any live ref after this call are no longer protected,
+        so a subsequent older_than() will include them in age-based expiry.
+
+        Args:
+            default_max_ref_age_ms: Fallback max age in milliseconds for refs that have no
+                per-ref max_ref_age_ms configured. If None, such refs are skipped.
+
+        Returns:
+            This for method chaining.
+        """
+        now_ms = int(datetime.now().timestamp() * 1000)
+        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 default_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 or (now_ms - snapshot.timestamp_ms) > effective_max_ref_age_ms:
+                self._ref_names_to_expire.add(name)
         return self

tests/table/test_expire_snapshots.py

@@ -323,7 +323,8 @@ def _make_commit_response(table: Table) -> CommitTableResponse:
 
 
 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."""
+    """A tag whose snapshot age exceeds max_ref_age_ms is removed; its orphaned snapshot
+    is also expired when older_than() is combined."""
     table_v2.catalog = MagicMock()
     table_v2.catalog.commit_table.return_value = _make_commit_response(table_v2)
 
@@ -332,7 +333,8 @@ def test_ref_expiration_removes_old_tag_and_snapshot(table_v2: Table) -> None:
     assert "test" in table_v2.metadata.refs
     assert table_v2.metadata.refs["test"].snapshot_id == OLD_SNAPSHOT
 
-    table_v2.maintenance.expire_snapshots().commit()
+    future = datetime.now() + timedelta(days=1)
+    table_v2.maintenance.expire_snapshots().remove_expired_refs().older_than(future).commit()
 
     args, _ = table_v2.catalog.commit_table.call_args
     updates = args[2]
@@ -354,17 +356,15 @@ def test_ref_expiration_removes_old_branch(table_v2: Table) -> None:
     table_v2.metadata = table_v2.metadata.model_copy(
         update={
             "refs": {
-                "main": SnapshotRef(
-                    **{"snapshot-id": CURRENT_SNAPSHOT, "type": SnapshotRefType.BRANCH}
-                ),
+                "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()
+    table_v2.maintenance.expire_snapshots().remove_expired_refs().commit()
 
     args, _ = table_v2.catalog.commit_table.call_args
     updates = args[2]
@@ -376,9 +376,7 @@ def test_ref_expiration_removes_old_branch(table_v2: Table) -> None:
 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": {
@@ -389,41 +387,33 @@ def test_main_branch_never_expires(table_v2: Table) -> None:
         }
     )
 
-    table_v2.maintenance.expire_snapshots().commit()
+    table_v2.maintenance.expire_snapshots().remove_expired_refs().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."""
+def test_remove_expired_refs_default_max_ref_age_ms(table_v2: Table) -> None:
+    """A ref without per-ref max_ref_age_ms uses the method parameter 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"},
+                "main": SnapshotRef(**{"snapshot-id": CURRENT_SNAPSHOT, "type": SnapshotRefType.BRANCH}),
+                # no per-ref max_ref_age_ms
+                "old-tag": SnapshotRef(**{"snapshot-id": OLD_SNAPSHOT, "type": SnapshotRefType.TAG}),
+            }
         }
     )
 
-    table_v2.maintenance.expire_snapshots().commit()
+    # 1 ms default — old-tag's snapshot (from 2018) will exceed it
+    table_v2.maintenance.expire_snapshots().remove_expired_refs(default_max_ref_age_ms=1).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 any(u.ref_name == "old-tag" for u in ref_updates)
     assert not any(u.ref_name == "main" for u in ref_updates)
 
 
@@ -448,12 +438,43 @@ def test_young_ref_is_retained(table_v2: Table) -> None:
         }
     )
 
-    table_v2.maintenance.expire_snapshots().commit()
+    table_v2.maintenance.expire_snapshots().remove_expired_refs().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"
+
+
+def test_remove_expired_refs_and_older_than_order_independent(table_v2: Table) -> None:
+    """remove_expired_refs().older_than() and older_than().remove_expired_refs() produce the same result."""
+    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-tag": SnapshotRef(
+                    **{"snapshot-id": OLD_SNAPSHOT, "type": SnapshotRefType.TAG, "max-ref-age-ms": 1}
+                ),
+            }
+        }
+    )
+
+    future = datetime.now() + timedelta(days=1)
+
+    def _get_expired_ids(order: str) -> set[int]:
+        table_v2.catalog.reset_mock()
+        expire = table_v2.maintenance.expire_snapshots()
+        if order == "refs_first":
+            expire.remove_expired_refs().older_than(future).commit()
+        else:
+            expire.older_than(future).remove_expired_refs().commit()
+        args, _ = table_v2.catalog.commit_table.call_args
+        snap_updates = [u for u in args[2] if isinstance(u, RemoveSnapshotsUpdate)]
+        return {sid for u in snap_updates for sid in u.snapshot_ids}
+
+    assert _get_expired_ids("refs_first") == _get_expired_ids("snapshots_first")