feat: Add support for rolling back to timestamp

Co-authored-by: Chinmay Bhat <12948588+chinmay-bhat@users.noreply.github.com>

Author: geruh

pyiceberg/table/snapshots.py

@@ -472,3 +472,24 @@ def ancestors_between(from_snapshot: Snapshot | None, to_snapshot: Snapshot, tab
                 break
     else:
         yield from ancestors_of(to_snapshot, table_metadata)
+
+
+def latest_ancestor_before_timestamp(table_metadata: TableMetadata, timestamp_ms: int) -> Snapshot | None:
+    """Find the latest ancestor snapshot whose timestamp is before the provided timestamp.
+
+    Args:
+        table_metadata: The table metadata for a table
+        timestamp_ms: lookup snapshots before this timestamp
+
+    Returns:
+        The latest ancestor snapshot older than the timestamp, or None if not found.
+    """
+    result: Snapshot | None = None
+    result_timestamp: int = 0
+
+    for ancestor in ancestors_of(table_metadata.current_snapshot(), table_metadata):
+        if timestamp_ms > ancestor.timestamp_ms > result_timestamp:
+            result = ancestor
+            result_timestamp = ancestor.timestamp_ms
+
+    return result

pyiceberg/table/update/snapshot.py

@@ -65,6 +65,7 @@
     SnapshotSummaryCollector,
     Summary,
     ancestors_of,
+    latest_ancestor_before_timestamp,
     update_snapshot_summaries,
 )
 from pyiceberg.table.update import (
@@ -1007,6 +1008,33 @@ def rollback_to_snapshot(self, snapshot_id: int) -> ManageSnapshots:
 
         return self.set_current_snapshot(snapshot_id=snapshot_id)
 
+    def rollback_to_timestamp(self, timestamp_ms: int) -> ManageSnapshots:
+        """Rollback the table to the latest snapshot before the given timestamp.
+
+        Finds the latest ancestor snapshot whose timestamp is before the given timestamp and rolls back to it.
+
+        Args:
+            timestamp_ms: Rollback to the latest snapshot before this timestamp in milliseconds.
+        Returns:
+            This for method chaining
+        Raises:
+            ValueError: If no valid snapshot exists older than the given timestamp.
+        """
+        self._commit_if_ref_updates_exist()
+
+        snapshot = latest_ancestor_before_timestamp(self._transaction.table_metadata, timestamp_ms)
+        if snapshot is None:
+            raise ValueError(f"Cannot roll back, no valid snapshot older than: {timestamp_ms}")
+
+        update, requirement = self._transaction._set_ref_snapshot(
+            snapshot_id=snapshot.snapshot_id,
+            ref_name=MAIN_BRANCH,
+            type=str(SnapshotRefType.BRANCH),
+        )
+        self._updates += update
+        self._requirements += requirement
+        return self
+
     def _is_current_ancestor(self, snapshot_id: int) -> bool:
         return snapshot_id in self._current_ancestors()
 

tests/integration/test_snapshot_operations.py

@@ -186,3 +186,32 @@ def test_rollback_to_snapshot(catalog: Catalog) -> None:
     tbl = catalog.load_table(identifier)
     restored_snapshot = tbl.current_snapshot()
     assert restored_snapshot and restored_snapshot.snapshot_id == current_snapshot_id
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
+def test_rollback_to_timestamp(catalog: Catalog) -> None:
+    identifier = "default.test_table_snapshot_operations"
+    tbl = catalog.load_table(identifier)
+
+    current_snapshot = tbl.current_snapshot()
+    assert current_snapshot is not None
+    assert current_snapshot.parent_snapshot_id is not None
+
+    parent_snapshot = tbl.metadata.snapshot_by_id(current_snapshot.parent_snapshot_id)
+    assert parent_snapshot is not None
+
+    # rollback_to_timestamp finds the latest ancestor with timestamp less than given timestamp
+    tbl.manage_snapshots().rollback_to_timestamp(timestamp_ms=current_snapshot.timestamp_ms).commit()
+
+    tbl = catalog.load_table(identifier)
+    updated_snapshot = tbl.current_snapshot()
+    assert updated_snapshot is not None
+    assert updated_snapshot.snapshot_id == parent_snapshot.snapshot_id
+
+    # restore table
+    tbl.manage_snapshots().set_current_snapshot(snapshot_id=current_snapshot.snapshot_id).commit()
+    tbl = catalog.load_table(identifier)
+    restored_snapshot = tbl.current_snapshot()
+    assert restored_snapshot is not None
+    assert restored_snapshot.snapshot_id == current_snapshot.snapshot_id

tests/table/test_manage_snapshots.py

@@ -303,3 +303,78 @@ def test_rollback_to_current_snapshot(table_v2: Table) -> None:
 
     table_v2.manage_snapshots().rollback_to_snapshot(snapshot_id=current_snapshot.snapshot_id).commit()
     table_v2.catalog.commit_table.assert_called_once()
+
+
+def test_rollback_to_timestamp() -> None:
+    from pyiceberg.table.metadata import TableMetadataV2
+
+    metadata_dict = {
+        "format-version": 2,
+        "table-uuid": "9c12d441-03fe-4693-9a96-a0705ddf69c1",
+        "location": "s3://bucket/test/location",
+        "last-sequence-number": 4,
+        "last-updated-ms": 1602638573590,
+        "last-column-id": 1,
+        "current-schema-id": 0,
+        "schemas": [{"type": "struct", "schema-id": 0, "fields": [{"id": 1, "name": "x", "required": True, "type": "long"}]}],
+        "default-spec-id": 0,
+        "partition-specs": [{"spec-id": 0, "fields": []}],
+        "last-partition-id": 999,
+        "default-sort-order-id": 0,
+        "sort-orders": [{"order-id": 0, "fields": []}],
+        "current-snapshot-id": 4,
+        "snapshots": [
+            {"snapshot-id": 1, "timestamp-ms": 1000, "sequence-number": 1, "manifest-list": "s3://a/1.avro"},
+            {
+                "snapshot-id": 2,
+                "parent-snapshot-id": 1,
+                "timestamp-ms": 2000,
+                "sequence-number": 2,
+                "manifest-list": "s3://a/2.avro",
+            },
+            {
+                "snapshot-id": 3,
+                "parent-snapshot-id": 2,
+                "timestamp-ms": 3000,
+                "sequence-number": 3,
+                "manifest-list": "s3://a/3.avro",
+            },
+            {
+                "snapshot-id": 4,
+                "parent-snapshot-id": 3,
+                "timestamp-ms": 4000,
+                "sequence-number": 4,
+                "manifest-list": "s3://a/4.avro",
+            },
+        ],
+    }
+
+    mock_catalog = MagicMock()
+    table = Table(
+        identifier=("db", "table"),
+        metadata=TableMetadataV2(**metadata_dict),
+        metadata_location="s3://bucket/test/metadata.json",
+        io=load_file_io(),
+        catalog=mock_catalog,
+    )
+    mock_catalog.commit_table.return_value = _mock_commit_response(table)
+
+    # verify we find the ancestor before timestamp 2500
+    table.manage_snapshots().rollback_to_timestamp(timestamp_ms=2500).commit()
+
+    updates = _get_updates(mock_catalog)
+    set_ref_updates = [u for u in updates if isinstance(u, SetSnapshotRefUpdate)]
+
+    assert len(set_ref_updates) == 1
+    assert set_ref_updates[0].snapshot_id == 2
+    assert set_ref_updates[0].ref_name == "main"
+
+
+def test_rollback_to_timestamp_no_valid_snapshot(table_v2: Table) -> None:
+    # The oldest snapshot is at timestamp 1515100955770
+    table_v2.catalog = MagicMock()
+
+    with pytest.raises(ValueError, match="Cannot roll back, no valid snapshot older than"):
+        table_v2.manage_snapshots().rollback_to_timestamp(timestamp_ms=1515100955770).commit()
+
+    table_v2.catalog.commit_table.assert_not_called()

tests/table/test_snapshots.py

@@ -30,6 +30,7 @@
     Summary,
     ancestors_between,
     ancestors_of,
+    latest_ancestor_before_timestamp,
     update_snapshot_summaries,
 )
 from pyiceberg.transforms import IdentityTransform
