diff --git a/src/crud/peer.py b/src/crud/peer.py index f37a8dd2..c3962d19 100644 --- a/src/crud/peer.py +++ b/src/crud/peer.py @@ -43,6 +43,7 @@ async def get_or_create_peers( peers: list[schemas.PeerCreate], *, _retry: bool = False, + _pending_invalidation: list[str] | None = None, ) -> GetOrCreateResult[list[models.Peer]]: """ Get an existing list of peers or create new peers if they don't exist. @@ -53,6 +54,8 @@ async def get_or_create_peers( workspace_name: Name of the workspace peers: List of peer creation schemas _retry: Whether to retry the operation + _pending_invalidation: Names of peers already mutated by a prior attempt, + whose cache keys must still be purged. See the retry branch below. Returns: GetOrCreateResult containing the list of peers and whether any were created @@ -123,11 +126,26 @@ async def get_or_create_peers( raise ConflictException( f"Unable to create or get peers: {peer_names}" ) from None - return await get_or_create_peers(db, workspace_name, peers, _retry=True) + # `begin_nested()` autoflushes the mutations above *before* opening the + # savepoint, so they are already committed-in-transaction and the rollback + # doesn't undo them — nor does it expire the now-clean ORM state. The retry + # would therefore compare already-updated values, find no change, and skip + # the purge. Carry the names forward so the invalidation can't be lost. + return await get_or_create_peers( + db, + workspace_name, + peers, + _retry=True, + _pending_invalidation=(_pending_invalidation or []) + + [p.name for p in changed_peers], + ) # Capture peer names eagerly so the closure holds plain strings, not ORM objects _cache_keys_to_invalidate = [ - peer_cache_key(workspace_name, p.name) for p in changed_peers + new_peers + peer_cache_key(workspace_name, name) + for name in dict.fromkeys( + (_pending_invalidation or []) + [p.name for p in changed_peers + new_peers] + ) ] async def _invalidate_peer_cache(): diff --git a/src/crud/scope.py b/src/crud/scope.py index ac6cf283..2ec1e22c 100644 --- a/src/crud/scope.py +++ b/src/crud/scope.py @@ -51,6 +51,7 @@ async def get_or_create_scopes( scopes: list[schemas.ScopeCreate], *, _retry: bool = False, + _pending_invalidation: list[str] | None = None, ) -> GetOrCreateResult[list[models.Peer]]: """ Get existing scopes or create new ones if they don't exist. @@ -66,6 +67,9 @@ async def get_or_create_scopes( db: Database session workspace_name: Name of the workspace scopes: List of scope creation schemas (unprefixed names) + _retry: Whether this is the retry attempt + _pending_invalidation: Names of scope peers already mutated by a prior + attempt, whose cache keys must still be purged. See the retry branch. Returns: GetOrCreateResult containing the backing peers and whether any were @@ -122,10 +126,24 @@ async def get_or_create_scopes( raise ConflictException( f"Unable to create or get scopes: {sorted(peer_names)}" ) from None - return await get_or_create_scopes(db, workspace_name, scopes, _retry=True) + # `begin_nested()` autoflushes the mutations above *before* opening the + # savepoint, so they survive the rollback and leave the ORM state clean — + # the retry would compare already-updated values, find no change, and skip + # the purge. Carry the names forward so the invalidation can't be lost. + return await get_or_create_scopes( + db, + workspace_name, + scopes, + _retry=True, + _pending_invalidation=(_pending_invalidation or []) + + [p.name for p in changed_peers], + ) _cache_keys_to_invalidate = [ - peer_cache_key(workspace_name, p.name) for p in changed_peers + new_peers + peer_cache_key(workspace_name, name) + for name in dict.fromkeys( + (_pending_invalidation or []) + [p.name for p in changed_peers + new_peers] + ) ] async def _invalidate_peer_cache(): diff --git a/src/routers/sessions.py b/src/routers/sessions.py index e59bbf76..6aff5fb3 100644 --- a/src/routers/sessions.py +++ b/src/routers/sessions.py @@ -41,7 +41,7 @@ router = APIRouter( # scopes facade so the observer mechanics stay internal. _SCOPES_ROUTE_GUIDANCE = ( "Scope membership is managed via the scopes routes " - "(/workspaces/{workspace_id}/scopes/{scope_id}/sessions) or the `scopes` " + "(/v3/workspaces/{workspace_id}/scopes/{scope_id}/sessions) or the `scopes` " "field at session creation." ) diff --git a/src/vector_store/__init__.py b/src/vector_store/__init__.py index 85fdd9ff..17eeec34 100644 --- a/src/vector_store/__init__.py +++ b/src/vector_store/__init__.py @@ -207,10 +207,10 @@ def _create_store_by_type(store_type: str) -> VectorStore: except ImportError as exc: raise RuntimeError( "VECTOR_STORE.TYPE is set to 'lancedb', but the 'lancedb' package " - "is not installed (for example on macOS Intel, where it is omitted " - "from dependencies because PyPI has no wheel). " - "Use TYPE 'pgvector' or 'turbopuffer', or install lancedb manually. " - f"Original import error: {exc}" + + "is not installed (for example on macOS Intel, where it is omitted " + + "from dependencies because PyPI has no wheel). " + + "Use TYPE 'pgvector' or 'turbopuffer', or install lancedb manually. " + + f"Original import error: {exc}" ) from exc return LanceDBVectorStore() diff --git a/tests/crud/test_get_or_create_retry_invalidation.py b/tests/crud/test_get_or_create_retry_invalidation.py new file mode 100644 index 00000000..f4576bf0 --- /dev/null +++ b/tests/crud/test_get_or_create_retry_invalidation.py @@ -0,0 +1,196 @@ +"""Regression tests for cache invalidation across the get_or_create retry path. + +`get_or_create_peers` / `get_or_create_scopes` mutate existing rows, then insert +new ones inside `db.begin_nested()`. A concurrent writer that creates one of those +rows first makes the insert raise `IntegrityError`, and the function retries. + +The subtlety: `begin_nested()` autoflushes the pending mutations *before* opening +the savepoint, so the rollback neither undoes them nor expires the ORM state. A +retry that recomputed "what changed" from that state would see no change and skip +the cache purge — while the row change still commits anyway, leaving the cache +stale until TTL. These tests pin the purge. + +The race is real (a second session committing a real row, producing a real +IntegrityError from the database); only its *timing* is made deterministic, by +hooking the one point that sits between the SELECT and the INSERT. +""" + +from unittest.mock import AsyncMock, patch + +import pytest +from nanoid import generate as generate_nanoid +from sqlalchemy.ext.asyncio import ( + AsyncEngine, + AsyncSession, + AsyncSessionTransaction, + async_sessionmaker, +) + +from src import crud, models, schemas +from src.crud.peer import peer_cache_key +from src.crud.scope import SCOPE_PEER_CONFIGURATION +from src.utils.scopes import scope_peer_name + + +class _RaceOnBeginNested: + """Commit a racing row on entry to `begin_nested()`, then delegate. + + That entry point is after the function's SELECT and metadata mutation but + before its INSERT flushes — precisely the window a real concurrent writer + has to slip through to trigger the IntegrityError retry. + """ + + _db: AsyncSession + _engine: AsyncEngine + _rows: list[models.Peer] + _real: AsyncSessionTransaction | None + fired: bool + + def __init__(self, db: AsyncSession, engine: AsyncEngine, rows: list[models.Peer]): + self._db = db + self._engine = engine + self._rows = rows + self._real = None + self.fired = False + + def __call__(self): + return self + + async def __aenter__(self): + if self._rows: + Session = async_sessionmaker(bind=self._engine, expire_on_commit=False) + async with Session() as other: + other.add_all(self._rows) + await other.commit() + self._rows = [] # race only once; the retry must succeed + self.fired = True + self._real = AsyncSession.begin_nested(self._db) + return await self._real.__aenter__() + + async def __aexit__(self, *exc_info: object): + assert self._real is not None + return await self._real.__aexit__(*exc_info) + + +@pytest.mark.asyncio +async def test_peer_retry_still_invalidates_mutated_peer( + db_session: AsyncSession, + db_engine: AsyncEngine, + sample_data: tuple[models.Workspace, models.Peer], +): + """A peer mutated before a losing race still gets its cache key purged.""" + test_workspace, existing_peer = sample_data + racer_name = str(generate_nanoid()) + + # Give the existing peer metadata we will then change, so it is a real update. + existing_peer.h_metadata = {"v": "old"} + await db_session.commit() + + race = _RaceOnBeginNested( + db_session, + db_engine, + [models.Peer(name=racer_name, workspace_name=test_workspace.name)], + ) + + with ( + patch("src.crud.peer.safe_cache_delete", new=AsyncMock()) as mock_delete, + patch.object(db_session, "begin_nested", race), + ): + result = await crud.get_or_create_peers( + db_session, + test_workspace.name, + [ + schemas.PeerCreate(name=existing_peer.name, metadata={"v": "new"}), + schemas.PeerCreate(name=racer_name), + ], + ) + await db_session.commit() + await result.post_commit() + + assert race.fired, "the race must actually have fired" + + purged = {call.args[0] for call in mock_delete.await_args_list} + assert ( + peer_cache_key(test_workspace.name, existing_peer.name) in purged + ), "the mutated peer's cache key must still be purged after the retry" + + # The mutation really did land — which is what makes a missed purge stale. + await db_session.refresh(existing_peer) + assert existing_peer.h_metadata == {"v": "new"} + + +@pytest.mark.asyncio +async def test_scope_retry_still_invalidates_mutated_scope( + db_session: AsyncSession, + db_engine: AsyncEngine, + sample_data: tuple[models.Workspace, models.Peer], +): + """Same guarantee for the scopes facade, which mirrors get_or_create_peers.""" + test_workspace, _ = sample_data + kept_scope, racing_scope = str(generate_nanoid()), str(generate_nanoid()) + + seeded = await crud.get_or_create_scopes( + db_session, + test_workspace.name, + [schemas.ScopeCreate(name=kept_scope, metadata={"v": "old"})], + ) + await db_session.commit() + await seeded.post_commit() + + # The racer creates the second scope's backing peer — as a *valid* scope peer, + # so the flow reaches the insert rather than tripping the legacy-collision 409. + race = _RaceOnBeginNested( + db_session, + db_engine, + [ + models.Peer( + name=scope_peer_name(racing_scope), + workspace_name=test_workspace.name, + configuration=dict(SCOPE_PEER_CONFIGURATION), + ) + ], + ) + + with ( + patch("src.crud.scope.safe_cache_delete", new=AsyncMock()) as mock_delete, + patch.object(db_session, "begin_nested", race), + ): + result = await crud.get_or_create_scopes( + db_session, + test_workspace.name, + [ + schemas.ScopeCreate(name=kept_scope, metadata={"v": "new"}), + schemas.ScopeCreate(name=racing_scope), + ], + ) + await db_session.commit() + await result.post_commit() + + assert race.fired, "the race must actually have fired" + + purged = {call.args[0] for call in mock_delete.await_args_list} + assert ( + peer_cache_key(test_workspace.name, scope_peer_name(kept_scope)) in purged + ), "the mutated scope peer's cache key must still be purged after the retry" + + +@pytest.mark.asyncio +async def test_peer_no_race_does_not_invalidate_unchanged_peer( + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], +): + """Baseline: with no race, an unchanged peer is not purged.""" + test_workspace, existing_peer = sample_data + existing_peer.h_metadata = {"v": "same"} + await db_session.commit() + + with patch("src.crud.peer.safe_cache_delete", new=AsyncMock()) as mock_delete: + result = await crud.get_or_create_peers( + db_session, + test_workspace.name, + [schemas.PeerCreate(name=existing_peer.name, metadata={"v": "same"})], + ) + await db_session.commit() + await result.post_commit() + + assert mock_delete.await_count == 0, "an unchanged peer must not be purged" diff --git a/tests/test_cache_redaction.py b/tests/test_cache_redaction.py index 0ff2bbc7..805b31c7 100644 --- a/tests/test_cache_redaction.py +++ b/tests/test_cache_redaction.py @@ -2,7 +2,9 @@ import pytest -from src.cache.client import _redact_cache_url +from src.cache.client import ( + _redact_cache_url, # pyright: ignore[reportPrivateUsage] +) class TestRedactCacheUrl: