diff --git a/src/palace/manager/sqlalchemy/refresh_equivalents.py b/src/palace/manager/sqlalchemy/refresh_equivalents.py index 1042335b45..ea82c99c86 100644 --- a/src/palace/manager/sqlalchemy/refresh_equivalents.py +++ b/src/palace/manager/sqlalchemy/refresh_equivalents.py @@ -12,7 +12,11 @@ from __future__ import annotations +from collections.abc import Collection, Iterable +from itertools import batched + from sqlalchemy import and_, delete, select, union +from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.orm import Session from palace.manager.sqlalchemy.model.identifier import ( @@ -21,6 +25,43 @@ RecursiveEquivalencyCache, ) +# Number of rows sent per INSERT statement. +INSERT_CHUNK_SIZE = 10_000 + + +def _insert_cache_rows(session: Session, rows: Iterable[dict[str, int]]) -> None: + """ + Insert (parent_identifier_id, identifier_id) rows into the cache. + + Rows that already exist are skipped rather than raising a unique violation. + + Rows are inserted in a stable ``(parent_identifier_id, identifier_id)`` + order so that two transactions inserting overlapping keys don't + deadlock against each other. + """ + ordered = sorted( + rows, key=lambda row: (row["parent_identifier_id"], row["identifier_id"]) + ) + for chunk in batched(ordered, INSERT_CHUNK_SIZE): + session.execute( + pg_insert(RecursiveEquivalencyCache) + .values(list(chunk)) + .on_conflict_do_nothing( + index_elements=["parent_identifier_id", "identifier_id"] + ) + ) + + +def _delete_cache_rows(session: Session, parent_ids: Collection[int]) -> None: + """Delete every cache row belonging to any of the given parent identifiers.""" + if not parent_ids: + return + session.execute( + delete(RecursiveEquivalencyCache).where( + RecursiveEquivalencyCache.parent_identifier_id.in_(parent_ids) + ) + ) + def process_identifier_ids(session: Session, identifier_ids: frozenset[int]) -> None: """ @@ -58,25 +99,19 @@ def process_identifier_ids(session: Session, identifier_ids: frozenset[int]) -> .add_columns(Identifier.id) ) chained_identifiers = session.execute(qu).fetchall() + if not chained_identifiers: + return - # Delete old cache entries for every affected parent, then insert fresh ones. - completed: set[int] = set() - new_rows: list[RecursiveEquivalencyCache] = [] - for link_id, parent_id in chained_identifiers: - if parent_id not in completed: - session.execute( - delete(RecursiveEquivalencyCache).where( - RecursiveEquivalencyCache.parent_identifier_id == parent_id - ) - ) - new_rows.append( - RecursiveEquivalencyCache( - parent_identifier_id=parent_id, identifier_id=link_id - ) - ) - completed.add(parent_id) - - session.add_all(new_rows) + # Delete the old cache entries for every affected parent, then insert the + # fresh ones. + _delete_cache_rows(session, {parent_id for _, parent_id in chained_identifiers}) + _insert_cache_rows( + session, + ( + {"parent_identifier_id": parent_id, "identifier_id": link_id} + for link_id, parent_id in chained_identifiers + ), + ) def add_identity_equivalents(session: Session, batch_size: int = 200) -> None: @@ -86,6 +121,10 @@ def add_identity_equivalents(session: Session, batch_size: int = 200) -> None: This ensures that queries against the cache always return at least the identifier itself, even when it has no equivalencies. + + An Identifier created by another transaction after this scan begins gets its + self-reference from the ``Identifier`` creation listener, which may commit + before we do; the insert tolerates that (see :func:`_insert_cache_rows`). """ missing_q = ( select(Identifier.id) @@ -100,11 +139,13 @@ def add_identity_equivalents(session: Session, batch_size: int = 200) -> None: .execution_options(yield_per=batch_size) ) - for (identifier_id,) in session.execute(missing_q): - session.add( - RecursiveEquivalencyCache( - parent_identifier_id=identifier_id, identifier_id=identifier_id - ) + for partition in session.execute(missing_q).partitions(): + _insert_cache_rows( + session, + ( + {"parent_identifier_id": identifier_id, "identifier_id": identifier_id} + for (identifier_id,) in partition + ), ) diff --git a/tests/manager/sqlalchemy/test_refresh_equivalents.py b/tests/manager/sqlalchemy/test_refresh_equivalents.py index 938a824f6f..eae92d3ead 100644 --- a/tests/manager/sqlalchemy/test_refresh_equivalents.py +++ b/tests/manager/sqlalchemy/test_refresh_equivalents.py @@ -1,8 +1,12 @@ +import pytest + from palace.manager.sqlalchemy.model.identifier import ( Equivalency, RecursiveEquivalencyCache, ) from palace.manager.sqlalchemy.refresh_equivalents import ( + _delete_cache_rows, + _insert_cache_rows, add_identity_equivalents, process_identifier_ids, refresh_equivalent_identifiers, @@ -89,6 +93,52 @@ def test_replaces_stale_cache( assert b.id not in recursive_equivalency_cache.cache_for(a.id) + def test_tolerates_rows_the_delete_did_not_remove( + self, + db: DatabaseTransactionFixture, + recursive_equivalency_cache: RecursiveEquivalencyCacheFixture, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + # A concurrent writer can commit a cache row for a parent after we have + # already deleted that parent's rows, so the rows we are about to insert + # may collide. Simulate that by disabling the delete entirely: the insert + # must skip the existing rows rather than raise a unique violation. + a = db.identifier() + b = db.identifier() + db.session.add(Equivalency(input_id=a.id, output_id=b.id, strength=1.0)) + db.session.flush() + + process_identifier_ids(db.session, frozenset([a.id, b.id])) + db.session.flush() + assert recursive_equivalency_cache.cache_for(a.id) == {a.id, b.id} + + monkeypatch.setattr( + "palace.manager.sqlalchemy.refresh_equivalents._delete_cache_rows", + lambda session, parent_ids: None, + ) + process_identifier_ids(db.session, frozenset([a.id, b.id])) + db.session.flush() + + assert recursive_equivalency_cache.cache_for(a.id) == {a.id, b.id} + assert recursive_equivalency_cache.cache_for(b.id) == {a.id, b.id} + + def test_no_chained_identifiers_is_noop( + self, + db: DatabaseTransactionFixture, + ) -> None: + # A non-empty set of identifier IDs that don't exist passes the empty + # check but yields no chains, so the function returns without touching + # the cache. + db.identifier() + db.session.flush() + before = db.session.query(RecursiveEquivalencyCache).count() + + process_identifier_ids(db.session, frozenset([-1, -2])) + db.session.flush() + + # Nothing was deleted or inserted. + assert db.session.query(RecursiveEquivalencyCache).count() == before + class TestAddIdentityEquivalents: def test_adds_self_references( @@ -120,6 +170,102 @@ def test_skips_existing(self, db: DatabaseTransactionFixture) -> None: # No duplicate rows should be added. assert after == before + def test_sweeps_across_multiple_partitions( + self, + db: DatabaseTransactionFixture, + recursive_equivalency_cache: RecursiveEquivalencyCacheFixture, + ) -> None: + # With more missing identifiers than one yield_per partition, the sweep + # inserts self-references while the streaming cursor is still open and + # spans several partitions. A batch_size of 2 over 5 identifiers forces + # at least three partitions, exercising the interleaved-insert path. + identifiers = [db.identifier() for _ in range(5)] + db.session.flush() + recursive_equivalency_cache.drop() + + add_identity_equivalents(db.session, batch_size=2) + db.session.flush() + + for identifier in identifiers: + assert recursive_equivalency_cache.cache_for(identifier.id) == { + identifier.id + } + + +class TestInsertCacheRows: + def test_skips_existing_rows( + self, + db: DatabaseTransactionFixture, + recursive_equivalency_cache: RecursiveEquivalencyCacheFixture, + ) -> None: + # A row committed by a concurrent writer — the Identifier creation + # listener, or a second refresh chain — must not abort the insert. + a = db.identifier() + b = db.identifier() + db.session.flush() + # a and b already have self-references from the creation listener. + + rows = [ + {"parent_identifier_id": a.id, "identifier_id": a.id}, + {"parent_identifier_id": a.id, "identifier_id": b.id}, + ] + _insert_cache_rows(db.session, rows) + db.session.flush() + + assert recursive_equivalency_cache.cache_for(a.id) == {a.id, b.id} + + # Re-inserting the same rows is a no-op. + _insert_cache_rows(db.session, rows) + db.session.flush() + + assert recursive_equivalency_cache.cache_for(a.id) == {a.id, b.id} + + def test_empty_input(self, db: DatabaseTransactionFixture) -> None: + _insert_cache_rows(db.session, []) + assert db.session.query(RecursiveEquivalencyCache).count() == 0 + + def test_batches_across_multiple_chunks( + self, + db: DatabaseTransactionFixture, + recursive_equivalency_cache: RecursiveEquivalencyCacheFixture, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + # More rows than a single INSERT chunk must all be inserted, across + # multiple statements. Shrink the chunk size so a handful of rows spans + # several chunks. + monkeypatch.setattr( + "palace.manager.sqlalchemy.refresh_equivalents.INSERT_CHUNK_SIZE", 2 + ) + parent = db.identifier() + others = [db.identifier() for _ in range(5)] + db.session.flush() + recursive_equivalency_cache.drop() + + rows = [ + {"parent_identifier_id": parent.id, "identifier_id": identifier.id} + for identifier in [parent, *others] + ] + _insert_cache_rows(db.session, rows) + db.session.flush() + + assert recursive_equivalency_cache.cache_for(parent.id) == { + parent.id, + *(identifier.id for identifier in others), + } + + +class TestDeleteCacheRows: + def test_empty_input(self, db: DatabaseTransactionFixture) -> None: + # An empty parent set is a no-op: no DELETE is issued, so existing rows + # (e.g. the creation listener's self-references) are left in place. + db.identifier() + db.session.flush() + before = db.session.query(RecursiveEquivalencyCache).count() + + _delete_cache_rows(db.session, set()) + + assert db.session.query(RecursiveEquivalencyCache).count() == before + class TestRefreshEquivalentIdentifiers: def test_full_refresh( @@ -141,3 +287,18 @@ def test_full_refresh( assert recursive_equivalency_cache.cache_for(b.id) == {a.id, b.id} # c should have a self-reference. assert recursive_equivalency_cache.cache_for(c.id) == {c.id} + + def test_no_equivalencies( + self, + db: DatabaseTransactionFixture, + recursive_equivalency_cache: RecursiveEquivalencyCacheFixture, + ) -> None: + # With no equivalencies at all, the chain computation is skipped and + # only the self-reference sweep runs. + a = db.identifier() + db.session.flush() + recursive_equivalency_cache.drop() + + refresh_equivalent_identifiers(db.session) + + assert recursive_equivalency_cache.cache_for(a.id) == {a.id}