diff --git a/mkdocs/docs/api.md b/mkdocs/docs/api.md index 0a2a72b6f0..8fcb63198d 100644 --- a/mkdocs/docs/api.md +++ b/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"`: diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index c0adce84dc..94944f8e1a 100644 --- a/pyiceberg/table/__init__.py +++ b/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 diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 7931edacdd..90ff410f3d 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/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 diff --git a/tests/integration/test_snapshot_operations.py b/tests/integration/test_snapshot_operations.py index 07fb77edbb..0eeac60e31 100644 --- a/tests/integration/test_snapshot_operations.py +++ b/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) diff --git a/tests/table/test_expire_snapshots.py b/tests/table/test_expire_snapshots.py index 106e5b786c..c62aff08fd 100644 --- a/tests/table/test_expire_snapshots.py +++ b/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