Skip to content

fix(document): apply transactions atomically #10275

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kirangadhave wants to merge 2 commits into
base: main
Choose a base branch
from
+292 −13
Open

fix(document): apply transactions atomically #10275

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
911caa5
fix(document): apply transactions atomically
kirangadhave Jul 22, 2026
67544ac
fix(document): stage transactions with copy-on-write
kirangadhave Jul 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 51 additions & 11 deletions marimo/_messaging/notebook/document.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,19 +149,33 @@ def __iter__(self) -> Iterator[CellId_t]:
def apply(self, tx: Transaction) -> Transaction:
"""Validate and apply *tx*, return it with `version` assigned.

Raises `ValueError` for validation failures and `KeyError`
when a change references a non-existent cell.
Raises `ValueError` when a change has an invalid target, anchor,
or intra-transaction conflict.
"""
if not tx.changes:
return structs_replace(tx, version=self._version)

_validate(tx.changes, self._cells)

staged = NotebookDocument(
[
structs_replace(
cell,
config=structs_replace(cell.config),
)
for cell in self._cells
]
)
for change in tx.changes:
self._apply_change(change)
staged._apply_change(change)
Comment thread
kirangadhave marked this conversation as resolved.
Outdated

self._version += 1
return structs_replace(tx, version=self._version)
next_version = self._version + 1
applied = structs_replace(tx, version=next_version)

# Commit only after staging and result construction both succeed.
self._cells = staged._cells
self._version = next_version
return applied

def _apply_change(self, change: DocumentChange) -> None:
# TODO: refactor to use match/case (min Python is 3.10) once
Expand Down Expand Up @@ -322,22 +336,32 @@ def notebook_document_context(
def _validate(
changes: tuple[DocumentChange, ...], cells: list[NotebookCell]
) -> None:
"""Check for conflicting changes. Raises `ValueError`."""
existing_ids = {c.id for c in cells}
created: set[CellId_t] = set()
"""Check targets, anchors, and conflicts. Raises `ValueError`."""
live_ids = {cell.id for cell in cells}
# Claimed IDs are never released within a transaction, preserving the
# rule that delete + recreate cannot reuse an ID in the same batch.
claimed_ids = set(live_ids)
deleted: set[CellId_t] = set()
updated: set[CellId_t] = set()
moved: set[CellId_t] = set()

for change in changes:
if isinstance(change, CreateCell):
if change.cell_id in existing_ids or change.cell_id in created:
if change.cell_id in claimed_ids:
raise ValueError(f"Cell {change.cell_id!r} already exists")
if change.before is not None and change.after is not None:
raise ValueError(
"CreateCell cannot specify both 'before' and 'after'"
)
created.add(change.cell_id)
anchor = (
change.before if change.before is not None else change.after
)
if anchor is not None and anchor not in live_ids:
raise ValueError(
f"CreateCell anchor {anchor!r} does not exist"
)
live_ids.add(change.cell_id)
claimed_ids.add(change.cell_id)

elif isinstance(change, DeleteCell):
if change.cell_id in deleted:
Expand All @@ -354,20 +378,34 @@ def _validate(
f"Cannot delete cell {change.cell_id!r} that is also "
f"moved in the same transaction"
)
if change.cell_id not in live_ids:
raise ValueError(f"Cell {change.cell_id!r} does not exist")
deleted.add(change.cell_id)
live_ids.remove(change.cell_id)

elif isinstance(change, MoveCell):
if change.cell_id in deleted:
raise ValueError(
f"Cannot move cell {change.cell_id!r} that is also "
f"deleted in the same transaction"
)
if change.cell_id not in live_ids:
raise ValueError(f"Cell {change.cell_id!r} does not exist")
if change.before is not None and change.after is not None:
raise ValueError(
"MoveCell cannot specify both 'before' and 'after'"
)
if change.before is None and change.after is None:
anchor = (
change.before if change.before is not None else change.after
)
if anchor is None:
raise ValueError("MoveCell requires 'before' or 'after'")
if anchor == change.cell_id:
raise ValueError(
f"MoveCell cannot anchor cell {change.cell_id!r} to itself"
)
if anchor not in live_ids:
raise ValueError(f"MoveCell anchor {anchor!r} does not exist")
moved.add(change.cell_id)

elif isinstance(change, ReorderCells):
Expand All @@ -379,6 +417,8 @@ def _validate(
f"Cannot update cell {change.cell_id!r} that is also "
f"deleted in the same transaction"
)
if change.cell_id not in live_ids:
raise ValueError(f"Cell {change.cell_id!r} does not exist")
updated.add(change.cell_id)

else:
Expand Down
191 changes: 189 additions & 2 deletions tests/_messaging/notebook/test_document.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,23 @@ def _ids(doc: NotebookDocument) -> list[str]:
return [str(cid) for cid in doc.cell_ids]


def _state(
doc: NotebookDocument,
) -> list[tuple[CellId_t, str, str, int | None, bool, bool, int]]:
return [
(
cell.id,
cell.code,
cell.name,
cell.config.column,
cell.config.disabled,
cell.config.hide_code,
cell.version,
)
for cell in doc.cells
]


# ------------------------------------------------------------------
# CreateCell
# ------------------------------------------------------------------
Expand Down Expand Up @@ -148,6 +165,42 @@ def test_after_pending(self) -> None:
)
assert _ids(doc) == snapshot(["a", "x", "y"])

