fix(retain): close to_unit_id deferred-FK race on memory_links inserts (#1882) (#1894)

The memory_links → memory_units FKs are DEFERRABLE INITIALLY DEFERRED
(migration 9f8e7d6c5b4a), so an INSERT into memory_links takes no lock on
the referenced parent rows until COMMIT. Temporal and ANN link inserts
reference a *pre-existing* neighbor unit as to_unit_id (graph maintenance
also references a pre-existing from_unit_id). A concurrent transaction that
commits a DELETE of that unit in the window between the link INSERT and our
COMMIT — consolidation pruning observation units, document re-tracking —
makes the deferred check fail at COMMIT with
fk_memory_links_to_unit_id_memory_units, failing the async op with no retry.

#1795/#1805 only removed one *deleter* (sibling async children sharing a
document_id) for the from_unit_id side; the to_unit_id side, and any other
deleter, stayed uncovered.

Fix: in the PostgreSQL bulk link insert, lock the referenced parent units
FOR KEY SHARE via a CTE in the *same* INSERT statement. The lock blocks a
concurrent DELETE until our transaction commits and is held through the
deferred check; the INSERT only takes links whose endpoints are in the
locked set, so endpoints that already vanished are dropped. Folding it into
the one INSERT keeps this to a single round-trip — no extra query and no
surrounding transaction — so retain's perf characteristics are unchanged.
A WHERE EXISTS guard can't fix this (the row passes the check, then is
deleted before the deferred check runs). Oracle's FK is immediate (no such
window) and keeps its existing exists_clause path.

Adds a deterministic regression test that hand-drives the connection
interleaving (no sleeps): insert link on A (uncommitted) → delete neighbor
on B → commit A. Pre-fix this raises the FK violation; post-fix B blocks on
A's lock and the link commits cleanly.

Author: nicoloboschi

hindsight-api-slim/hindsight_api/engine/db/ops_postgresql.py

@@ -200,6 +200,23 @@ async def bulk_insert_links(
         exists_clause: str,
         chunk_size: int = 5000,
     ) -> None:
+        # exists_clause is unused on PostgreSQL: the memory_links → memory_units
+        # FKs are DEFERRABLE INITIALLY DEFERRED, so an INSERT takes no lock on the
+        # referenced parent rows until COMMIT — a concurrent committed DELETE in
+        # that window (consolidation pruning observations, document re-tracking)
+        # trips fk_memory_links_{to,from}_unit_id_memory_units at COMMIT (#1882),
+        # and a WHERE EXISTS guard can't prevent it (the row passes the check,
+        # then is deleted before the deferred check runs). Instead a CTE locks the
+        # referenced units FOR KEY SHARE in the *same statement*: the lock blocks a
+        # concurrent DELETE until our transaction commits and is held through the
+        # deferred check, and the INSERT only takes links whose endpoints are in
+        # the locked set, so rows that already vanished are dropped. Folding it
+        # into the one INSERT keeps this to a single round-trip — no extra query
+        # and no surrounding transaction needed. (Oracle's immediate FK has no
+        # such window and uses exists_clause via its own bulk_insert_links.)
+        from ..schema import fq_table
+
+        mu_table = fq_table("memory_units")
         from_ids = [lnk[0] for lnk in sorted_links]
         to_ids = [lnk[1] for lnk in sorted_links]
         types = [lnk[2] for lnk in sorted_links]
