fix(ingest): retry on bibkey UNIQUE collision (audit C-15)

The id-keyed upsert (ON CONFLICT (id)) does not catch the papers.bibkey UNIQUE
constraint, so a second paper whose auto-generated key matches an existing one
(two same-author/year works → e.g. 'BobenkoLutz2023') failed the whole ingest
with a UniqueViolation. Surfaced by the canonical-id re-ingest: 2305.10988
collided with 2310.17529.

ingest now catches a bibkey UniqueViolation, rolls back, and retries the upsert
with an a/b/c… suffix. Non-bibkey violations re-raise unchanged. Mocks don't
raise, so the existing happy-path tests are untouched; a new test drives the
collision-and-retry path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Tarik Moussa
2026-06-15 22:11:14 +02:00
parent 45c3b16dd7
commit b52303d22b
2 changed files with 77 additions and 28 deletions

View File

@@ -8,6 +8,7 @@ from dataclasses import dataclass
from pathlib import Path from pathlib import Path
import numpy as np import numpy as np
import psycopg
from codex.config import get_settings from codex.config import get_settings
from codex.db import get_conn from codex.db import get_conn
@@ -125,9 +126,7 @@ def ingest_paper(
# --------------------------------------------------------------- # ---------------------------------------------------------------
# 3. Upsert paper (INSERT … ON CONFLICT (id) DO UPDATE SET …) # 3. Upsert paper (INSERT … ON CONFLICT (id) DO UPDATE SET …)
# --------------------------------------------------------------- # ---------------------------------------------------------------
with get_conn() as conn: upsert_sql = """
conn.execute(
"""
INSERT INTO papers (id, openalex_id, bibkey, title, authors, year, abstract, INSERT INTO papers (id, openalex_id, bibkey, title, authors, year, abstract,
source_path, abstract_emb) source_path, abstract_emb)
VALUES (%(id)s, %(openalex_id)s, %(bibkey)s, %(title)s, %(authors)s, %(year)s, VALUES (%(id)s, %(openalex_id)s, %(bibkey)s, %(title)s, %(authors)s, %(year)s,
@@ -141,7 +140,16 @@ def ingest_paper(
abstract = EXCLUDED.abstract, abstract = EXCLUDED.abstract,
source_path = COALESCE(EXCLUDED.source_path, papers.source_path), source_path = COALESCE(EXCLUDED.source_path, papers.source_path),
abstract_emb = EXCLUDED.abstract_emb abstract_emb = EXCLUDED.abstract_emb
""", """
with get_conn() as conn:
# The id-keyed upsert does not catch a papers.bibkey UNIQUE collision (two
# same-author/year papers generate the same key). On such a collision retry
# with an a/b/c… suffix instead of failing the whole ingest (audit C-15).
base_bibkey = paper.bibkey
for attempt in range(27):
try:
conn.execute(
upsert_sql,
{ {
"id": paper.id, "id": paper.id,
"openalex_id": paper.openalex_id, "openalex_id": paper.openalex_id,
@@ -154,6 +162,12 @@ def ingest_paper(
"abstract_emb": abstract_emb, "abstract_emb": abstract_emb,
}, },
) )
break
except psycopg.errors.UniqueViolation as exc:
if base_bibkey is None or "bibkey" not in str(exc).lower() or attempt == 26:
raise
conn.rollback()
paper.bibkey = f"{base_bibkey}{chr(ord('a') + attempt)}"
# Register openalex_id in paper_identifiers (alias table). # Register openalex_id in paper_identifiers (alias table).
# Every ingest records the primary openalex_id so the graph JOIN # Every ingest records the primary openalex_id so the graph JOIN

View File

@@ -11,6 +11,7 @@ from typing import Any
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock, patch
import numpy as np import numpy as np
import psycopg
import pytest import pytest
from codex.ingest import IngestResult, _make_bibkey, ingest_paper from codex.ingest import IngestResult, _make_bibkey, ingest_paper
@@ -525,6 +526,40 @@ def test_ingest_sets_bibkey_when_none(tmp_path: Any) -> None:
assert upsert_params["bibkey"] == "AliceBob2023" assert upsert_params["bibkey"] == "AliceBob2023"
def test_ingest_disambiguates_bibkey_on_unique_violation() -> None:
"""A papers.bibkey UNIQUE collision is retried with an a/b/c suffix (audit C-15)."""
paper = _make_paper() # bibkey auto-generates to "AliceBob2023"
mock_conn = MagicMock()
state = {"first": True}
def execute_side_effect(sql: str, params: Any = None) -> MagicMock:
# The first papers INSERT collides on the bibkey; the retry must succeed.
if "INSERT INTO papers" in sql and state["first"]:
state["first"] = False
raise psycopg.errors.UniqueViolation(
'duplicate key value violates unique constraint "papers_bibkey_key"'
)
return MagicMock()
mock_conn.execute.side_effect = execute_side_effect
with (
patch("codex.ingest.openalex.fetch_paper", return_value=paper),
patch("codex.ingest.openalex.fetch_citations", return_value=[]),
patch("codex.ingest.get_embedder", return_value=_fake_embedder()),
patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)),
):
ingest_paper(paper.id)
papers_inserts = [
c for c in mock_conn.execute.call_args_list if "INSERT INTO papers" in c[0][0]
]
assert len(papers_inserts) == 2, "expected one retry after the bibkey collision"
assert papers_inserts[0][0][1]["bibkey"] == "AliceBob2023"
assert papers_inserts[1][0][1]["bibkey"] == "AliceBob2023a"
mock_conn.rollback.assert_called_once()
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------