common/db.py

"""SQLite metadata store with FTS5 keyword search for ACL Anthology papers."""
import sqlite3

SCHEMA = """
CREATE TABLE IF NOT EXISTS papers (
    id TEXT PRIMARY KEY,
    title TEXT NOT NULL,
    abstract TEXT NOT NULL DEFAULT '',
    authors TEXT NOT NULL DEFAULT '',
    venue TEXT NOT NULL DEFAULT '',
    year INTEGER,
    url TEXT NOT NULL DEFAULT '',
    bibtex TEXT NOT NULL DEFAULT '',
    pdf_url TEXT NOT NULL DEFAULT '',
    active INTEGER NOT NULL DEFAULT 1,
    faiss_id INTEGER
);

CREATE VIRTUAL TABLE IF NOT EXISTS papers_fts USING fts5(
    id UNINDEXED,
    title,
    abstract,
    tokenize='porter'
);
"""


# WAL + synchronous=NORMAL lets writers commit without an fsync per commit
# (fsync only happens at checkpoint). On slow-fsync filesystems like WSL's
# /mnt/c, this is the difference between minutes and hours for a full sync.
_PRAGMAS = [
    "PRAGMA journal_mode=WAL",
    "PRAGMA synchronous=NORMAL",
    "PRAGMA temp_store=MEMORY",
]


