From 29feaf71cbca7e227bef04b95caef0731680901f Mon Sep 17 00:00:00 2001
From: Borodin Gregory <grihabor@gmail.com>
Date: Thu, 23 Jan 2025 22:40:40 +0100
Subject: [PATCH] Make sure refs are removed when removing snapshots

---
 pyiceberg/table/update/__init__.py | 12 +++++++++++-
 tests/table/test_init.py           |  9 ++++++---
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/pyiceberg/table/update/__init__.py b/pyiceberg/table/update/__init__.py
index 74dc4644bf..2577ed03d4 100644
--- a/pyiceberg/table/update/__init__.py
+++ b/pyiceberg/table/update/__init__.py
@@ -468,8 +468,18 @@ def _(update: RemoveSnapshotsUpdate, base_metadata: TableMetadata, context: _Tab
         for s in base_metadata.snapshots
         if s.snapshot_id not in update.snapshot_ids
     ]
+
+    remove_ref_updates = (
+        RemoveSnapshotRefUpdate(ref_name=ref_name)
+        for ref_name, ref in base_metadata.refs.items()
+        if ref.snapshot_id in update.snapshot_ids
+    )
+    new_metadata = base_metadata
+    for remove_ref in remove_ref_updates:
+        new_metadata = _apply_table_update(remove_ref, new_metadata, context)
+
     context.add_update(update)
-    return base_metadata.model_copy(update={"snapshots": snapshots})
+    return new_metadata.model_copy(update={"snapshots": snapshots})
 
 
 @_apply_table_update.register(RemoveSnapshotRefUpdate)
diff --git a/tests/table/test_init.py b/tests/table/test_init.py
index aa27a6f3bd..e76f3de9a6 100644
--- a/tests/table/test_init.py
+++ b/tests/table/test_init.py
@@ -795,15 +795,18 @@ def test_update_metadata_set_snapshot_ref(table_v2: Table) -> None:
 
 
 def test_update_remove_snapshots(table_v2: Table) -> None:
-    update = RemoveSnapshotsUpdate(
-        snapshot_ids=[3051729675574597004],
-    )
+    # assert fixture data to easily understand the test assumptions
+    assert len(table_v2.metadata.snapshots) == 2
+    assert len(table_v2.metadata.refs) == 2
+    update = RemoveSnapshotsUpdate(snapshot_ids=[3051729675574597004])
     new_metadata = update_table_metadata(table_v2.metadata, (update,))
     assert len(new_metadata.snapshots) == 1
     assert new_metadata.snapshots[0].snapshot_id == 3055729675574597004
     assert new_metadata.snapshots[0].parent_snapshot_id == None
     assert new_metadata.current_snapshot_id == 3055729675574597004
     assert new_metadata.last_updated_ms > table_v2.metadata.last_updated_ms
+    assert len(new_metadata.refs) == 1
+    assert new_metadata.refs["main"].snapshot_id == 3055729675574597004
 
 
 def test_update_remove_snapshots_doesnt_exist(table_v2: Table) -> None: