Fixed bug: concurrent same name branch and tag writes

Author: vinjai

pyiceberg/table/__init__.py

@@ -349,7 +349,7 @@ def set_ref_snapshot(
             ),
         )
 
-        requirements = (AssertRefSnapshotId(snapshot_id=parent_snapshot_id, ref=ref_name),)
+        requirements = (AssertRefSnapshotId(snapshot_id=parent_snapshot_id, ref=ref_name, ref_type=type),)
         return self._apply(updates, requirements)
 
     def _set_ref_snapshot(
@@ -380,6 +380,7 @@ def _set_ref_snapshot(
             AssertRefSnapshotId(
                 snapshot_id=self.table_metadata.refs[ref_name].snapshot_id if ref_name in self.table_metadata.refs else None,
                 ref=ref_name,
+                ref_type=type,
             ),
         )
 

pyiceberg/table/update/__init__.py

@@ -592,6 +592,7 @@ class AssertRefSnapshotId(ValidatableTableRequirement):
 
     type: Literal["assert-ref-snapshot-id"] = Field(default="assert-ref-snapshot-id")
     ref: str = Field(...)
+    ref_type: SnapshotRefType = Field(...)
     snapshot_id: Optional[int] = Field(default=None, alias="snapshot-id")
 
     def validate(self, base_metadata: Optional[TableMetadata]) -> None:
@@ -607,6 +608,8 @@ def validate(self, base_metadata: Optional[TableMetadata]) -> None:
                 raise CommitFailedException(
                     f"Requirement failed: {ref_type} {self.ref} has changed: expected id {self.snapshot_id}, found {snapshot_ref.snapshot_id}"
                 )
+            elif ref_type != self.ref_type:
+                raise CommitFailedException(f"Requirement failed: {ref_type} {self.ref} can't be changed to type {self.ref_type}")
         elif self.snapshot_id is not None:
             raise CommitFailedException(f"Requirement failed: branch or tag {self.ref} is missing, expected {self.snapshot_id}")
 

pyiceberg/table/update/snapshot.py

@@ -287,6 +287,7 @@ def _commit(self) -> UpdatesAndRequirements:
                     if self._branch in self._transaction.table_metadata.refs
                     else self._transaction.table_metadata.current_snapshot_id,
                     ref=self._branch,
+                    ref_type=SnapshotRefType.BRANCH,
                 ),
             ),
         )

tests/table/test_init.py

@@ -49,7 +49,7 @@
     _match_deletes_to_data_file,
 )
 from pyiceberg.table.metadata import INITIAL_SEQUENCE_NUMBER, TableMetadataUtil, TableMetadataV2, _generate_snapshot_id
-from pyiceberg.table.refs import SnapshotRef
+from pyiceberg.table.refs import SnapshotRef, SnapshotRefType
 from pyiceberg.table.snapshots import (
     MetadataLogEntry,
     Operation,
@@ -982,28 +982,43 @@ def test_assert_table_uuid(table_v2: Table) -> None:
 
 def test_assert_ref_snapshot_id(table_v2: Table) -> None:
     base_metadata = table_v2.metadata
-    AssertRefSnapshotId(ref="main", snapshot_id=base_metadata.current_snapshot_id).validate(base_metadata)
+    AssertRefSnapshotId(ref="main", snapshot_id=base_metadata.current_snapshot_id, ref_type=SnapshotRefType.BRANCH).validate(
+        base_metadata
+    )
 
     with pytest.raises(CommitFailedException, match="Requirement failed: current table metadata is missing"):
-        AssertRefSnapshotId(ref="main", snapshot_id=1).validate(None)
+        AssertRefSnapshotId(ref="main", snapshot_id=1, ref_type=SnapshotRefType.BRANCH).validate(None)
 
     with pytest.raises(
         CommitFailedException,
         match="Requirement failed: branch main was created concurrently",
     ):
-        AssertRefSnapshotId(ref="main", snapshot_id=None).validate(base_metadata)
+        AssertRefSnapshotId(ref="main", snapshot_id=None, ref_type=SnapshotRefType.BRANCH).validate(base_metadata)
 
     with pytest.raises(
         CommitFailedException,
         match="Requirement failed: branch main has changed: expected id 1, found 3055729675574597004",
     ):
-        AssertRefSnapshotId(ref="main", snapshot_id=1).validate(base_metadata)
+        AssertRefSnapshotId(ref="main", snapshot_id=1, ref_type=SnapshotRefType.BRANCH).validate(base_metadata)
+
+    with pytest.raises(
+        CommitFailedException,
+        match="Requirement failed: branch or tag not_exist_branch is missing, expected 1",
+    ):
+        AssertRefSnapshotId(ref="not_exist_branch", snapshot_id=1, ref_type=SnapshotRefType.BRANCH).validate(base_metadata)
+
+    with pytest.raises(
+        CommitFailedException,
+        match="Requirement failed: branch or tag not_exist_tag is missing, expected 1",
+    ):
+        AssertRefSnapshotId(ref="not_exist_tag", snapshot_id=1, ref_type=SnapshotRefType.TAG).validate(base_metadata)
 
+    # existing Tag in metadata: test
     with pytest.raises(
         CommitFailedException,
-        match="Requirement failed: branch or tag not_exist is missing, expected 1",
+        match="Requirement failed: tag test can't be changed to type branch",
     ):
-        AssertRefSnapshotId(ref="not_exist", snapshot_id=1).validate(base_metadata)
+        AssertRefSnapshotId(ref="test", snapshot_id=3051729675574597004, ref_type=SnapshotRefType.BRANCH).validate(base_metadata)
 
 
 def test_assert_last_assigned_field_id(table_v2: Table) -> None: