diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 8d79472af5..bfb252c6e0 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -305,9 +305,14 @@ def _stage( # For the requirements, it does not make sense to add a requirement more than once # For example, you cannot assert that the current schema has two different IDs - existing_requirements = {type(requirement) for requirement in self._requirements} + # Ref requirements are the exception: each ref needs its own assertion + def key(requirement: TableRequirement) -> tuple[type, str | None]: + return type(requirement), requirement.ref if isinstance(requirement, AssertRefSnapshotId) else None + + existing_requirements = {key(requirement) for requirement in self._requirements} for new_requirement in requirements: - if type(new_requirement) not in existing_requirements: + if key(new_requirement) not in existing_requirements: + existing_requirements.add(key(new_requirement)) self._requirements = self._requirements + (new_requirement,) return self diff --git a/tests/table/test_init.py b/tests/table/test_init.py index f6b984f1a3..0bb5a139a5 100644 --- a/tests/table/test_init.py +++ b/tests/table/test_init.py @@ -1437,6 +1437,16 @@ def test_assert_ref_snapshot_id(table_v2: Table) -> None: assert AssertRefSnapshotId(ref="main").model_dump_json() == expected_json +def test_stage_keeps_ref_requirement_per_ref(table_v2: Table) -> None: + tx = table_v2.transaction() + main = AssertRefSnapshotId(ref=MAIN_BRANCH, snapshot_id=3055729675574597004) + test = AssertRefSnapshotId(ref="test", snapshot_id=3051729675574597004) + + tx._stage((), (main, test, main, test)) + + assert tx._requirements == (main, test) + + def test_assert_last_assigned_field_id(table_v2: Table) -> None: base_metadata = table_v2.metadata AssertLastAssignedFieldId(last_assigned_field_id=base_metadata.last_column_id).validate(base_metadata)