Fix race condition in get_or_create_* under concurrent ingestion

Every get_or_create_* helper (genre, keyword, language, country, company, collection, person) followed check-then-act: look up by natural key, insert if missing. That was safe as long as nothing ingested concurrently for overlapping data -- which stopped being true the moment ingest_series_task could cascade season ingestion (include_seasons), since arq runs multiple queued jobs concurrently. Two seasons of the same series sharing a guest star, crew member, or company can now race to create it: both see "not found", both try to INSERT, and the loser hits a unique-constraint violation instead of finding the winner's row. Reproduced live: cascading Better Call Saul (60059) with include_seasons crashed a season job with "duplicate key value violates unique constraint uq_external_id_source" on a recurring guest actor shared across two seasons ingesting at the same time. _create_or_recover wraps each 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. Added unit tests for the retry mechanism itself; genuinely reproducing the underlying race in an automated test would need real Postgres MVCC timing, which isn't something a deterministic test can force reliably. Co-Authored-By: Claude Sonnet 5 <[email protected]>
2026-09-05 02:27:07 -05:00
parent 64169c4aee
commit dd3ff9764b
2 changed files with 261 additions and 100 deletions
@@ -1,13 +1,28 @@
 """Get-or-create helpers for shared reference/lookup data used during
 ingestion: genres, keywords, languages, countries, companies, collections,
-people, and external ID links."""
+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.
+"""

 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
@@ -25,50 +40,95 @@ 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)
-    result = await session.execute(select(Genre).where(Genre.slug == slug))
-    genre = result.scalar_one_or_none()
-    if genre is not None:
+
+    async def _lookup() -> Genre | None:
+        result = await session.execute(select(Genre).where(Genre.slug == slug))
+        return result.scalar_one_or_none()
+
+    existing = await _lookup()
+    if existing is not None:
+        return existing
+
+    async def _create() -> Genre:
+        genre = Genre(name=name, slug=slug)
+        session.add(genre)
+        await session.flush()
        return genre
-    genre = Genre(name=name, slug=slug)
-    session.add(genre)
-    await session.flush()
-    return genre
+
+    return await _create_or_recover(session, _create, _lookup)


 async def get_or_create_keyword(session: AsyncSession, name: str) -> Keyword:
    slug = slugify(name)
-    result = await session.execute(select(Keyword).where(Keyword.slug == slug))
-    keyword = result.scalar_one_or_none()
-    if keyword is not None:
+
+    async def _lookup() -> Keyword | None:
+        result = await session.execute(select(Keyword).where(Keyword.slug == slug))
+        return result.scalar_one_or_none()
+
+    existing = await _lookup()
+    if existing is not None:
+        return existing
+
+    async def _create() -> Keyword:
+        keyword = Keyword(name=name, slug=slug)
+        session.add(keyword)
+        await session.flush()
        return keyword
-    keyword = Keyword(name=name, slug=slug)
-    session.add(keyword)
-    await session.flush()
-    return keyword
+
+    return await _create_or_recover(session, _create, _lookup)


 async def get_or_create_language(session: AsyncSession, code: str, english_name: str) -> Language:
-    result = await session.execute(select(Language).where(Language.code == code))
-    language = result.scalar_one_or_none()
-    if language is not None:
+    async def _lookup() -> Language | None:
+        result = await session.execute(select(Language).where(Language.code == code))
+        return result.scalar_one_or_none()
+
+    existing = await _lookup()
+    if existing is not None:
+        return existing
+
+    async def _create() -> Language:
+        language = Language(code=code, english_name=english_name)
+        session.add(language)
+        await session.flush()
        return language
-    language = Language(code=code, english_name=english_name)
-    session.add(language)
-    await session.flush()
-    return language
+
+    return await _create_or_recover(session, _create, _lookup)


 async def get_or_create_country(session: AsyncSession, code: str, name: str) -> Country:
-    result = await session.execute(select(Country).where(Country.code == code))
-    country = result.scalar_one_or_none()
-    if country is not None:
+    async def _lookup() -> Country | None:
+        result = await session.execute(select(Country).where(Country.code == code))
+        return result.scalar_one_or_none()
+
+    existing = await _lookup()
+    if existing is not None:
+        return existing
+
+    async def _create() -> Country:
+        country = Country(code=code, name=name)
+        session.add(country)
+        await session.flush()
        return country
-    country = Country(code=code, name=name)
-    session.add(country)
-    await session.flush()
-    return country
+
+    return await _create_or_recover(session, _create, _lookup)


 async def find_company_by_external_id(
@@ -95,22 +155,29 @@ async def get_or_create_company(
    company_type: CompanyType = CompanyType.PRODUCTION_COMPANY,
    country_code: str | None = None,
 ) -> Company:
-    company = await find_company_by_external_id(session, provider_id, external_id)
-    if company is not None:
-        return company
-    company = Company(name=name, company_type=company_type, country_code=country_code)
-    session.add(company)
-    await session.flush()
-    session.add(
-        ExternalID(
-            entity_type="company",
-            entity_id=company.id,
-            provider_id=provider_id,
-            external_id=external_id,
+    async def _lookup() -> Company | None:
+        return await find_company_by_external_id(session, provider_id, external_id)
+
+    existing = await _lookup()
+    if existing is not None:
+        return existing
+
+    async def _create() -> Company:
+        company = Company(name=name, company_type=company_type, country_code=country_code)
+        session.add(company)
+        await session.flush()
+        session.add(
+            ExternalID(
+                entity_type="company",
+                entity_id=company.id,
+                provider_id=provider_id,
+                external_id=external_id,
+            )
        )
-    )
-    await session.flush()
-    return company
+        await session.flush()
+        return company
+
+    return await _create_or_recover(session, _create, _lookup)


 async def find_collection_by_external_id(
@@ -136,30 +203,37 @@ async def get_or_create_collection(
    name: str,
    overview: str | None = None,
 ) -> Collection:
-    collection = await find_collection_by_external_id(session, provider_id, external_id)
-    if collection is not None:
-        return collection
-    base_slug = slugify(name)
-    slug = base_slug
-    suffix = 1
-    while (
-        await session.execute(select(Collection.id).where(Collection.slug == slug))
-    ).scalar_one_or_none() is not None:
-        suffix += 1
-        slug = f"{base_slug}-{suffix}"
-    collection = Collection(name=name, overview=overview, slug=slug)
-    session.add(collection)
-    await session.flush()
-    session.add(
-        ExternalID(
-            entity_type="collection",
-            entity_id=collection.id,
-            provider_id=provider_id,
-            external_id=external_id,
+    async def _lookup() -> Collection | None:
+        return await find_collection_by_external_id(session, provider_id, external_id)
+
+    existing = await _lookup()
+    if existing is not None:
+        return existing
+
+    async def _create() -> Collection:
+        base_slug = slugify(name)
+        slug = base_slug
+        suffix = 1
+        while (
+            await session.execute(select(Collection.id).where(Collection.slug == slug))
+        ).scalar_one_or_none() is not None:
+            suffix += 1
+            slug = f"{base_slug}-{suffix}"
+        collection = Collection(name=name, overview=overview, slug=slug)
+        session.add(collection)
+        await session.flush()
+        session.add(
+            ExternalID(
+                entity_type="collection",
+                entity_id=collection.id,
+                provider_id=provider_id,
+                external_id=external_id,
+            )
        )
-    )
-    await session.flush()
-    return collection
+        await session.flush()
+        return collection
+
+    return await _create_or_recover(session, _create, _lookup)


 async def find_person_by_external_id(
@@ -185,26 +259,33 @@ async def get_or_create_person(
    name: str,
    **extra_fields: object,
 ) -> Person:
-    person = await find_person_by_external_id(session, provider_id, external_id)
-    if person is not None:
+    async def _lookup() -> Person | None:
+        return await find_person_by_external_id(session, provider_id, external_id)
+
+    existing = await _lookup()
+    if existing is not None:
        for key, value in extra_fields.items():
            if value is not None:
-                setattr(person, key, value)
+                setattr(existing, key, value)
+        await session.flush()
+        return existing
+
+    async def _create() -> Person:
+        person = Person(name=name, **extra_fields)
+        session.add(person)
+        await session.flush()
+        session.add(
+            ExternalID(
+                entity_type="person",
+                entity_id=person.id,
+                provider_id=provider_id,
+                external_id=external_id,
+            )
+        )
        await session.flush()
        return person
-    person = Person(name=name, **extra_fields)
-    session.add(person)
-    await session.flush()
-    session.add(
-        ExternalID(
-            entity_type="person",
-            entity_id=person.id,
-            provider_id=provider_id,
-            external_id=external_id,
-        )
-    )
-    await session.flush()
-    return person
+
+    return await _create_or_recover(session, _create, _lookup)


 async def find_entity_by_external_id(
@@ -247,24 +328,32 @@ async def link_external_id(
    external_id: str,
    url: str | None = None,
 ) -> None:
-    stmt = select(ExternalID).where(
-        ExternalID.entity_type == entity_type,
-        ExternalID.entity_id == entity_id,
-        ExternalID.provider_id == provider_id,
-    )
-    result = await session.execute(stmt)
-    existing = result.scalar_one_or_none()
+    async def _lookup() -> ExternalID | None:
+        stmt = select(ExternalID).where(
+            ExternalID.entity_type == entity_type,
+            ExternalID.entity_id == entity_id,
+            ExternalID.provider_id == provider_id,
+        )
+        result = await session.execute(stmt)
+        return result.scalar_one_or_none()
+
+    existing = await _lookup()
    if existing is not None:
        existing.external_id = external_id
        existing.url = url
-    else:
-        session.add(
-            ExternalID(
-                entity_type=entity_type,
-                entity_id=entity_id,
-                provider_id=provider_id,
-                external_id=external_id,
-                url=url,
-            )
+        await session.flush()
+        return
+
+    async def _create() -> ExternalID:
+        record = ExternalID(
+            entity_type=entity_type,
+            entity_id=entity_id,
+            provider_id=provider_id,
+            external_id=external_id,
+            url=url,
        )
-    await session.flush()
+        session.add(record)
+        await session.flush()
+        return record
+
+    await _create_or_recover(session, _create, _lookup)