def test_missing_after_anchor_raises(self) -> None:
doc = _doc("a")
with pytest.raises(
ValueError,
match="CreateCell anchor 'missing' does not exist",
):
doc.apply(
_tx(
CreateCell(
cell_id=CellId_t("new"),
code="x",
name="__",
config=CellConfig(),
after=CellId_t("missing"),
)
)
)

def test_missing_before_anchor_raises(self) -> None:
doc = _doc("a")
with pytest.raises(
ValueError,
match="CreateCell anchor 'missing' does not exist",
):
doc.apply(
_tx(
CreateCell(
cell_id=CellId_t("new"),
code="x",
name="__",
config=CellConfig(),
before=CellId_t("missing"),
)
)
)

def test_duplicate_id_raises(self) -> None:
doc = _doc("a")
with pytest.raises(ValueError, match="already exists"):
Expand Down Expand Up @@ -251,7 +304,7 @@ def test_delete_all(self) -> None:

def test_delete_not_found(self) -> None:
doc = _doc("a")
with pytest.raises(KeyError):
with pytest.raises(ValueError, match="Cell 'missing' does not exist"):
doc.apply(_tx(DeleteCell(cell_id=CellId_t("missing"))))


Expand All @@ -276,6 +329,66 @@ def test_no_anchor_raises(self) -> None:
with pytest.raises(ValueError, match="before.*after"):
doc.apply(_tx(MoveCell(cell_id=CellId_t("a"))))

def test_missing_cell_raises(self) -> None:
doc = _doc("a")
with pytest.raises(ValueError, match="Cell 'missing' does not exist"):
doc.apply(
_tx(
MoveCell(
cell_id=CellId_t("missing"),
after=CellId_t("a"),
)
)
)

def test_missing_anchor_raises(self) -> None:
doc = _doc("a", "b")
with pytest.raises(
ValueError,
match="MoveCell anchor 'missing' does not exist",
):
doc.apply(
_tx(
MoveCell(
cell_id=CellId_t("a"),
after=CellId_t("missing"),
)
)
)

def test_self_anchor_raises(self) -> None:
doc = _doc("a", "b")
with pytest.raises(
ValueError,
match="MoveCell cannot anchor cell 'a' to itself",
):
doc.apply(
_tx(
MoveCell(
cell_id=CellId_t("a"),
after=CellId_t("a"),
)
)
)

def test_anchor_created_earlier_in_batch(self) -> None:
doc = _doc("a", "b")
doc.apply(
_tx(
CreateCell(
cell_id=CellId_t("new"),
code="x",
name="__",
config=CellConfig(),
),
MoveCell(
cell_id=CellId_t("a"),
after=CellId_t("new"),
),
)
)
assert _ids(doc) == snapshot(["b", "new", "a"])


# ------------------------------------------------------------------
# SetCode
Expand All @@ -291,7 +404,7 @@ def test_update_code(self) -> None:

def test_not_found(self) -> None:
doc = _doc("a")
with pytest.raises(KeyError):
with pytest.raises(ValueError, match="Cell 'missing' does not exist"):
doc.apply(_tx(SetCode(cell_id=CellId_t("missing"), code="x")))


Expand Down Expand Up @@ -522,6 +635,80 @@ def test_valid_mixed_ops(self) -> None:
assert _ids(doc) == snapshot(["a", "new"])


# ------------------------------------------------------------------
# Atomicity
# ------------------------------------------------------------------


class TestAtomicity:
def test_invalid_late_anchor_does_not_mutate_document(self) -> None:
doc = _doc("a", "b")
doc.apply(_tx(SetCode(cell_id=CellId_t("a"), code="original")))
before_state = _state(doc)
before_version = doc.version

with pytest.raises(
ValueError,
match="MoveCell anchor 'missing' does not exist",
):
doc.apply(
_tx(
SetName(cell_id=CellId_t("a"), name="renamed"),
CreateCell(
cell_id=CellId_t("new"),
code="new code",
name="__",
config=CellConfig(disabled=True),
),
MoveCell(
cell_id=CellId_t("b"),
after=CellId_t("missing"),
),
)
)

assert _state(doc) == before_state
assert doc.version == before_version

def test_unexpected_apply_failure_does_not_mutate_document(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
doc = _doc("a", "b")
before_state = _state(doc)
before_version = doc.version
original_apply_change = NotebookDocument._apply_change
applied_changes = 0

def fail_after_second_change(
staged: NotebookDocument, change: DocumentChange
) -> None:
nonlocal applied_changes
original_apply_change(staged, change)
applied_changes += 1
if applied_changes == 2:
raise RuntimeError("injected apply failure")

monkeypatch.setattr(
NotebookDocument, "_apply_change", fail_after_second_change
)

with pytest.raises(RuntimeError, match="injected apply failure"):
doc.apply(
_tx(
SetName(cell_id=CellId_t("a"), name="renamed"),
CreateCell(
cell_id=CellId_t("new"),
code="new code",
name="__",
config=CellConfig(disabled=True),
),
)
)

assert _state(doc) == before_state
assert doc.version == before_version


# ------------------------------------------------------------------
# Versioning
# ------------------------------------------------------------------
Expand Down
Loading