cherry-pick: Support Remove Branch or Tag APIs (#822)

Author: grihabor

pyiceberg/table/update/__init__.py

@@ -472,6 +472,23 @@ def _(update: RemoveSnapshotsUpdate, base_metadata: TableMetadata, context: _Tab
     return base_metadata.model_copy(update={"snapshots": snapshots})
 
 
+@_apply_table_update.register(RemoveSnapshotRefUpdate)
+def _(update: RemoveSnapshotRefUpdate, base_metadata: TableMetadata, context: _TableMetadataUpdateContext) -> TableMetadata:
+    if (existing_ref := base_metadata.refs.get(update.ref_name, None)) is None:
+        return base_metadata
+
+    if base_metadata.snapshot_by_id(existing_ref.snapshot_id) is None:
+        raise ValueError(f"Cannot remove {update.ref_name} ref with unknown snapshot {existing_ref.snapshot_id}")
+
+    if update.ref_name == MAIN_BRANCH:
+        raise ValueError("Cannot remove main branch")
+
+    metadata_refs = {**base_metadata.refs}
+    metadata_refs.pop(update.ref_name, None)
+    context.add_update(update)
+    return base_metadata.model_copy(update={"refs": metadata_refs})
+
+
 @_apply_table_update.register(AddSortOrderUpdate)
 def _(update: AddSortOrderUpdate, base_metadata: TableMetadata, context: _TableMetadataUpdateContext) -> TableMetadata:
     context.add_update(update)

pyiceberg/table/update/snapshot.py

@@ -65,6 +65,7 @@
 from pyiceberg.table.update import (
     AddSnapshotUpdate,
     AssertRefSnapshotId,
+    RemoveSnapshotRefUpdate,
     SetSnapshotRefUpdate,
     TableRequirement,
     TableUpdate,
@@ -749,6 +750,27 @@ def _commit(self) -> UpdatesAndRequirements:
         """Apply the pending changes and commit."""
         return self._updates, self._requirements
 
+    def _remove_ref_snapshot(self, ref_name: str) -> ManageSnapshots:
+        """Remove a snapshot ref.
+        Args:
+            ref_name: branch / tag name to remove
+        Stages the updates and requirements for the remove-snapshot-ref.
+        Returns
+            This method for chaining
+        """
+        updates = (RemoveSnapshotRefUpdate(ref_name=ref_name),)
+        requirements = (
+            AssertRefSnapshotId(
+                snapshot_id=self._transaction.table_metadata.refs[ref_name].snapshot_id
+                if ref_name in self._transaction.table_metadata.refs
+                else None,
+                ref=ref_name,
+            ),
+        )
+        self._updates += updates
+        self._requirements += requirements
+        return self
+
     def create_tag(self, snapshot_id: int, tag_name: str, max_ref_age_ms: Optional[int] = None) -> ManageSnapshots:
         """
         Create a new tag pointing to the given snapshot id.
@@ -771,6 +793,16 @@ def create_tag(self, snapshot_id: int, tag_name: str, max_ref_age_ms: Optional[i
         self._requirements += requirement
         return self
 
+    def remove_tag(self, tag_name: str) -> ManageSnapshots:
+        """
+        Remove a tag.
+        Args:
+            tag_name (str): name of tag to remove
+        Returns:
+            This for method chaining
+        """
+        return self._remove_ref_snapshot(ref_name=tag_name)
+
     def create_branch(
         self,
         snapshot_id: int,
@@ -802,3 +834,13 @@ def create_branch(
         self._updates += update
         self._requirements += requirement
         return self
+
+    def remove_branch(self, branch_name: str) -> ManageSnapshots:
+        """
+        Remove a branch.
+        Args:
+            branch_name (str): name of branch to remove
+        Returns:
+            This for method chaining
+        """
+        return self._remove_ref_snapshot(ref_name=branch_name)

tests/integration/test_snapshot_operations.py

@@ -40,3 +40,35 @@ def test_create_branch(catalog: Catalog) -> None:
     branch_snapshot_id = tbl.history()[-2].snapshot_id
     tbl.manage_snapshots().create_branch(snapshot_id=branch_snapshot_id, branch_name="branch123").commit()
     assert tbl.metadata.refs["branch123"] == SnapshotRef(snapshot_id=branch_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_remove_tag(catalog: Catalog) -> None:
+    identifier = "default.test_table_snapshot_operations"
+    tbl = catalog.load_table(identifier)
+    assert len(tbl.history()) > 3
+    # first, create the tag to remove
+    tag_name = "tag_to_remove"
+    tag_snapshot_id = tbl.history()[-3].snapshot_id
+    tbl.manage_snapshots().create_tag(snapshot_id=tag_snapshot_id, tag_name=tag_name).commit()
+    assert tbl.metadata.refs[tag_name] == SnapshotRef(snapshot_id=tag_snapshot_id, snapshot_ref_type="tag")
+    # now, remove the tag
+    tbl.manage_snapshots().remove_tag(tag_name=tag_name).commit()
+    assert tbl.metadata.refs.get(tag_name, None) is None
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
+def test_remove_branch(catalog: Catalog) -> None:
+    identifier = "default.test_table_snapshot_operations"
+    tbl = catalog.load_table(identifier)
+    assert len(tbl.history()) > 2
+    # first, create the branch to remove
+    branch_name = "branch_to_remove"
+    branch_snapshot_id = tbl.history()[-2].snapshot_id
+    tbl.manage_snapshots().create_branch(snapshot_id=branch_snapshot_id, branch_name=branch_name).commit()
+    assert tbl.metadata.refs[branch_name] == SnapshotRef(snapshot_id=branch_snapshot_id, snapshot_ref_type="branch")
+    # now, remove the branch
+    tbl.manage_snapshots().remove_branch(branch_name=branch_name).commit()
+    assert tbl.metadata.refs.get(branch_name, None) is None