@@ -208,24 +225,37 @@ async def bulk_insert_links(
 
         for chunk_start in range(0, len(sorted_links), chunk_size):
             chunk_end = min(chunk_start + chunk_size, len(sorted_links))
+            chunk_from = from_ids[chunk_start:chunk_end]
+            chunk_to = to_ids[chunk_start:chunk_end]
+            # Distinct referenced parents, sorted so concurrent inserters acquire
+            # the row-share locks in a consistent order (avoids deadlocks; same
+            # convention as the (from, to) link sort).
+            referenced = sorted({str(x) for x in chunk_from} | {str(x) for x in chunk_to})
             await conn.execute(
                 f"""
+                WITH locked AS (
+                    SELECT id FROM {mu_table}
+                    WHERE id = ANY($7::uuid[])
+                    ORDER BY id
+                    FOR KEY SHARE
+                )
                 INSERT INTO {table}
                     (from_unit_id, to_unit_id, link_type, weight, entity_id, bank_id)
                 SELECT f, t, tp, w, e, $6
                 FROM unnest($1::uuid[], $2::uuid[], $3::text[], $4::float8[], $5::uuid[])
-                    AS t(f, t, tp, w, e)
-                {exists_clause}
+                    AS u(f, t, tp, w, e)
+                WHERE f IN (SELECT id FROM locked) AND t IN (SELECT id FROM locked)
                 ON CONFLICT (from_unit_id, to_unit_id, link_type,
                              COALESCE(entity_id, '{nil_entity_uuid}'::uuid))
                 DO NOTHING
                 """,
-                from_ids[chunk_start:chunk_end],
-                to_ids[chunk_start:chunk_end],
+                chunk_from,
+                chunk_to,
                 types[chunk_start:chunk_end],
                 weights[chunk_start:chunk_end],
                 entity_ids[chunk_start:chunk_end],
                 bank_id,
+                referenced,
                 timeout=300,
             )
 

hindsight-api-slim/tests/test_memory_links_to_unit_id_concurrent_delete.py

@@ -0,0 +1,126 @@
+"""Regression for #1882 — the ``to_unit_id`` side of the memory_links FK race.
+
+Mirror of #1795 (which covered ``from_unit_id``). Temporal / ANN link inserts
+reference a **pre-existing neighbor** unit as ``to_unit_id`` and bulk-insert with
+``skip_exists_check=True`` (see ``create_temporal_links_batch_per_fact`` /
+``compute_semantic_links_ann``). The FK
+``fk_memory_links_to_unit_id_memory_units`` is ``DEFERRABLE INITIALLY DEFERRED``
+(migration ``9f8e7d6c5b4a``), so the INSERT takes **no lock** on the parent
+``memory_units`` row — the check is pushed to COMMIT. A concurrent committed
+DELETE of that neighbor (consolidation pruning observation units at
+``consolidator.py``; document re-tracking at ``fact_storage.py``) then makes the
+deferred check fail at COMMIT with the exact violation reported in #1882.
+
+This is normally a timing race, but it is made **fully deterministic** here by
+driving the two connections by hand (no sleeps):
+
+  1. conn A opens a txn and inserts a temporal link ``(from, neighbor)`` via the
+     real ``_bulk_insert_links`` path — but does NOT commit.
+  2. conn B deletes the neighbor unit and commits.
+  3. conn A commits.
+
+Pre-fix: step 2 succeeds instantly (A holds no lock on the neighbor), the
+neighbor is gone, and step 3 raises ``ForeignKeyViolationError`` → this test
+fails, reproducing #1882.
+
+Post-fix (A locks referenced parents ``FOR KEY SHARE`` before inserting): step 2
+blocks on A's lock and is cut off by ``statement_timeout``; the neighbor
+survives and step 3 commits cleanly → this test passes.
+"""
+
+import uuid
+
+import asyncpg
+import pytest
+
+from hindsight_api.engine.db.ops_postgresql import PostgreSQLOps
+from hindsight_api.engine.db.postgresql import PostgresConnection
+from hindsight_api.engine.retain.link_utils import _bulk_insert_links
+
+
+async def _insert_unit(conn: asyncpg.Connection, bank_id: str) -> str:
+    """Insert one committed memory_unit (autocommit) and return its id as text."""
+    return await conn.fetchval(
+        """
+        INSERT INTO memory_units (bank_id, text, event_date, fact_type)
+        VALUES ($1, 'unit', now(), 'world')
+        RETURNING id::text
+        """,
+        bank_id,
+    )
+
+
+@pytest.mark.asyncio
+async def test_to_unit_id_survives_concurrent_neighbor_delete(pg0_db_url):
+    bank_id = f"fk-to-race-{uuid.uuid4().hex}"
+    ops = PostgreSQLOps()
+
+    # Three independent connections: a setup/cleanup conn, the retain txn (A),
+    # and the concurrent deleter (B). Independent connections give us full,
+    # deterministic control over transaction boundaries.
+    setup = await asyncpg.connect(pg0_db_url)
+    conn_a = await asyncpg.connect(pg0_db_url)
+    conn_b = await asyncpg.connect(pg0_db_url)
+    try:
+        # Two committed units: `from_id` plays the freshly-inserted fact,
+        # `neighbor_id` plays the pre-existing temporal neighbor (to_unit_id).
+        from_id = await _insert_unit(setup, bank_id)
+        neighbor_id = await _insert_unit(setup, bank_id)
+
+        # --- Step 1: conn A inserts the temporal link, does NOT commit ---
+        tx_a = conn_a.transaction()
+        await tx_a.start()
+        await _bulk_insert_links(
+            PostgresConnection(conn_a),
+            [(from_id, neighbor_id, "temporal", 1.0, None)],
+            bank_id=bank_id,
+            skip_exists_check=True,  # the racy path: no EXISTS guard
+            ops=ops,
+        )
+
+        # --- Step 2: conn B deletes the neighbor and commits (autocommit) ---
+        # statement_timeout bounds the post-fix path, where A holds FOR KEY
+        # SHARE on the neighbor and this DELETE must block.
+        await conn_b.execute("SET statement_timeout = '1000ms'")
+        neighbor_deleted = False
+        try:
+            await conn_b.execute("DELETE FROM memory_units WHERE id = $1", uuid.UUID(neighbor_id))
+            neighbor_deleted = True
+        except asyncpg.QueryCanceledError:
+            # Blocked on A's lock until statement_timeout — the fixed behavior.
+            neighbor_deleted = False
+
+        # --- Step 3: conn A commits — deferred FK check fires here ---
+        fk_error: Exception | None = None
+        try:
+            await tx_a.commit()
+        except asyncpg.ForeignKeyViolationError as exc:
+            fk_error = exc
+            await conn_a.execute("ROLLBACK")
+
+        assert fk_error is None, (
+            "memory_links.to_unit_id FK violation at COMMIT (#1882): a concurrent "
+            "DELETE removed the referenced neighbor unit between the deferred-FK "
+            "link INSERT and COMMIT. The link insert must lock referenced parent "
+            f"units (FOR KEY SHARE) or drop vanished ones. neighbor_deleted="
+            f"{neighbor_deleted}. {fk_error}"
+        )
+
+        # Post-fix end state: the neighbor was protected, so the link persists.
+        assert not neighbor_deleted, (
+            "Expected the concurrent DELETE to block on A's FOR KEY SHARE lock "
+            "and be cut off by statement_timeout; it succeeded instead, meaning "
+            "the parent row was not locked before the link insert."
+        )
+        link_count = await setup.fetchval(
+            "SELECT count(*) FROM memory_links WHERE from_unit_id = $1 AND to_unit_id = $2",
+            uuid.UUID(from_id),
+            uuid.UUID(neighbor_id),
+        )
+        assert link_count == 1
+    finally:
+        await setup.execute("DELETE FROM memory_links WHERE bank_id = $1", bank_id)
+        await setup.execute("DELETE FROM memory_units WHERE bank_id = $1", bank_id)
+        await setup.close()
+        await conn_a.close()
+        await conn_b.close()