diff --git a/app/repositories/concurrency.py b/app/repositories/concurrency.py new file mode 100644 index 0000000..e5a0a84 --- /dev/null +++ b/app/repositories/concurrency.py @@ -0,0 +1,42 @@ +"""Shared helper for check-then-act get-or-create races (section 15/28). + +Every get_or_create_* across the repository layer follows check-then-act +(look up by natural key, insert if missing) rather than a single atomic +UPSERT. That's fine as long as nothing ingests concurrently for +overlapping data — which isn't a safe assumption once more than one +ingestion job can run at the same time (arq's own concurrency, season +ingestion cascading via `include_seasons`, and this bulk-import script's +worker pool all create exactly this situation): two concurrent callers can +both see "missing" and both try to INSERT, and the loser's INSERT hits a +unique-constraint violation instead of finding the winner's row. + +`create_or_recover` wraps the insert in a SAVEPOINT and, on that specific +failure, rolls back just the savepoint and re-queries for what the +concurrent caller committed — recovering instead of crashing. Originally +lived only in `reference_repository.py`; moved here once +`provider_repository.py:get_or_create_provider` needed the exact same +protection (found live: 50 concurrent bulk-import workers all racing to +create the one-and-only "tmdb" Provider row on first run crashed 49 of +them with `uq_providers_slug` violations).""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable + +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession + + +async def create_or_recover[T]( + session: AsyncSession, + create: Callable[[], Awaitable[T]], + recover: Callable[[], Awaitable[T | None]], +) -> T: + try: + async with session.begin_nested(): + return await create() + except IntegrityError: + existing = await recover() + if existing is None: + raise + return existing diff --git a/app/repositories/provider_repository.py b/app/repositories/provider_repository.py index e6d94a9..5b85d62 100644 --- a/app/repositories/provider_repository.py +++ b/app/repositories/provider_repository.py @@ -10,6 +10,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from app.models.enums import IngestionJobStatus, MatchConfidence from app.models.provider import IngestionJob, Provider, ProviderRecord +from app.repositories.concurrency import create_or_recover async def get_provider_by_slug(session: AsyncSession, slug: str) -> Provider | None: @@ -20,13 +21,20 @@ async def get_provider_by_slug(session: AsyncSession, slug: str) -> Provider | N async def get_or_create_provider( session: AsyncSession, slug: str, name: str, base_url: str | None = None ) -> Provider: - provider = await get_provider_by_slug(session, slug) - if provider is not None: + async def _lookup() -> Provider | None: + return await get_provider_by_slug(session, slug) + + existing = await _lookup() + if existing is not None: + return existing + + async def _create() -> Provider: + provider = Provider(slug=slug, name=name, base_url=base_url) + session.add(provider) + await session.flush() return provider - provider = Provider(slug=slug, name=name, base_url=base_url) - session.add(provider) - await session.flush() - return provider + + return await create_or_recover(session, _create, _lookup) async def get_provider_record( @@ -96,20 +104,27 @@ async def create_ingestion_job( entity_type: str | None = None, external_id: str | None = None, ) -> IngestionJob: - existing = await get_ingestion_job_by_idempotency_key(session, idempotency_key) + async def _lookup() -> IngestionJob | None: + return await get_ingestion_job_by_idempotency_key(session, idempotency_key) + + existing = await _lookup() if existing is not None: return existing - job = IngestionJob( - provider_id=provider_id, - job_type=job_type, - entity_type=entity_type, - external_id=external_id, - idempotency_key=idempotency_key, - status=IngestionJobStatus.PENDING, - ) - session.add(job) - await session.flush() - return job + + async def _create() -> IngestionJob: + job = IngestionJob( + provider_id=provider_id, + job_type=job_type, + entity_type=entity_type, + external_id=external_id, + idempotency_key=idempotency_key, + status=IngestionJobStatus.PENDING, + ) + session.add(job) + await session.flush() + return job + + return await create_or_recover(session, _create, _lookup) async def get_ingestion_job(session: AsyncSession, job_id: uuid.UUID) -> IngestionJob | None: diff --git a/app/repositories/reference_repository.py b/app/repositories/reference_repository.py index 68fdf83..786ef36 100644 --- a/app/repositories/reference_repository.py +++ b/app/repositories/reference_repository.py @@ -3,26 +3,17 @@ ingestion: genres, keywords, languages, countries, companies, collections, people, and external ID links. Every get_or_create_* here follows check-then-act (look up by natural key, -insert if missing) rather than a single atomic UPSERT. That's fine as long -as nothing ingests concurrently for overlapping data — which stopped being -true once season ingestion could be cascaded (section 15's `include_seasons`, -app/workers/tasks.py): arq runs multiple queued jobs concurrently, so two -seasons of the same series sharing a guest star, crew member, or company can -race to create it, and the loser's INSERT hits a unique-constraint violation -instead of finding the winner's row. `_create_or_recover` below wraps the -insert in a SAVEPOINT and, on that specific failure, rolls back just the -savepoint and re-queries for what the concurrent job committed — recovering -instead of aborting the whole ingestion job. +insert if missing) rather than a single atomic UPSERT — see +`app/repositories/concurrency.py:create_or_recover` for why that's only +safe wrapped in a SAVEPOINT-based race recovery, not on its own. """ from __future__ import annotations import re import uuid -from collections.abc import Awaitable, Callable from sqlalchemy import select -from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from app.models.collection import Collection @@ -33,6 +24,7 @@ from app.models.genre import Genre from app.models.keyword import Keyword from app.models.language_country import Country, Language from app.models.person import Person +from app.repositories.concurrency import create_or_recover as _create_or_recover def slugify(value: str) -> str: @@ -40,21 +32,6 @@ def slugify(value: str) -> str: return slug or "n-a" -async def _create_or_recover[T]( - session: AsyncSession, - create: Callable[[], Awaitable[T]], - recover: Callable[[], Awaitable[T | None]], -) -> T: - try: - async with session.begin_nested(): - return await create() - except IntegrityError: - existing = await recover() - if existing is None: - raise - return existing - - async def get_or_create_genre(session: AsyncSession, name: str) -> Genre: slug = slugify(name) diff --git a/tests/integration/test_provider_repository.py b/tests/integration/test_provider_repository.py new file mode 100644 index 0000000..4e54d30 --- /dev/null +++ b/tests/integration/test_provider_repository.py @@ -0,0 +1,44 @@ +"""Integration test for provider_repository.get_or_create_provider's +concurrency safety -- found broken live: 50 concurrent bulk-import workers +(app/scripts/bulk_import.py) all racing to create the one-and-only "tmdb" +Provider row on first run crashed 49 of them with a uq_providers_slug +violation, since it was a plain check-then-insert with no race recovery +(unlike its sibling get_or_create_* helpers in reference_repository.py, +which already used create_or_recover for exactly this reason -- now moved +to app/repositories/concurrency.py and shared by both). + +Needs genuinely separate DB connections/transactions to reproduce -- the +db_session fixture's single connection/transaction can't race against +itself, so this uses session_scope() per worker (matching +test_ingestion_retry.py), each a real, independent connection.""" + +from __future__ import annotations + +import asyncio +import uuid + +import pytest +from sqlalchemy import delete +from sqlalchemy.ext.asyncio import AsyncSession + +from app.core.database import session_scope +from app.models.provider import Provider +from app.repositories.provider_repository import get_or_create_provider + + +@pytest.mark.asyncio +async def test_get_or_create_provider_survives_concurrent_races(db_session: AsyncSession): + # db_session is unused directly but keeps the Postgres-reachability + # skip guard (see tests/conftest.py), matching test_ingestion_retry.py. + slug = f"race-test-{uuid.uuid4().hex[:8]}" + + async def _worker() -> Provider: + async with session_scope() as session: + return await get_or_create_provider(session, slug, "Race Test Provider") + + try: + results = await asyncio.gather(*(_worker() for _ in range(20))) + assert len({provider.id for provider in results}) == 1 + finally: + async with session_scope() as session: + await session.execute(delete(Provider).where(Provider.slug == slug))