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:
@@ -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,35 +126,48 @@ def ingest_paper(
|
|||||||
# ---------------------------------------------------------------
|
# ---------------------------------------------------------------
|
||||||
# 3. Upsert paper (INSERT … ON CONFLICT (id) DO UPDATE SET …)
|
# 3. Upsert paper (INSERT … ON CONFLICT (id) DO UPDATE SET …)
|
||||||
# ---------------------------------------------------------------
|
# ---------------------------------------------------------------
|
||||||
|
upsert_sql = """
|
||||||
|
INSERT INTO papers (id, openalex_id, bibkey, title, authors, year, abstract,
|
||||||
|
source_path, abstract_emb)
|
||||||
|
VALUES (%(id)s, %(openalex_id)s, %(bibkey)s, %(title)s, %(authors)s, %(year)s,
|
||||||
|
%(abstract)s, %(source_path)s, %(abstract_emb)s)
|
||||||
|
ON CONFLICT (id) DO UPDATE SET
|
||||||
|
openalex_id = EXCLUDED.openalex_id,
|
||||||
|
bibkey = COALESCE(papers.bibkey, EXCLUDED.bibkey),
|
||||||
|
title = EXCLUDED.title,
|
||||||
|
authors = EXCLUDED.authors,
|
||||||
|
year = EXCLUDED.year,
|
||||||
|
abstract = EXCLUDED.abstract,
|
||||||
|
source_path = COALESCE(EXCLUDED.source_path, papers.source_path),
|
||||||
|
abstract_emb = EXCLUDED.abstract_emb
|
||||||
|
"""
|
||||||
with get_conn() as conn:
|
with get_conn() as conn:
|
||||||
conn.execute(
|
# 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
|
||||||
INSERT INTO papers (id, openalex_id, bibkey, title, authors, year, abstract,
|
# with an a/b/c… suffix instead of failing the whole ingest (audit C-15).
|
||||||
source_path, abstract_emb)
|
base_bibkey = paper.bibkey
|
||||||
VALUES (%(id)s, %(openalex_id)s, %(bibkey)s, %(title)s, %(authors)s, %(year)s,
|
for attempt in range(27):
|
||||||
%(abstract)s, %(source_path)s, %(abstract_emb)s)
|
try:
|
||||||
ON CONFLICT (id) DO UPDATE SET
|
conn.execute(
|
||||||
openalex_id = EXCLUDED.openalex_id,
|
upsert_sql,
|
||||||
bibkey = COALESCE(papers.bibkey, EXCLUDED.bibkey),
|
{
|
||||||
title = EXCLUDED.title,
|
"id": paper.id,
|
||||||
authors = EXCLUDED.authors,
|
"openalex_id": paper.openalex_id,
|
||||||
year = EXCLUDED.year,
|
"bibkey": paper.bibkey,
|
||||||
abstract = EXCLUDED.abstract,
|
"title": paper.title,
|
||||||
source_path = COALESCE(EXCLUDED.source_path, papers.source_path),
|
"authors": paper.authors,
|
||||||
abstract_emb = EXCLUDED.abstract_emb
|
"year": paper.year,
|
||||||
""",
|
"abstract": paper.abstract,
|
||||||
{
|
"source_path": source_path,
|
||||||
"id": paper.id,
|
"abstract_emb": abstract_emb,
|
||||||
"openalex_id": paper.openalex_id,
|
},
|
||||||
"bibkey": paper.bibkey,
|
)
|
||||||
"title": paper.title,
|
break
|
||||||
"authors": paper.authors,
|
except psycopg.errors.UniqueViolation as exc:
|
||||||
"year": paper.year,
|
if base_bibkey is None or "bibkey" not in str(exc).lower() or attempt == 26:
|
||||||
"abstract": paper.abstract,
|
raise
|
||||||
"source_path": source_path,
|
conn.rollback()
|
||||||
"abstract_emb": abstract_emb,
|
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
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user