def init_db(path: str) -> sqlite3.Connection:
    conn = sqlite3.connect(path, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    conn.executescript(SCHEMA)
    _migrate(conn)
    for pragma in _PRAGMAS:
        conn.execute(pragma)
    conn.commit()
    return conn


def _migrate(conn: sqlite3.Connection) -> None:
    """Add columns introduced after the first shipped schema to an existing DB.

    `CREATE TABLE IF NOT EXISTS` won't add columns to a table that already
    exists, so an older downloaded snapshot (without `bibtex`) would be missing
    the column until this runs. Keeping migrations here means the service works
    against the current HF snapshot without a fresh full resync.
    """
    columns = {row["name"] for row in conn.execute("PRAGMA table_info(papers)")}
    if "bibtex" not in columns:
        conn.execute("ALTER TABLE papers ADD COLUMN bibtex TEXT NOT NULL DEFAULT ''")
    if "pdf_url" not in columns:
        conn.execute("ALTER TABLE papers ADD COLUMN pdf_url TEXT NOT NULL DEFAULT ''")


def upsert_papers(conn: sqlite3.Connection, papers: list[dict]) -> None:
    """Upsert many papers in a single transaction.

    Batched so we pay the commit cost once per batch instead of once per
    paper; the per-paper commit (and its fsync) was the dominant cost of a
    full sync. `upsert_paper` delegates here with a one-element list.
    """
    if not papers:
        return
    rows = []
    fts_rows = []
    for paper in papers:
        params = {**paper, "active": int(paper["active"])}
        params.setdefault("faiss_id", None)
        params.setdefault("bibtex", "")
        params.setdefault("pdf_url", "")
        rows.append(params)
        fts_rows.append((paper["id"], paper["title"], paper["abstract"]))
    with conn:  # one transaction: commit on success, rollback on error
        conn.executemany(
            """
            INSERT INTO papers (id, title, abstract, authors, venue, year, url, bibtex, pdf_url, active, faiss_id)
            VALUES (:id, :title, :abstract, :authors, :venue, :year, :url, :bibtex, :pdf_url, :active, :faiss_id)
            ON CONFLICT(id) DO UPDATE SET
                title=excluded.title, abstract=excluded.abstract, authors=excluded.authors,
                venue=excluded.venue, year=excluded.year, url=excluded.url, bibtex=excluded.bibtex,
                pdf_url=excluded.pdf_url,
                active=excluded.active, faiss_id=COALESCE(excluded.faiss_id, papers.faiss_id)
            """,
            rows,
        )
        conn.executemany("DELETE FROM papers_fts WHERE id = ?", [(r["id"],) for r in rows])
        conn.executemany(
            "INSERT INTO papers_fts (id, title, abstract) VALUES (?, ?, ?)",
            fts_rows,
        )


def upsert_paper(conn: sqlite3.Connection, paper: dict) -> None:
    upsert_papers(conn, [paper])


def update_metadata_fields(conn: sqlite3.Connection, papers: list[dict]) -> None:
    """Keep the bibtex/pdf_url columns in sync with the anthology source,
    without re-embedding.

    Embedding is driven by `compute_content_hash` (title, abstract, authors,
    venue, year, url) in `sync.delta`; bibtex and pdf_url are deliberately
    excluded from that hash because re-embedding appends a new FAISS vector
    and orphans the old one (`add_vector` never overwrites). The embedding
    delta path handles new/changed papers; this metadata-only write keeps
    these fields current for *unchanged* papers at no embedding cost.

    Only rows where at least one field differs are written, so in steady
    state this is a no-op. Papers without either a `bibtex` or `pdf_url` key
    (e.g. test mocks) are skipped entirely. For a paper missing just one of
    the two keys, that field is passed as `None` and `COALESCE`d against its
    existing column value, so it's left untouched instead of being clobbered
    with "".
    """
    rows = [
        (p.get("bibtex"), p.get("pdf_url"), p["id"])
        for p in papers if "bibtex" in p or "pdf_url" in p
    ]
    if not rows:
        return
    with conn:
        conn.executemany(
            """
            UPDATE papers SET
                bibtex = COALESCE(?, bibtex),
                pdf_url = COALESCE(?, pdf_url)
            WHERE id = ? AND (bibtex IS NOT COALESCE(?, bibtex) OR pdf_url IS NOT COALESCE(?, pdf_url))
            """,
            [(b, p, i, b, p) for b, p, i in rows],
        )


def mark_inactive_ids(conn: sqlite3.Connection, paper_ids: list[str]) -> None:
    """Mark many papers inactive in a single transaction (batched for the
    same reason as `upsert_papers`)."""
    if not paper_ids:
        return
    ids = [(pid,) for pid in paper_ids]
    with conn:
        conn.executemany("UPDATE papers SET active = 0 WHERE id = ?", ids)
        conn.executemany("DELETE FROM papers_fts WHERE id = ?", ids)


def mark_inactive(conn: sqlite3.Connection, paper_id: str) -> None:
    mark_inactive_ids(conn, [paper_id])


def get_paper_id_by_faiss_id(conn: sqlite3.Connection, faiss_id: int) -> str | None:
    row = conn.execute(
        "SELECT id FROM papers WHERE faiss_id = ? AND active = 1", (faiss_id,)
    ).fetchone()
    return row["id"] if row else None


def get_paper(conn: sqlite3.Connection, paper_id: str) -> dict | None:
    row = conn.execute("SELECT * FROM papers WHERE id = ? AND active = 1", (paper_id,)).fetchone()
    return dict(row) if row else None


def keyword_search(conn: sqlite3.Connection, query: str, limit: int = 10) -> list[dict]:
    rows = conn.execute(
        """
        SELECT p.* FROM papers_fts
        JOIN papers p ON p.id = papers_fts.id
        WHERE papers_fts MATCH ? AND p.active = 1
        ORDER BY rank
        LIMIT ?
        """,
        (query, limit),
    ).fetchall()
    return [dict(r) for r in rows]


def get_papers_by_ids(conn: sqlite3.Connection, ids: list[str]) -> list[dict]:
    if not ids:
        return []
    placeholders = ",".join("?" * len(ids))
    rows = conn.execute(
        f"SELECT * FROM papers WHERE id IN ({placeholders}) AND active = 1", ids
    ).fetchall()
    by_id = {r["id"]: dict(r) for r in rows}
    return [by_id[i] for i in ids if i in by_id]