mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
[F108] atomic replace_chunks: single-txn delete+insert closes reindex race
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user