diff --git a/lacing/store/base.py b/lacing/store/base.py index d2e53ae1..963ea8f4 100644 --- a/lacing/store/base.py +++ b/lacing/store/base.py @@ -55,6 +55,17 @@ def remove(self, annotation_id) -> Annotation | None: """Remove and return the annotation with this id, or None if absent.""" ... + def update(self, annotation: Annotation) -> None: + """Replace the stored annotation that has ``annotation.id``, atomically. + + One commit: a concurrent reader sees the old annotation or the new one, + never neither (``remove(id)`` then ``add(...)`` leaves a gap in + between). The replacement may move to a different interval or tier. + Raises ``KeyError`` when no annotation has that id — use :meth:`add` + for a new one. + """ + ... + def all(self) -> Iterator[Annotation]: """Iterate every annotation in the store, order unspecified.""" ... diff --git a/lacing/store/mapping.py b/lacing/store/mapping.py index 9b86753a..bf117fd6 100644 --- a/lacing/store/mapping.py +++ b/lacing/store/mapping.py @@ -199,18 +199,27 @@ def __delitem__(self, key: TimeInterval) -> None: # --- annotation-level ----------------------------------------------------- def add(self, annotation: Annotation) -> None: - if annotation.id in self._ids: - raise ValueError(f"annotation id {annotation.id} already in the store") - self._put(annotation) - self._ids.add(annotation.id) - super().add(annotation) + with self._lock: + if annotation.id in self._ids: + raise ValueError(f"annotation id {annotation.id} already in the store") + self._put(annotation) + self._ids.add(annotation.id) + super().add(annotation) def remove(self, annotation_id: UUID) -> Annotation | None: - if annotation_id not in self._ids: - return None - del self._mapping[str(annotation_id)] - self._ids.discard(annotation_id) - return super().remove(annotation_id) + with self._lock: + if annotation_id not in self._ids: + return None + del self._mapping[str(annotation_id)] + self._ids.discard(annotation_id) + return super().remove(annotation_id) + + def update(self, annotation: Annotation) -> None: + with self._lock: + if annotation.id not in self._ids: + raise KeyError(annotation.id) + self._put(annotation) # one key overwritten: never absent in the mapping + super().update(annotation) # --- tier registry -------------------------------------------------------- diff --git a/lacing/store/memory.py b/lacing/store/memory.py index cf4b7838..ec875481 100644 --- a/lacing/store/memory.py +++ b/lacing/store/memory.py @@ -8,6 +8,7 @@ from __future__ import annotations +import threading from collections.abc import Iterable, Iterator from fractions import Fraction from uuid import UUID @@ -25,6 +26,11 @@ def _key(interval: TimeInterval) -> tuple[Fraction, Fraction]: return interval.start.to_fraction(), interval.end.to_fraction() +def _index_of(bucket: list[Annotation], ann: Annotation) -> int: + """Position of ``ann`` in ``bucket`` by identity (equal annotations may repeat).""" + return next(i for i, a in enumerate(bucket) if a is ann) + + class MemoryStore: """``IntervalAnnotationStore`` implementation over ``intervaltree``. @@ -34,6 +40,9 @@ class MemoryStore: """ def __init__(self) -> None: + # Serialises writers and the snapshots readers take, so ``update`` is + # atomic for a reader on another thread (lacing#66). + self._lock = threading.RLock() self._tree: IntervalTree = IntervalTree() # interval-keyed bucket: (start_frac, end_frac) -> list of annotations self._buckets: dict[tuple[Fraction, Fraction], list[Annotation]] = {} @@ -42,6 +51,16 @@ def __init__(self) -> None: self._timeless: dict[UUID, Annotation] = {} self._tiers: dict[str, Tier] = {} + # --- copy / pickle --------------------------------------------------------- + + def __getstate__(self) -> dict: + # A lock cannot be copied or pickled; each copy gets its own. + return {k: v for k, v in self.__dict__.items() if k != "_lock"} + + def __setstate__(self, state: dict) -> None: + self.__dict__.update(state) + self._lock = threading.RLock() + # --- helpers -------------------------------------------------------------- def _ensure_node(self, key: tuple[Fraction, Fraction]) -> list[Annotation]: @@ -64,25 +83,28 @@ def _drop_node(self, key: tuple[Fraction, Fraction]) -> None: # --- MutableMapping interface ------------------------------------------- def __getitem__(self, key: TimeInterval) -> list[Annotation]: - bucket = self._buckets.get(_key(key)) - if bucket is None: - raise KeyError(key) - return list(bucket) # defensive copy + with self._lock: + bucket = self._buckets.get(_key(key)) + if bucket is None: + raise KeyError(key) + return list(bucket) # defensive copy def __setitem__(self, key: TimeInterval, value: list[Annotation]) -> None: k = _key(key) - # Replace; ensure tree node exists exactly when bucket is non-empty. - if not value: - self._drop_node(k) - return - self._ensure_node(k) - self._buckets[k] = list(value) + with self._lock: + # Replace; ensure tree node exists exactly when bucket is non-empty. + if not value: + self._drop_node(k) + return + self._ensure_node(k) + self._buckets[k] = list(value) def __delitem__(self, key: TimeInterval) -> None: k = _key(key) - if k not in self._buckets: - raise KeyError(key) - self._drop_node(k) + with self._lock: + if k not in self._buckets: + raise KeyError(key) + self._drop_node(k) def __iter__(self) -> Iterator[TimeInterval]: for ann_list in self._buckets.values(): @@ -97,40 +119,83 @@ def __len__(self) -> int: # --- annotation-level convenience --------------------------------------- def add(self, annotation: Annotation) -> None: - iv = annotation.interval - if iv is None: - self._timeless[annotation.id] = annotation - return - bucket = self._ensure_node(_key(iv)) - bucket.append(annotation) + self._insert(annotation) + + def _insert(self, annotation: Annotation) -> None: + """``add`` without subclass checks (``update`` must not trip a duplicate-id guard).""" + with self._lock: + iv = annotation.interval + if iv is None: + self._timeless[annotation.id] = annotation + return + bucket = self._ensure_node(_key(iv)) + bucket.append(annotation) def remove(self, annotation_id: UUID) -> Annotation | None: - # Search timeless first. - if annotation_id in self._timeless: - return self._timeless.pop(annotation_id) - for k, bucket in list(self._buckets.items()): - for i, a in enumerate(bucket): - if a.id == annotation_id: - bucket.pop(i) - if not bucket: - self._drop_node(k) - return a - return None + with self._lock: + # Search timeless first. + if annotation_id in self._timeless: + return self._timeless.pop(annotation_id) + for k, bucket in list(self._buckets.items()): + for i, a in enumerate(bucket): + if a.id == annotation_id: + bucket.pop(i) + if not bucket: + self._drop_node(k) + return a + return None + + def update(self, annotation: Annotation) -> None: + with self._lock: + old = next((a for a in self.all() if a.id == annotation.id), None) + if old is None: + raise KeyError(annotation.id) + old_iv, new_iv = old.interval, annotation.interval + if old_iv is None and new_iv is None: + self._timeless[annotation.id] = annotation # one slot write + elif ( + old_iv is not None + and new_iv is not None + and _key(old_iv) == _key(new_iv) + ): + bucket = self._buckets[_key(new_iv)] + bucket[_index_of(bucket, old)] = annotation # one slot write + else: + # The copy lands somewhere else: put the new one in *before* + # taking the old out, so the id is never absent even to a + # reader that does not take the lock. + self._insert(annotation) + self._take_out(old) + + def _take_out(self, old: Annotation) -> None: + """Remove exactly ``old`` (by identity) from where ``add`` put it.""" + if old.interval is None: + del self._timeless[old.id] + return + k = _key(old.interval) + bucket = self._buckets[k] + bucket.pop(_index_of(bucket, old)) + if not bucket: + self._drop_node(k) def all(self) -> Iterator[Annotation]: - for bucket in self._buckets.values(): - yield from bucket - yield from self._timeless.values() + with self._lock: + snapshot = [a for bucket in self._buckets.values() for a in bucket] + snapshot.extend(self._timeless.values()) + yield from snapshot # --- Allen-relation queries --------------------------------------------- + def _bucket_snapshot(self) -> list[list[Annotation]]: + """Copies of the non-empty buckets, taken under the lock (see ``update``).""" + with self._lock: + return [list(b) for b in self._buckets.values() if b] + def _query( self, query: TimeInterval, relation: AllenRelation ) -> Iterator[Annotation]: pred = PREDICATE_BY_RELATION[relation] - for bucket in self._buckets.values(): - if not bucket: - continue + for bucket in self._bucket_snapshot(): iv = bucket[0].interval if iv is not None and pred(iv, query): yield from bucket @@ -138,9 +203,7 @@ def _query( def intersects(self, query: TimeInterval) -> Iterator[Annotation]: from lacing.allen import intersects as _intersects - for bucket in self._buckets.values(): - if not bucket: - continue + for bucket in self._bucket_snapshot(): iv = bucket[0].interval if iv is not None and _intersects(iv, query): yield from bucket @@ -170,9 +233,7 @@ def relate( self, query: TimeInterval, relations: Iterable[AllenRelation] ) -> Iterator[Annotation]: rels = set(relations) - for bucket in self._buckets.values(): - if not bucket: - continue + for bucket in self._bucket_snapshot(): iv = bucket[0].interval if iv is None: continue diff --git a/lacing/store/postgres.py b/lacing/store/postgres.py index cae78275..4a7ad70c 100644 --- a/lacing/store/postgres.py +++ b/lacing/store/postgres.py @@ -643,6 +643,34 @@ def remove(self, annotation_id: UUID) -> Annotation | None: ) return ann + def update(self, annotation: Annotation) -> None: + """Replace the row with this id in ONE ``UPDATE`` statement (lacing#66). + + A single statement is atomic in any isolation mode, so a concurrent + reader sees the old row or the new one — unlike ``remove`` then ``add``, + which commit twice. Raises ``KeyError`` if the id is not in this + ``(owner, project)`` and ``TierOverlapError`` like :meth:`add`. + """ + cols, values = self._annotation_to_row(annotation) + # The key and the tenant are what the row is found by; they never change. + fixed = {"id", "owner_id", "project_id"} + changing = [(c, v) for c, v in zip(cols, values) if c not in fixed] + scope, params = self._scope() + assignments = ", ".join(f"{c} = %s" for c, _ in changing) + try: + with self._conn.cursor() as cur: + cur.execute( + f"UPDATE annotations SET {assignments} WHERE id = %s AND {scope}", + (*(v for _, v in changing), annotation.id, *params), + ) + if cur.rowcount == 0: + raise KeyError(annotation.id) + except self._psycopg.errors.ExclusionViolation as exc: + raise TierOverlapError( + f"annotation {annotation.id} would overlap an existing one in tier " + f"{annotation.tier!r} (per-tier EXCLUDE constraint is active)" + ) from exc + def all(self) -> Iterator[Annotation]: scope, params = self._scope() with self._conn.cursor() as cur: diff --git a/lacing/store/sqlite.py b/lacing/store/sqlite.py index 13c54af2..653cd14f 100644 --- a/lacing/store/sqlite.py +++ b/lacing/store/sqlite.py @@ -344,14 +344,64 @@ def remove(self, annotation_id: UUID) -> Annotation | None: cur.execute("DELETE FROM annotations WHERE id = ?", (str(annotation_id),)) return ann + def update(self, annotation: Annotation) -> None: + """Replace the row with this id (and its R*Tree entry) in one transaction. + + Readers take the same lock (:meth:`_fetchall`), so on a connection + shared between threads — where a reader would otherwise see this + transaction's uncommitted state — they get the old annotation or the + new one, never neither (lacing#66). + """ + cols, values = _annotation_to_row(annotation) + assignments = [(c, v) for c, v in zip(cols, values) if c != "id"] + with self._lock: + cur = self._conn.cursor() + cur.execute("BEGIN IMMEDIATE") + try: + row = cur.execute( + "SELECT rowid FROM annotations WHERE id = ?", (str(annotation.id),) + ).fetchone() + if row is None: + raise KeyError(annotation.id) + rowid = row["rowid"] + cur.execute( + f"UPDATE annotations SET {', '.join(f'{c} = ?' for c, _ in assignments)} " + "WHERE id = ?", + (*(v for _, v in assignments), str(annotation.id)), + ) + cur.execute("DELETE FROM annotations_rtree WHERE rowid = ?", (rowid,)) + iv = annotation.interval + if iv is not None: + start_s = float(iv.start.to_fraction()) + end_s = float(iv.end.to_fraction()) + cur.execute( + "INSERT INTO annotations_rtree (rowid, start_seconds, end_seconds) " + "VALUES (?, ?, ?)", + ( + rowid, + math.nextafter(start_s, -math.inf), + math.nextafter(end_s, math.inf), + ), + ) + cur.execute("COMMIT") + except BaseException: + cur.execute("ROLLBACK") + raise + + def _fetchall(self, sql: str, params: tuple = ()) -> list[sqlite3.Row]: + """Run a read under the store lock, so it never lands inside a transaction + another thread has open on this connection (see :meth:`update`).""" + with self._lock: + return self._conn.execute(sql, params).fetchall() + def _fetch_one(self, annotation_id: UUID) -> Annotation | None: - row = self._conn.execute( + rows = self._fetchall( "SELECT * FROM annotations WHERE id = ?", (str(annotation_id),) - ).fetchone() - return None if row is None else _row_to_annotation(row) + ) + return _row_to_annotation(rows[0]) if rows else None def all(self) -> Iterator[Annotation]: - for row in self._conn.execute("SELECT * FROM annotations"): + for row in self._fetchall("SELECT * FROM annotations"): yield _row_to_annotation(row) def __len__(self) -> int: @@ -390,11 +440,11 @@ def __contains__(self, key: object) -> bool: return row is not None def __getitem__(self, key: TimeInterval) -> list[Annotation]: - rows = self._conn.execute( + rows = self._fetchall( "SELECT * FROM annotations WHERE " " start_value = ? AND start_rate = ? AND end_value = ? AND end_rate = ?", (key.start.value, key.start.rate, key.end.value, key.end.rate), - ).fetchall() + ) if not rows: raise KeyError(key) return [_row_to_annotation(r) for r in rows] @@ -489,7 +539,7 @@ def relate( # --- tier filters -------------------------------------------------- def by_tier(self, tier_name: str) -> Iterator[Annotation]: - for row in self._conn.execute( + for row in self._fetchall( "SELECT * FROM annotations WHERE tier = ?", (tier_name,) ): yield _row_to_annotation(row) @@ -553,7 +603,7 @@ def _candidates(self, query: TimeInterval) -> Iterator[Annotation]: "JOIN annotations_rtree r ON a.rowid = r.rowid " "WHERE r.start_seconds <= ? AND r.end_seconds >= ?" ) - for row in self._conn.execute(sql, (widened_end, widened_start)): + for row in self._fetchall(sql, (widened_end, widened_start)): yield _row_to_annotation(row) def _candidates_filtered( diff --git a/tests/test_store_postgres.py b/tests/test_store_postgres.py index 902d8574..0649d938 100644 --- a/tests/test_store_postgres.py +++ b/tests/test_store_postgres.py @@ -194,11 +194,11 @@ def test_remove_missing_returns_none(self, postgres_store): class TestAllen: def _populate(self, store): - store.add(_ann(_ti(0, 10))) # before query [30, 40) - store.add(_ann(_ti(15, 25))) # overlaps query - store.add(_ann(_ti(30, 40))) # equals query - store.add(_ann(_ti(33, 37))) # during query - store.add(_ann(_ti(50, 60))) # after query + store.add(_ann(_ti(0, 10))) # before query [30, 40) + store.add(_ann(_ti(15, 25))) # overlaps query + store.add(_ann(_ti(30, 40))) # equals query + store.add(_ann(_ti(33, 37))) # during query + store.add(_ann(_ti(50, 60))) # after query def test_intersects(self, postgres_store): self._populate(postgres_store) @@ -739,3 +739,82 @@ def test_mixed_provenance_refs_round_trip_through_jsonb(self, postgres_store): ) assert annotation_ids == [parent.id] assert asset_ids == [asset] + + +class TestUpdate: + """``update`` replaces by id in one statement (lacing#66).""" + + def test_replaces_and_can_move(self, postgres_store): + a = _ann(_ti(0, 2), text="old") + postgres_store.add(a) + moved = a.model_copy( + update={ + "reference": MediaRef(asset_id="blake3:test", interval=_ti(10, 12)), + "body": {"text": "new"}, + } + ) + postgres_store.update(moved) + got = list(postgres_store.all()) + assert [(x.id, x.body["text"]) for x in got] == [(a.id, "new")] + assert [x.id for x in postgres_store.intersects(_ti(10, 12))] == [a.id] + assert list(postgres_store.intersects(_ti(0, 2))) == [] + + def test_unknown_id_raises_and_changes_nothing(self, postgres_store): + a = _ann(_ti(0, 2), text="old") + postgres_store.add(a) + with pytest.raises(KeyError): + postgres_store.update(_ann(_ti(0, 2), text="ghost")) + assert [x.body["text"] for x in postgres_store.all()] == ["old"] + + def test_update_into_an_overlap_raises_tier_overlap_error(self, postgres_store): + postgres_store.add_tier(Tier("words"), enforce_no_overlap=True) + postgres_store.add(_ann(_ti(0, 100), tier="words")) + b = _ann(_ti(200, 300), tier="words") + postgres_store.add(b) + moved = b.model_copy( + update={ + "reference": MediaRef(asset_id="blake3:test", interval=_ti(50, 150)) + } + ) + with pytest.raises(TierOverlapError): + postgres_store.update(moved) + + def test_reader_on_another_connection_never_finds_the_id_missing( + self, postgres_store, postgresql + ): + import threading + + info = postgresql.info + reader_store = PostgresStore( + { + "host": info.host, + "port": info.port, + "user": info.user, + "password": info.password or "", + "dbname": info.dbname, + }, + rate=24000, + use_pool=False, + ) + target = _ann(_ti(0, 2), text="0") + postgres_store.add(target) + missing: list[str] = [] + stop = threading.Event() + + def reader(): + while not stop.is_set(): + if target.id not in {a.id for a in reader_store.all()}: + missing.append("all()") + + t = threading.Thread(target=reader) + t.start() + try: + for i in range(1, 101): + postgres_store.update( + target.model_copy(update={"body": {"text": str(i)}}) + ) + finally: + stop.set() + t.join() + reader_store.close() + assert missing == [] diff --git a/tests/test_store_update.py b/tests/test_store_update.py new file mode 100644 index 00000000..7b486699 --- /dev/null +++ b/tests/test_store_update.py @@ -0,0 +1,153 @@ +"""``store.update(annotation)`` — replace by id in one transaction (lacing#66). + +``remove(id)`` then ``add(...)`` is two commits: a reader between them sees no +annotation with that id, which is how a "read or create" caller made a second +singleton row. ``update`` must leave no such gap on any backend. The Postgres +backend's case lives in ``test_store_postgres.py`` (it needs a server). +""" + +from __future__ import annotations + +import sys +import threading +from uuid import uuid4 + +import pytest + +from lacing.model import Annotation, MediaRef, Provenance +from lacing.store import MappingStore, MemoryStore, SqliteStore +from lacing.tier import Tier +from lacing.time import RationalTime, TimeInterval + +#: Replacements the writer thread makes in the concurrency test. +N_REPLACEMENTS = 300 + + +def _ti(s: int, e: int) -> TimeInterval: + return TimeInterval(RationalTime(s), RationalTime(e)) + + +def _ann(interval: TimeInterval, *, id=None, tier="words", text="x") -> Annotation: + return Annotation( + id=id or uuid4(), + tier=tier, + reference=MediaRef(asset_id="blake3:test", interval=interval), + body={"text": text}, + body_schema_uri="annot://schema/word/v1", + provenance=Provenance( + was_generated_by="user:test", + was_attributed_to="test", + generated_at_time=RationalTime(0), + ), + ) + + +@pytest.fixture(params=["memory", "mapping", "sqlite"]) +def store(request, tmp_path): + if request.param == "memory": + s = MemoryStore() + elif request.param == "mapping": + s = MappingStore({}) + else: + s = SqliteStore(tmp_path / "u.annot", check_same_thread=False) + s.add_tier(Tier("words")) + yield s + if request.param == "sqlite": + s.close() + + +def _texts(store) -> dict: + return {a.id: a.body["text"] for a in store.all()} + + +def test_update_replaces_in_place(store): + a = _ann(_ti(0, 2), text="old") + other = _ann(_ti(0, 2), text="keep") + store.add(a) + store.add(other) + store.update(_ann(_ti(0, 2), id=a.id, text="new")) + assert _texts(store) == {a.id: "new", other.id: "keep"} + + +def test_update_can_move_the_annotation(store): + a = _ann(_ti(0, 2), text="old") + store.add(a) + store.update(_ann(_ti(10, 12), id=a.id, text="moved")) + assert _texts(store) == {a.id: "moved"} + assert [x.id for x in store.intersects(_ti(10, 12))] == [a.id] + assert list(store.intersects(_ti(0, 2))) == [] + assert _ti(0, 2) not in store + + +def test_update_unknown_id_raises_and_changes_nothing(store): + a = _ann(_ti(0, 2), text="old") + store.add(a) + with pytest.raises(KeyError): + store.update(_ann(_ti(0, 2), text="ghost")) + assert _texts(store) == {a.id: "old"} + + +@pytest.mark.parametrize("move", [False, True], ids=["same-interval", "moving"]) +def test_a_reader_never_finds_the_id_missing(store, move): + """The acceptance test: a reader thread and a replacing thread. + + The reader asks for the id by whole-store scan *and* by an interval query + that covers both places the annotation lives, so a gap in either the row or + the interval index shows up. + """ + target = _ann(_ti(0, 2), text="0") + store.add(target) + everywhere = _ti(0, 100) + missing: list[str] = [] + reads = [0] + stop = threading.Event() + + def reader(): + # A reader that dies must fail the test, not end it quietly. + try: + while not stop.is_set(): + if target.id not in {a.id for a in store.all()}: + missing.append("all()") + if target.id not in {a.id for a in store.intersects(everywhere)}: + missing.append("intersects()") + reads[0] += 1 + except Exception as exc: # noqa: BLE001 - reported by the assert below + missing.append(f"reader died: {exc!r}") + + old_interval = sys.getswitchinterval() + sys.setswitchinterval(1e-6) # make a thread switch inside the gap likely + t = threading.Thread(target=reader) + t.start() + try: + for i in range(1, N_REPLACEMENTS + 1): + iv = _ti(10, 12) if move and i % 2 else _ti(0, 2) + store.update(_ann(iv, id=target.id, text=str(i))) + finally: + stop.set() + t.join() + sys.setswitchinterval(old_interval) + assert missing == [] + assert reads[0] > 0 # the reader really ran + assert _texts(store) == {target.id: str(N_REPLACEMENTS)} + + +def test_mapping_store_update_persists(tmp_path): + backing: dict = {} + s = MappingStore(backing) + a = _ann(_ti(0, 2), text="old") + s.add(a) + s.update(_ann(_ti(4, 6), id=a.id, text="new")) + reopened = MappingStore(backing) + assert _texts(reopened) == {a.id: "new"} + assert len(backing) == 1 + + +def test_stores_with_a_lock_still_copy_and_pickle(): + import copy + import pickle + + s = MemoryStore() + s.add(_ann(_ti(0, 2))) + for clone in (copy.deepcopy(s), pickle.loads(pickle.dumps(s))): + assert [a.id for a in clone.all()] == [a.id for a in s.all()] + clone.update(next(iter(clone.all())).model_copy(update={"body": {"text": "y"}}))