diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 5ca3eb7a73..a41359e0a4 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/pyiceberg/table/update/snapshot.py @@ -1112,7 +1112,12 @@ def create_tag(self, snapshot_id: int, tag_name: str, max_ref_age_ms: int | None Returns: This for method chaining + + Raises: + ValueError: If a branch or tag with this name already exists. """ + if tag_name in self._effective_refs(): + raise ValueError(f"Ref {tag_name} already exists") update, requirement = self._transaction._set_ref_snapshot( snapshot_id=snapshot_id, ref_name=tag_name, @@ -1131,7 +1136,16 @@ def remove_tag(self, tag_name: str) -> ManageSnapshots: tag_name (str): name of tag to remove Returns: This for method chaining + + Raises: + NoSuchSnapshotRefError: If the tag does not exist. + SnapshotRefTypeError: If the name belongs to a branch. """ + ref = self._effective_refs().get(tag_name) + if ref is None: + raise NoSuchSnapshotRefError(f"Tag does not exist: {tag_name}") + if ref.snapshot_ref_type != SnapshotRefType.TAG: + raise SnapshotRefTypeError(f"Ref {tag_name} is a branch, not a tag") return self._remove_ref_snapshot(ref_name=tag_name) def create_branch( @@ -1153,7 +1167,12 @@ def create_branch( min_snapshots_to_keep (Optional[int]): min number of snapshots to keep for the branch Returns: This for method chaining + + Raises: + ValueError: If a branch or tag with this name already exists. """ + if branch_name in self._effective_refs(): + raise ValueError(f"Ref {branch_name} already exists") update, requirement = self._transaction._set_ref_snapshot( snapshot_id=snapshot_id, ref_name=branch_name, @@ -1174,7 +1193,19 @@ def remove_branch(self, branch_name: str) -> ManageSnapshots: branch_name (str): name of branch to remove Returns: This for method chaining + + Raises: + ValueError: If the branch is main. + NoSuchSnapshotRefError: If the branch does not exist. + SnapshotRefTypeError: If the name belongs to a tag. """ + if branch_name == MAIN_BRANCH: + raise ValueError("Cannot remove main branch") + ref = self._effective_refs().get(branch_name) + if ref is None: + raise NoSuchSnapshotRefError(f"Branch does not exist: {branch_name}") + if ref.snapshot_ref_type != SnapshotRefType.BRANCH: + raise SnapshotRefTypeError(f"Ref {branch_name} is a tag, not a branch") return self._remove_ref_snapshot(ref_name=branch_name) def set_current_snapshot(self, snapshot_id: int | None = None, ref_name: str | None = None) -> ManageSnapshots: diff --git a/tests/table/test_manage_snapshots.py b/tests/table/test_manage_snapshots.py index 37607132ac..4fd6c06ae3 100644 --- a/tests/table/test_manage_snapshots.py +++ b/tests/table/test_manage_snapshots.py @@ -29,7 +29,7 @@ from pyiceberg.table import CommitTableResponse, Table from pyiceberg.table.metadata import TableMetadataUtil from pyiceberg.table.refs import SnapshotRef, SnapshotRefType -from pyiceberg.table.update import SetSnapshotRefUpdate, TableRequirement, TableUpdate +from pyiceberg.table.update import RemoveSnapshotRefUpdate, SetSnapshotRefUpdate, TableRequirement, TableUpdate # The two snapshots in the table_v2 fixture: CHILD's parent is PARENT. PARENT_SNAPSHOT_ID = 3051729675574597004 @@ -82,6 +82,79 @@ def table_v2_main_behind(example_table_metadata_v2: dict[str, Any]) -> Table: ) +def test_create_ref_rejects_existing_name(table_v2: Table) -> None: + table_v2.catalog = MagicMock() + + # "test" is a tag in the fixture; neither kind may take an existing name. + with pytest.raises(ValueError, match="Ref test already exists"): + table_v2.manage_snapshots().create_branch(snapshot_id=CHILD_SNAPSHOT_ID, branch_name="test") + with pytest.raises(ValueError, match="Ref test already exists"): + table_v2.manage_snapshots().create_tag(snapshot_id=CHILD_SNAPSHOT_ID, tag_name="test") + + table_v2.catalog.commit_table.assert_not_called() + assert table_v2.metadata.refs["test"].snapshot_id == PARENT_SNAPSHOT_ID + + +def test_create_ref_rejects_name_staged_earlier_in_chain(table_v2: Table) -> None: + table_v2.catalog = MagicMock() + ms = table_v2.manage_snapshots().create_branch(snapshot_id=PARENT_SNAPSHOT_ID, branch_name="dev") + + with pytest.raises(ValueError, match="Ref dev already exists"): + ms.create_branch(snapshot_id=CHILD_SNAPSHOT_ID, branch_name="dev") + with pytest.raises(ValueError, match="Ref dev already exists"): + ms.create_tag(snapshot_id=CHILD_SNAPSHOT_ID, tag_name="dev") + + table_v2.catalog.commit_table.assert_not_called() + + +def test_remove_ref_rejects_main_missing_and_wrong_kind(table_v2_main_behind: Table) -> None: + table_v2_main_behind.catalog = MagicMock() + ms = table_v2_main_behind.manage_snapshots() + + with pytest.raises(ValueError, match="Cannot remove main branch"): + ms.remove_branch("main") + with pytest.raises(NoSuchSnapshotRefError, match="Branch does not exist: missing"): + ms.remove_branch("missing") + with pytest.raises(SnapshotRefTypeError, match="Ref test is a tag, not a branch"): + ms.remove_branch("test") + with pytest.raises(NoSuchSnapshotRefError, match="Tag does not exist: missing"): + ms.remove_tag("missing") + with pytest.raises(SnapshotRefTypeError, match="Ref audit is a branch, not a tag"): + ms.remove_tag("audit") + + table_v2_main_behind.catalog.commit_table.assert_not_called() + + +def test_remove_then_create_same_name_in_one_chain(table_v2_main_behind: Table) -> None: + table = table_v2_main_behind + mock_catalog = MagicMock() + mock_catalog.commit_table.return_value = _mock_commit_response(table) + table.catalog = mock_catalog + + table.manage_snapshots().remove_branch("audit").create_branch(snapshot_id=PARENT_SNAPSHOT_ID, branch_name="audit").commit() + + updates = _get_updates(mock_catalog) + assert [type(update) for update in updates] == [RemoveSnapshotRefUpdate, SetSnapshotRefUpdate] + assert isinstance(updates[1], SetSnapshotRefUpdate) + assert updates[1].snapshot_id == PARENT_SNAPSHOT_ID + + +def test_remove_then_create_same_name_as_other_kind_in_one_chain(table_v2_main_behind: Table) -> None: + table = table_v2_main_behind + mock_catalog = MagicMock() + mock_catalog.commit_table.return_value = _mock_commit_response(table) + table.catalog = mock_catalog + + # "test" is a tag; once removed in the chain the name is free for a branch. + table.manage_snapshots().remove_tag("test").create_branch(snapshot_id=CHILD_SNAPSHOT_ID, branch_name="test").commit() + + updates = _get_updates(mock_catalog) + assert [type(update) for update in updates] == [RemoveSnapshotRefUpdate, SetSnapshotRefUpdate] + assert isinstance(updates[1], SetSnapshotRefUpdate) + assert updates[1].type == SnapshotRefType.BRANCH + assert updates[1].snapshot_id == CHILD_SNAPSHOT_ID + + def test_set_current_snapshot_basic(table_v2: Table) -> None: snapshot_one = 3051729675574597004