From 502ac56784d14832e79a279f56f7b2ae707d048b Mon Sep 17 00:00:00 2001 From: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com> Date: Sun, 1 Nov 2026 10:19:05 -0800 Subject: [PATCH] Table: assert every written ref in a multi-branch transaction Motivation: Transaction._stage deduplicated requirements by type only, so a transaction writing two branches sent AssertRefSnapshotId for the first branch only. A concurrent change to the second branch was not detected and was overwritten. Approach: deduplicate AssertRefSnapshotId by (type, ref); other requirement types keep per-type deduplication. Duplicates within a single _stage call are also dropped. Validation: added test_stage_keeps_ref_requirement_per_ref, which fails without the change and passes with it (uv run pytest tests/table/test_init.py -k stage_keeps). The reproduction script from the issue prints [1, 2, 3] before and [1, 3] after. make lint passes. Four other tests in tests/table/test_init.py fail identically on the base commit in this environment (missing optional modules). The retry path with two branches is not covered by a dedicated test. Report: https://github.com/apache/iceberg-python/issues/4183 Signed-off-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com> Assisted-by: claude-sonnet-5-5 (via Claude Code) --- pyiceberg/table/__init__.py | 9 +++++++-- tests/table/test_init.py | 10 ++++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) 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)