"""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]