From 27b48dd64e4595c0f9a1fd8eee08fa942bde68c9 Mon Sep 17 00:00:00 2001 From: Renn F Date: Sun, 28 Jun 2026 22:13:53 +0200 Subject: [PATCH] [F108] atomic replace_chunks: single-txn delete+insert closes reindex race --- roboco/services/optimal_brain/indexes/base.py | 12 +- roboco/services/optimal_brain/vector_store.py | 49 +++++++ .../test_replace_chunks_atomic.py | 134 ++++++++++++++++++ .../optimal_brain/test_replace_on_reingest.py | 24 +++- 4 files changed, 214 insertions(+), 5 deletions(-) create mode 100644 tests/unit/services/optimal_brain/test_replace_chunks_atomic.py diff --git a/roboco/services/optimal_brain/indexes/base.py b/roboco/services/optimal_brain/indexes/base.py index b82bcf99..624961c9 100644 --- a/roboco/services/optimal_brain/indexes/base.py +++ b/roboco/services/optimal_brain/indexes/base.py @@ -465,8 +465,16 @@ class BaseIndexPlugin(ABC): ) if self.replace_on_reingest: - await store.delete_by_source(doc.source) - await store.add_chunks(chunks_with_embeddings) + # Atomic delete + insert in one transaction (F108): the prior + # separate delete_by_source + add_chunks awaited on two pool + # connections, so concurrent re-indexes of the same source + # interleaved and produced duplicate chunk rows. replace_chunks + # does both on a single connection inside a transaction, so the + # whole replace is atomic (last committer wins, no duplicates, and + # a failed insert reverts the delete — no data loss). + await store.replace_chunks(doc.source, chunks_with_embeddings) + else: + await store.add_chunks(chunks_with_embeddings) return len(chunks) def _prepare_docs_for_batch( diff --git a/roboco/services/optimal_brain/vector_store.py b/roboco/services/optimal_brain/vector_store.py index 2d66011a..b453d7a4 100644 --- a/roboco/services/optimal_brain/vector_store.py +++ b/roboco/services/optimal_brain/vector_store.py @@ -236,6 +236,55 @@ class VectorStore: source, ) + async def replace_chunks(self, source: str, chunks: list[Chunk]) -> None: + """Atomically replace every chunk row for *source* with *chunks*. + + Deletes the source's existing rows and inserts the new embedded chunks + on a SINGLE connection inside a SINGLE transaction, so the whole + replace is atomic. ``delete_by_source`` + ``add_chunks`` were two + separate awaits on two separate pool connections; two concurrent + re-indexes of the same source interleaved across those connections and + produced duplicate chunk rows, and an insert failure after a + successful delete lost the source's index rows. Wrapping both in one + transaction closes the race (concurrent replacers serialize on the + row locks; the last committer wins with no duplicates) and reverts the + delete if the insert fails (no data loss). + + Chunks without an ``embedding`` are silently skipped (matches + ``add_chunks``); an empty ``chunks`` list still clears the source + (matches the prior delete-then-no-op-add behavior). + """ + records = [ + ( + chunk.text, + chunk.source, + _vec_to_str(chunk.embedding), + json.dumps(chunk.metadata or {}), + ) + for chunk in chunks + if chunk.embedding is not None + ] + pool = self._require_pool() + # One acquire, one transaction: the DELETE and INSERT share a single + # connection and commit together (or roll back together on failure). + async with pool.acquire() as conn, conn.transaction(): + await conn.execute( + self._q("DELETE FROM {table} WHERE source = $1"), + source, + ) + if records: + await conn.executemany( + self._q( + """ + INSERT INTO {table} + (content, source, embedding, metadata) + VALUES + ($1, $2, $3::vector, $4::jsonb) + """ + ), + records, + ) + # ------------------------------------------------------------------ # Read # ------------------------------------------------------------------ diff --git a/tests/unit/services/optimal_brain/test_replace_chunks_atomic.py b/tests/unit/services/optimal_brain/test_replace_chunks_atomic.py new file mode 100644 index 00000000..3d4a258a --- /dev/null +++ b/tests/unit/services/optimal_brain/test_replace_chunks_atomic.py @@ -0,0 +1,134 @@ +"""F108 — ``VectorStore.replace_chunks`` must be a single atomic transaction. + +The replace-on-reingest path used to be ``delete_by_source`` (one pool +connection) followed by ``add_chunks`` (a *second* pool connection). Two +concurrent re-indexes of the same source interleaved across those two +connections and produced duplicate chunk rows; an add failure after a +successful delete also lost the source's index rows. The fix is a single +``replace_chunks(source, chunks)`` that deletes + inserts on ONE connection +inside ONE asyncpg transaction, so the whole replace is atomic. + +These tests mock the asyncpg pool/connection to assert the atomicity +invariant (single acquire, transaction entered, delete + insert on the +same connection) without standing up a pgvector DB. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from roboco.services.optimal_brain.text_chunker import Chunk +from roboco.services.optimal_brain.vector_store import VectorStore + + +def _chunk(text: str, source: str, embedding: list[float]) -> Chunk: + return Chunk(text=text, source=source, metadata={}, embedding=embedding) + + +def _make_store_with_conn() -> tuple[VectorStore, MagicMock, MagicMock]: + """Build a VectorStore wired to a mock pool that yields one mock conn. + + Returns ``(store, conn, pool)`` so a test can assert call order on the + same connection instance that the transaction/DELETE/INSERT ran on. + """ + conn = MagicMock() + # asyncpg's transaction() is an async context manager. + tx = MagicMock() + tx.__aenter__ = AsyncMock(return_value=tx) + tx.__aexit__ = AsyncMock(return_value=None) + conn.transaction = MagicMock(return_value=tx) + conn.execute = AsyncMock() + conn.executemany = AsyncMock() + + pool = MagicMock() + cm = MagicMock() + cm.__aenter__ = AsyncMock(return_value=conn) + cm.__aexit__ = AsyncMock(return_value=None) + pool.acquire = MagicMock(return_value=cm) + + store = VectorStore( + dsn="postgresql://test", table_name="chunks_test", vector_dimension=3 + ) + store._pool = pool # injected for the test + return store, conn, pool + + +@pytest.mark.asyncio +async def test_replace_chunks_deletes_then_inserts_in_one_transaction() -> None: + """DELETE + INSERT run on a SINGLE connection inside a transaction.""" + store, conn, pool = _make_store_with_conn() + source = "roboco://standards/general/std-1" + chunks = [ + _chunk("aaa" * 100, source, [0.1, 0.2, 0.3]), + _chunk("bbb" * 100, source, [0.4, 0.5, 0.6]), + ] + + await store.replace_chunks(source, chunks) + + # Exactly one connection was acquired (not one for delete + one for add). + assert pool.acquire.call_count == 1 + # The transaction was entered on that same connection. + assert conn.transaction.call_count == 1 + # DELETE by source, then the batch INSERT, both on the same connection. + assert conn.execute.await_count == 1 + delete_sql, delete_source = conn.execute.await_args.args + assert "DELETE FROM" in delete_sql + assert "source = $1" in delete_sql + assert delete_source == source + assert conn.executemany.await_count == 1 + insert_sql, insert_records = conn.executemany.await_args.args + assert "INSERT INTO" in insert_sql + assert len(insert_records) == len(chunks) + + +@pytest.mark.asyncio +async def test_replace_chunks_with_no_chunks_still_clears_source() -> None: + """An empty reingest still deletes the source's existing rows (matches + the prior ``delete_by_source`` then no-op ``add_chunks`` behavior).""" + store, conn, _pool = _make_store_with_conn() + source = "roboco://standards/general/std-1" + + await store.replace_chunks(source, []) + + assert conn.execute.await_count == 1 # DELETE ran + assert conn.executemany.await_count == 0 # nothing to insert + + +@pytest.mark.asyncio +async def test_replace_chunks_atomic_on_insert_failure() -> None: + """If the INSERT fails, the transaction rolls back — the source's old + rows are NOT left deleted (no retrieval data loss). The DELETE and + INSERT share one transaction, so a mid-replace failure reverts both.""" + store, conn, _pool = _make_store_with_conn() + source = "roboco://standards/general/std-1" + chunks = [_chunk("aaa" * 100, source, [0.1, 0.2, 0.3])] + # The INSERT raises inside the transaction. + conn.executemany.side_effect = RuntimeError("insert blew up") + + with pytest.raises(RuntimeError): + await store.replace_chunks(source, chunks) + + # The transaction context was entered; the raised error propagates out of + # the `async with conn.transaction()` block, so asyncpg rolls it back. + assert conn.transaction.call_count == 1 + assert conn.execute.await_count == 1 # DELETE did run (inside the txn)… + assert conn.executemany.await_count == 1 # …but the INSERT raised → rollback + + +@pytest.mark.asyncio +async def test_replace_chunks_skips_chunks_without_embeddings() -> None: + """Chunks lacking an embedding are dropped before insert (matches + ``add_chunks``).""" + store, conn, _pool = _make_store_with_conn() + source = "roboco://standards/general/std-1" + chunks = [ + _chunk("with-emb", source, [0.1, 0.2, 0.3]), + Chunk(text="no-emb", source=source, metadata={}), # no embedding + ] + + await store.replace_chunks(source, chunks) + + assert conn.execute.await_count == 1 # DELETE ran + _sql, records = conn.executemany.await_args.args + assert len(records) == 1 # only the embedded chunk inserted diff --git a/tests/unit/services/optimal_brain/test_replace_on_reingest.py b/tests/unit/services/optimal_brain/test_replace_on_reingest.py index 7443ccb9..fd4e62f3 100644 --- a/tests/unit/services/optimal_brain/test_replace_on_reingest.py +++ b/tests/unit/services/optimal_brain/test_replace_on_reingest.py @@ -42,19 +42,37 @@ def test_conversations_plugin_appends_not_replaces() -> None: @pytest.mark.asyncio -async def test_reingest_deletes_existing_source_chunks_when_replacing( +async def test_reingest_replaces_source_chunks_atomically( monkeypatch: pytest.MonkeyPatch, ) -> None: + """Re-ingest must replace the source's chunks in ONE atomic call (F108). + + The replace path used to be two separate awaits — ``delete_by_source`` + then ``add_chunks`` — each acquiring its own pool connection. Two + concurrent re-indexes of the same source interleaved (A.delete, B.delete, + A.add, B.add) and produced duplicate chunk rows. The fix collapses the + delete + insert into a single transactional ``replace_chunks`` call so + the whole replace is atomic (concurrent replacers serialize; the last + committer wins with no duplicates). + """ plugin = StandardsIndexPlugin() source = "roboco://standards/general/std-1" store = _wire_plugin(plugin, source, monkeypatch) + store.replace_chunks = AsyncMock() doc = Document(content="x" * 250, source=source, metadata={}) count = await plugin._chunk_filter_embed_store(doc, {}) assert count == 1 - store.delete_by_source.assert_awaited_once_with(source) - store.add_chunks.assert_awaited_once() + store.replace_chunks.assert_awaited_once() + # The atomic path must NOT also issue the separate delete/add calls — + # that would re-open the non-atomic race the single call closes. + store.delete_by_source.assert_not_awaited() + store.add_chunks.assert_not_awaited() + # The source and the embedded chunks are passed through verbatim. + called_source, called_chunks = store.replace_chunks.await_args.args + assert called_source == source + assert len(called_chunks) == 1 @pytest.mark.asyncio