@@ -456,3 +457,79 @@ def test_ancestors_between(table_v2_with_extensive_snapshots: Table) -> None:
         )
         == 2000
     )
+
+
+def test_latest_ancestor_before_timestamp() -> None:
+    from pyiceberg.table.metadata import TableMetadataV2
+
+    # Create metadata with 4 snapshots at ordered timestamps
+    metadata = TableMetadataV2(
+        **{
+            "format-version": 2,
+            "table-uuid": "9c12d441-03fe-4693-9a96-a0705ddf69c1",
+            "location": "s3://bucket/test/location",
+            "last-sequence-number": 4,
+            "last-updated-ms": 1602638573590,
+            "last-column-id": 1,
+            "current-schema-id": 0,
+            "schemas": [{"type": "struct", "schema-id": 0, "fields": [{"id": 1, "name": "x", "required": True, "type": "long"}]}],
+            "default-spec-id": 0,
+            "partition-specs": [{"spec-id": 0, "fields": []}],
+            "last-partition-id": 999,
+            "default-sort-order-id": 0,
+            "sort-orders": [{"order-id": 0, "fields": []}],
+            "current-snapshot-id": 4,
+            "snapshots": [
+                {
+                    "snapshot-id": 1,
+                    "timestamp-ms": 1000,
+                    "sequence-number": 1,
+                    "summary": {"operation": "append"},
+                    "manifest-list": "s3://a/1.avro",
+                },
+                {
+                    "snapshot-id": 2,
+                    "parent-snapshot-id": 1,
+                    "timestamp-ms": 2000,
+                    "sequence-number": 2,
+                    "summary": {"operation": "append"},
+                    "manifest-list": "s3://a/2.avro",
+                },
+                {
+                    "snapshot-id": 3,
+                    "parent-snapshot-id": 2,
+                    "timestamp-ms": 3000,
+                    "sequence-number": 3,
+                    "summary": {"operation": "append"},
+                    "manifest-list": "s3://a/3.avro",
+                },
+                {
+                    "snapshot-id": 4,
+                    "parent-snapshot-id": 3,
+                    "timestamp-ms": 4000,
+                    "sequence-number": 4,
+                    "summary": {"operation": "append"},
+                    "manifest-list": "s3://a/4.avro",
+                },
+            ],
+        }
+    )
+
+    result = latest_ancestor_before_timestamp(metadata, 3500)
+    assert result is not None
+    assert result.snapshot_id == 3
+
+    result = latest_ancestor_before_timestamp(metadata, 2500)
+    assert result is not None
+    assert result.snapshot_id == 2
+
+    result = latest_ancestor_before_timestamp(metadata, 5000)
+    assert result is not None
+    assert result.snapshot_id == 4
+
+    result = latest_ancestor_before_timestamp(metadata, 3000)
+    assert result is not None
+    assert result.snapshot_id == 2
+
+    result = latest_ancestor_before_timestamp(metadata, 1000)
+    assert result is None