update tests

Author: chinmay-bhat

dev/provision.py

@@ -389,3 +389,51 @@
         VALUES (4)
         """
     )
+
+    spark.sql(
+        f"""
+        CREATE OR REPLACE TABLE {catalog_name}.default.test_table_rollback_to_snapshot_id (
+            timestamp int,
+            number integer
+        )
+        USING iceberg
+        TBLPROPERTIES (
+            'format-version'='2'
+        );
+        """
+    )
+
+    spark.sql(
+        f"""
+        INSERT INTO {catalog_name}.default.test_table_rollback_to_snapshot_id
+        VALUES (200, 1)
+        """
+    )
+
+    spark.sql(
+        f"""
+        INSERT INTO {catalog_name}.default.test_table_rollback_to_snapshot_id
+        VALUES (202, 2)
+        """
+    )
+
+    spark.sql(
+        f"""
+        DELETE FROM {catalog_name}.default.test_table_rollback_to_snapshot_id
+        WHERE number = 2
+        """
+    )
+
+    spark.sql(
+        f"""
+        INSERT INTO {catalog_name}.default.test_table_rollback_to_snapshot_id
+        VALUES (204, 3)
+        """
+    )
+
+    spark.sql(
+        f"""
+        INSERT INTO {catalog_name}.default.test_table_rollback_to_snapshot_id
+        VALUES (206, 4)
+        """
+    )

tests/integration/test_snapshot_operations.py

@@ -49,37 +49,38 @@ def test_manage_snapshots_context_manager(catalog: Catalog) -> None:
     tbl = catalog.load_table(identifier)
     assert len(tbl.history()) > 3
     current_snapshot_id = tbl.current_snapshot().snapshot_id  # type: ignore
-    rollback_snapshot_id = tbl.history()[-4].snapshot_id
+    expected_snapshot_id = tbl.history()[-4].snapshot_id
     with tbl.manage_snapshots() as ms:
         ms.create_tag(snapshot_id=current_snapshot_id, tag_name="testing")
-        ms.rollback_to_snapshot(snapshot_id=rollback_snapshot_id)
+        ms.set_current_snapshot(snapshot_id=expected_snapshot_id)
+        ms.create_tag(snapshot_id=expected_snapshot_id, tag_name="testing2")
     assert tbl.current_snapshot().snapshot_id is not current_snapshot_id  # type: ignore
     assert tbl.metadata.refs["testing"].snapshot_id == current_snapshot_id
-    assert tbl.metadata.refs["main"] == SnapshotRef(snapshot_id=rollback_snapshot_id, snapshot_ref_type="branch")
+    assert tbl.metadata.refs["main"] == SnapshotRef(snapshot_id=expected_snapshot_id, snapshot_ref_type="branch")
+    assert tbl.metadata.refs["testing2"].snapshot_id == expected_snapshot_id
 
 
 @pytest.mark.integration
 @pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
 def test_rollback_to_snapshot(catalog: Catalog) -> None:
-    identifier = "default.test_table_snapshot_operations"
+    identifier = "default.test_table_rollback_to_snapshot_id"
     tbl = catalog.load_table(identifier)
     assert len(tbl.history()) > 3
-    rollback_snapshot_id = tbl.history()[-3].snapshot_id
+    rollback_snapshot_id = tbl.current_snapshot().parent_snapshot_id  # type: ignore
     current_snapshot_id = tbl.current_snapshot().snapshot_id  # type: ignore
-    tbl.manage_snapshots().rollback_to_snapshot(snapshot_id=rollback_snapshot_id).commit()
+    tbl.manage_snapshots().rollback_to_snapshot(snapshot_id=rollback_snapshot_id).commit()  # type: ignore
     assert tbl.current_snapshot().snapshot_id is not current_snapshot_id  # type: ignore
     assert tbl.metadata.refs["main"] == SnapshotRef(snapshot_id=rollback_snapshot_id, snapshot_ref_type="branch")
 
 
 @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"
+    identifier = "default.test_table_rollback_to_snapshot_id"
     tbl = catalog.load_table(identifier)
-    assert len(tbl.history()) > 3
-    current_snapshot_id = tbl.current_snapshot().snapshot_id  # type: ignore
-    timestamp = tbl.history()[-2].timestamp_ms
-    expected_snapshot_id = tbl.history()[-3].snapshot_id
+    assert len(tbl.history()) > 4
+    current_snapshot_id, timestamp = tbl.history()[-1].snapshot_id, tbl.history()[-1].timestamp_ms
+    expected_snapshot_id = tbl.snapshot_by_id(current_snapshot_id).parent_snapshot_id  # type: ignore
     # not inclusive of rollback_timestamp
     tbl.manage_snapshots().rollback_to_timestamp(timestamp=timestamp).commit()
     assert tbl.current_snapshot().snapshot_id is not current_snapshot_id  # type: ignore
@@ -107,7 +108,7 @@ def test_set_current_snapshot_with_ref_name(catalog: Catalog) -> None:
     assert len(tbl.history()) > 3
     current_snapshot_id = tbl.current_snapshot().snapshot_id  # type: ignore
     expected_snapshot_id = tbl.history()[-3].snapshot_id
-    tbl.manage_snapshots().create_tag(snapshot_id=expected_snapshot_id, tag_name="test-tag19").commit()
-    tbl.manage_snapshots().set_current_snapshot(ref_name="test-tag19").commit()
+    tbl.manage_snapshots().create_tag(snapshot_id=expected_snapshot_id, tag_name="test-tag").commit()
+    tbl.manage_snapshots().set_current_snapshot(ref_name="test-tag").commit()
     assert tbl.current_snapshot().snapshot_id is not current_snapshot_id  # type: ignore
     assert tbl.metadata.refs["main"] == SnapshotRef(snapshot_id=expected_snapshot_id, snapshot_ref_type="branch")