Fix UniqueViolation in equivalent_identifiers_refresh cache rebuild (PP-4786)

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
+            ),
         )
 
 

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}