From 6ea150690d0eb468cedadf5de4807a90e24a8d53 Mon Sep 17 00:00:00 2001 From: Codex Bot Date: Sat, 7 Nov 2026 15:38:42 +0800 Subject: [PATCH] Support native PostgreSQL metadata ingestion and provenance writes --- .gitea/workflows/check.yml | 4 + README.md | 2 +- docs/operations.md | 10 ++ docs/usage.md | 19 ++- paper_atlas/db/backend.py | 23 ++++ paper_atlas/db/store.py | 29 ++-- paper_atlas/db/submissions.py | 16 ++- paper_atlas/ingestion/arxiv.py | 60 ++++---- paper_atlas/ingestion/backfill.py | 14 +- paper_atlas/ingestion/records.py | 26 ++-- paper_atlas/postgres/writer.py | 87 ++++++++++++ tests/test_postgres_ingestion.py | 220 ++++++++++++++++++++++++++++++ 12 files changed, 438 insertions(+), 72 deletions(-) create mode 100644 paper_atlas/postgres/writer.py create mode 100644 tests/test_postgres_ingestion.py diff --git a/.gitea/workflows/check.yml b/.gitea/workflows/check.yml index 0eeabc1..8b64c2c 100644 --- a/.gitea/workflows/check.yml +++ b/.gitea/workflows/check.yml @@ -49,6 +49,10 @@ jobs: subprocess.run([str(venv/'bin/paper-atlas'),'--help'],cwd=outside,check=True) subprocess.run([python,'-I','-B','-m','paper_atlas','collect','cvf','--help'],cwd=outside,check=True) subprocess.run([python,'-I','-B','-m','paper_atlas','maintain','postgres','--help'],cwd=outside,check=True) + subprocess.run([python,'-I','-B','-m','paper_atlas','import','records','--help'],cwd=outside,check=True) + subprocess.run([python,'-I','-B','-m','paper_atlas','import','arxiv','--help'],cwd=outside,check=True) + subprocess.run([python,'-I','-B','-m','paper_atlas','maintain','backfill','--help'],cwd=outside,check=True) + subprocess.run([python,'-I','-B','-c','from paper_atlas.postgres.writer import PostgresStore, WriteConnection'],cwd=outside,check=True) smoke = "from paper_atlas.db.store import Store; from paper_atlas.config import DEFAULT_REGISTRY; import json; assert {'COLM','MM','MLSys'} <= {v['acronym'] for v in json.loads(DEFAULT_REGISTRY.read_text())['venues']}; s=Store('installed.sqlite'); s.conn.close()" subprocess.run([python,'-I','-B','-c',smoke],cwd=outside,check=True) subprocess.run([python,'-I','-B','-m','paper_atlas','query','--db','installed.sqlite','--stats'],cwd=outside,check=True) diff --git a/README.md b/README.md index 43ac6bc..09dcf66 100644 --- a/README.md +++ b/README.md @@ -61,7 +61,7 @@ paper-atlas collect colm --cache /data/paper-index/raw/http-cache \ `paper-atlas collect --help` 查看来源；`paper-atlas import --help` 查看导入方式。完整批次步骤见[使用指南](docs/usage.md)。 -共享实例名为 `foundry-postgres`，数据位于 `foundry_data` 库的 `research` schema；其他项目可按权限复用同一组表。库名通过 libpq service 配置、schema 通过 search_path 或 `--pg-schema` 指定，不写死在查询中。PostgreSQL 凭据放在服务器权限为 `0600` 的 libpq service 文件中；命令行只使用服务名称。`paper-atlas maintain postgres --help` 提供迁移入口：先完整核验 21 张业务表，再发布目标 schema。SQLite FTS5 内部表由 PostgreSQL GIN 索引替代。切换后 SQLite 为备份，不能继续用旧导入器写入它；后续生产导入须适配 PostgreSQL，不能把 SQLite 更新当作生产更新。 +共享实例名为 `foundry-postgres`，数据位于 `foundry_data` 库的 `research` schema；其他项目可按权限复用同一组表。库名通过 libpq service 配置、schema 通过 search_path 或 `--pg-schema` 指定，不写死在查询中。PostgreSQL 凭据放在服务器权限为 `0600` 的 libpq service 文件中；命令行只使用服务名称。`paper-atlas maintain postgres --help` 提供迁移入口：先完整核验 21 张业务表，再发布目标 schema。SQLite FTS5 内部表由 PostgreSQL GIN 索引替代。切换后 SQLite 为只读备份；`import records`、`import arxiv` 和 `maintain backfill` 使用 `--pg-service foundry-research` 写入主库，整个写入会话只允许一个协调写入者，会议文件失败时整体回滚并保留失败批次记录。 ## 数据如何组织 diff --git a/docs/operations.md b/docs/operations.md index 766c8ba..5299390 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -30,6 +30,16 @@ SQLite 允许文本包含 `U+0000`（NUL），PostgreSQL `text` 不允许。迁 远程集成测试通过 `PAPER_ATLAS_TEST_PG_SERVICE` 显式指定，使用随机前缀 schema 并仅清理自己的 fixture；没有该配置时会跳过，应另外报告“未运行”，不能用 SQLite 测试代替 PostgreSQL 验证。 +## PostgreSQL 增量写入 + +`import records`、`import arxiv`、`maintain backfill` 使用 `--pg-service` 选择主库，`--pg-schema` 可覆盖配置的 search_path。写入只接受已经发布、拥有完整业务表和受支持编码标记的 schema，不自动建库或引导空表。SQLite 样本仍使用 `--db`，已切换备份继续禁止写入。 + +写入连接按数据库/schema 获取会话级 advisory lock，文件提交、回滚以及 arXiv 分批提交均不释放；关闭连接才释放，其他写入者立即得到明确的占用错误。只读查询继续使用独立一致性快照。锁是所有调用方必须遵守的协议，其他项目不要绕过接口并发修改同一论文和投稿；[PostgreSQL 会话锁说明](https://www.postgresql.org/docs/17/explicit-locking.html#ADVISORY-LOCKS)解释了其事务边界。 + +论文身份匹配、跨来源合并、投稿/决定来源和标签维护复用同一套领域逻辑。PostgreSQL 使用原生 RETURNING 和 identity 序列分配 ID，继承迁移时设定的历史 ID 高水位；失败和冲突尝试可以留下序列空洞，不能把最大 ID 当作行数。绑定文本按 `nul-soh-v1` 编码，结果解码，GIN 随论文写入自动更新。 + +会议文件的数据变化在同一事务内完成，manifest 数量不符等错误会回滚整份文件，`ingest_runs` 另保留 failed、已读数和零写入数；身份回填整个队列一次提交。arXiv 每 500 行提交数据和 checkpoint，异常只回滚尚未提交部分，重试必须仍为同一压缩包。大批次开始前核对独立审计清单、当前主库冲突和磁盘空间；完成后对数量、实际来源/决定关联及检索做数据验收，代码测试成功不代表生产导入成功。 + 代码以固定 commit 部署，已经启动的任务继续使用其原版本。目录重构不要求迁移数据库，也不更改已有 JSONL 和 manifest 格式。 ## arXiv 原始编号与类别证据 diff --git a/docs/usage.md b/docs/usage.md index c8d9895..18ea896 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -18,7 +18,7 @@ paper-atlas query --pg-service foundry-research --paper 123 CLI 自动还原迁移中的 NUL/SOH 文本转义，详情中的来源摘要和筛选参数保留原字符。直接 SQL 读取可能看到 `U+0001` 转义序列；跨项目 ETL 应按迁移元数据中的 `text_codec` 使用 `paper_atlas.postgres.text.encode/decode`。存储格式和无损哈希验收见[运行维护](operations.md#postgresql-迁移)。 -迁移与切换见[运行维护](operations.md#postgresql-迁移)。下述 SQLite 导入命令仅用于未切换的索引或独立样本；切换后的生产源拒绝旧写入。尚未适配的覆盖审计可只读检查冻结备份，不能当作 PostgreSQL 当前状态。 +迁移与切换见[运行维护](operations.md#postgresql-迁移)。`import records`、`import arxiv` 和 `maintain backfill` 均支持 `--pg-service` 和可选 `--pg-schema`；与 `--db` 二选一。SQLite 导入仅用于未切换的索引或独立样本。尚未适配 PostgreSQL 的维护、覆盖审计命令仍可只读检查冻结备份，不能当作主库的当前状态。 ## 1. 下载固定快照 @@ -43,10 +43,23 @@ paper-atlas audit public-openreview --help ## 3. 串行导入 -首次 arXiv 导入使用 `paper-atlas import arxiv --db ... --archive ...`。现有数据库合并已核验的会议文件使用： +PostgreSQL 主库逐文件导入已经核验的资料： ```sh -paper-atlas maintain merge --db /data/paper-index/papers.sqlite \ +paper-atlas import records --pg-service foundry-research \ + /data/paper-index/collected/selected-file.jsonl +paper-atlas maintain backfill --pg-service foundry-research \ + --queue /data/paper-index/collected/audited-identity-queue.json +``` + +文件和 queue 须替换为实际核验通过的路径。会议导入按文件提交，失败时该文件的论文、标签和决定更改全部回滚，失败批次日志保留；重跑会按来源标识与决定证据去重。回填使用已有源记录，只合并经独立身份佐证的投稿，不重新导入文本。所有写入者须遵守同一数据库/schema 的会话锁；其他项目也应通过相同的写入接口更新共享论文和投稿。 + +未来需要导入新的 arXiv 元数据包时使用 `paper-atlas import arxiv --pg-service foundry-research --archive ... --storage-path ...`。每 500 行提交断点，恢复同一文件时保留已提交 ID 和标签证据；`--storage-path` 指定运行主机上要监测的文件系统，在数据库主机执行时应指向数据库卷。已完成的历史压缩包不需要重复导入。 + +SQLite 独立样本的首次 arXiv 导入使用 `paper-atlas import arxiv --db ... --archive ...`。SQLite 协调器仍适用于未切换的索引： + +```sh +paper-atlas maintain merge --db /data/paper-index/samples/example.sqlite \ --archive /data/paper-index/raw/arxiv-metadata.zip \ /data/paper-index/collected/selected-file.jsonl ``` diff --git a/paper_atlas/db/backend.py b/paper_atlas/db/backend.py index 3faf2ad..09bb6b3 100644 --- a/paper_atlas/db/backend.py +++ b/paper_atlas/db/backend.py @@ -2,6 +2,29 @@ from pathlib import Path +def database_arguments(parser): + target = parser.add_mutually_exclusive_group(required=True) + target.add_argument('--db', help='Writable SQLite sample; never the frozen production backup') + target.add_argument('--pg-service', help='Protected libpq service for the published PostgreSQL dataset') + parser.add_argument('--pg-schema', help='Optional schema; otherwise use the configured search_path') + + +def open_store(db=None, *, pg_service=None, pg_schema=None): + if bool(db) == bool(pg_service) or (pg_schema and not pg_service): + raise ValueError('Choose exactly one of SQLite db or PostgreSQL service; schema requires PostgreSQL') + if pg_service: + from paper_atlas.postgres.writer import PostgresStore + return PostgresStore(pg_service, pg_schema) + from paper_atlas.db.store import Store + return Store(db) + + +def inserted_id(connection, sql, params=()): + if getattr(connection, 'postgres', False): + return connection.execute(sql + ' RETURNING id', params).fetchone()[0] + return connection.execute(sql, params).lastrowid + + def require_sqlite_writable(path): marker = Path(str(Path(path).resolve()) + '.postgresql-cutover.json') if marker.exists(): diff --git a/paper_atlas/db/store.py b/paper_atlas/db/store.py index 9495ecc..155b1bf 100644 --- a/paper_atlas/db/store.py +++ b/paper_atlas/db/store.py @@ -6,7 +6,7 @@ import sqlite3 import unicodedata from pathlib import Path from paper_atlas.db.submissions import attach_submission,migrate_sources,write_decision -from paper_atlas.db.backend import require_sqlite_writable +from paper_atlas.db.backend import require_sqlite_writable, inserted_id def normalized(value): @@ -38,7 +38,7 @@ class Store: migrate_sources(self.conn) def source(self,name,url): - self.conn.execute('INSERT OR IGNORE INTO sources(name,homepage) VALUES(?,?)',(name,url)) + self.conn.execute('INSERT INTO sources(name,homepage) VALUES(?,?) ON CONFLICT DO NOTHING',(name,url)) return self.conn.execute('SELECT id FROM sources WHERE name=?',(name,)).fetchone()[0] def merge(self,target,other,reason): @@ -55,7 +55,7 @@ class Store: c.execute('UPDATE identifiers SET paper_id=? WHERE paper_id=?',(target,other)) c.execute('UPDATE source_records SET paper_id=? WHERE paper_id=?',(target,other)) c.execute('UPDATE submissions SET paper_id=? WHERE paper_id=?',(target,other)) - c.execute('INSERT OR IGNORE INTO paper_tags SELECT ?,tag_id,source_record_id,evidence_key FROM paper_tags WHERE paper_id=?',(target,other)) + c.execute('INSERT INTO paper_tags SELECT ?,tag_id,source_record_id,evidence_key FROM paper_tags WHERE paper_id=? ON CONFLICT DO NOTHING',(target,other)) c.execute('DELETE FROM paper_tags WHERE paper_id=?',(other,)) c.execute('UPDATE paper_merges SET surviving_paper_id=? WHERE surviving_paper_id=?',(target,other)) # Reattach pending candidate evidence without losing colliding rows. @@ -65,7 +65,7 @@ class Store: left = target if left == other else left right = target if right == other else right if left != right: - c.execute('INSERT OR IGNORE INTO match_candidates VALUES(?,?,?,?,?,?)',(min(left,right),max(left,right),method,score,evidence,status)) + c.execute('INSERT INTO match_candidates VALUES(?,?,?,?,?,?) ON CONFLICT DO NOTHING',(min(left,right),max(left,right),method,score,evidence,status)) c.execute('DELETE FROM papers WHERE id=?',(other,)) def record(self,source,external_id,url,title,abstract,authors,identifiers=None,metadata=None,published=None): @@ -111,18 +111,19 @@ class Store: c.execute("UPDATE papers SET title=?,abstract=COALESCE(NULLIF(?,''),abstract),title_key=?,authors_json=CASE WHEN ?='[]' THEN authors_json ELSE ? END WHERE id=? AND title=? AND COALESCE(abstract,'')=COALESCE(?,'')",(title,abstract,key,json.dumps(authors,ensure_ascii=False),json.dumps(authors,ensure_ascii=False),paper,old[1],old_abstract)) c.execute("UPDATE papers SET abstract=CASE WHEN abstract IS NULL OR abstract='' THEN ? ELSE abstract END WHERE id=?",(abstract,paper)) else: - paper = c.execute('INSERT INTO papers(title,abstract,title_key,authors_json,first_published) VALUES(?,?,?,?,?)',(title,abstract,key,json.dumps(authors,ensure_ascii=False),published)).lastrowid + paper = inserted_id(c,'INSERT INTO papers(title,abstract,title_key,authors_json,first_published) VALUES(?,?,?,?,?)',(title,abstract,key,json.dumps(authors,ensure_ascii=False),published)) for scheme,value in ids: - c.execute('INSERT OR IGNORE INTO identifiers VALUES(?,?,?)',(scheme,value,paper)) + c.execute('INSERT INTO identifiers VALUES(?,?,?) ON CONFLICT DO NOTHING',(scheme,value,paper)) metadata=dict(metadata or {});metadata['record_identifiers']=dict(ids) if old and not abstract: previous={'abstract':old[2],'observed_at':old[3]} if old[2] else json.loads(old[4]).get('last_known_abstract') if previous:metadata['last_known_abstract']=previous - c.execute('INSERT INTO source_records(source_id,external_id,paper_id,url,title,abstract,authors_json,metadata_json) VALUES(?,?,?,?,?,?,?,?) ON CONFLICT(source_id,external_id) DO UPDATE SET paper_id=excluded.paper_id,url=excluded.url,title=excluded.title,abstract=excluded.abstract,authors_json=excluded.authors_json,metadata_json=excluded.metadata_json,retrieved_at=CURRENT_TIMESTAMP',(source_id,external_id,paper,url,title,abstract,json.dumps(authors,ensure_ascii=False),json.dumps(metadata,ensure_ascii=False,default=str))) + from paper_atlas.db.submissions import timestamp + c.execute('INSERT INTO source_records(source_id,external_id,paper_id,url,title,abstract,authors_json,metadata_json,retrieved_at) VALUES(?,?,?,?,?,?,?,?,?) ON CONFLICT(source_id,external_id) DO UPDATE SET paper_id=excluded.paper_id,url=excluded.url,title=excluded.title,abstract=excluded.abstract,authors_json=excluded.authors_json,metadata_json=excluded.metadata_json,retrieved_at=excluded.retrieved_at',(source_id,external_id,paper,url,title,abstract,json.dumps(authors,ensure_ascii=False),json.dumps(metadata,ensure_ascii=False,default=str),timestamp())) record = c.execute('SELECT id FROM source_records WHERE source_id=? AND external_id=?',(source_id,external_id)).fetchone()[0] for other,_,_ in title_matches: if other != paper and c.execute('SELECT 1 FROM papers WHERE id=?',(other,)).fetchone(): - c.execute('INSERT OR IGNORE INTO match_candidates VALUES(?,?,?,?,?,?)',(min(paper,other),max(paper,other),'exact_title_unconfirmed',None,json.dumps({'title':title}), 'pending')) + c.execute('INSERT INTO match_candidates VALUES(?,?,?,?,?,?) ON CONFLICT DO NOTHING',(min(paper,other),max(paper,other),'exact_title_unconfirmed',None,json.dumps({'title':title}), 'pending')) return paper,record def record_tags(self,paper,record,source,tags): @@ -137,18 +138,18 @@ class Store: c=self.conn c.execute('DELETE FROM paper_tags WHERE source_record_id=? AND evidence_key=?',(record,evidence)) for namespace,value in sorted(values): - c.execute('INSERT OR IGNORE INTO tags(namespace,value) VALUES(?,?)',(namespace,value)) + c.execute('INSERT INTO tags(namespace,value) VALUES(?,?) ON CONFLICT DO NOTHING',(namespace,value)) tag=c.execute('SELECT id FROM tags WHERE namespace=? AND value=?',(namespace,value)).fetchone()[0] - c.execute('INSERT OR IGNORE INTO paper_tags VALUES(?,?,?,?)',(paper,tag,record,evidence)) + c.execute('INSERT INTO paper_tags VALUES(?,?,?,?) ON CONFLICT DO NOTHING',(paper,tag,record,evidence)) def decision(self,paper,record,venue,year,track,external_id,url,status,raw_value,aliases=None,identities=None,observed_at=None,decision_date=None,source_decision_id='',submitted_at=None): c = self.conn - existing=c.execute('SELECT id FROM venues WHERE acronym=? COLLATE NOCASE',(venue,)).fetchone() + existing=c.execute('SELECT id FROM venues WHERE lower(acronym)=lower(?)',(venue,)).fetchone() if existing:venue_id=existing[0] - else:venue_id=c.execute("INSERT INTO venues(acronym,name,field,venue_type) VALUES(?,?,?,'conference')",(venue,venue,'unclassified')).lastrowid + else:venue_id=inserted_id(c,"INSERT INTO venues(acronym,name,field,venue_type) VALUES(?,?,?,'conference')",(venue,venue,'unclassified')) for alias in aliases or []: - c.execute('INSERT OR IGNORE INTO venue_aliases VALUES(?,?)',(venue_id,alias)) - c.execute('INSERT OR IGNORE INTO editions(venue_id,year,track) VALUES(?,?,?)',(venue_id,int(year),track)) + c.execute('INSERT INTO venue_aliases VALUES(?,?) ON CONFLICT DO NOTHING',(venue_id,alias)) + c.execute('INSERT INTO editions(venue_id,year,track) VALUES(?,?,?) ON CONFLICT DO NOTHING',(venue_id,int(year),track)) edition = c.execute('SELECT id FROM editions WHERE venue_id=? AND year=? AND track=?',(venue_id,int(year),track)).fetchone()[0] # A previous event in this row may already have merged the input paper. paper=c.execute('SELECT paper_id FROM source_records WHERE id=?',(record,)).fetchone()[0] diff --git a/paper_atlas/db/submissions.py b/paper_atlas/db/submissions.py index 2e017f4..5b12ee2 100644 --- a/paper_atlas/db/submissions.py +++ b/paper_atlas/db/submissions.py @@ -58,13 +58,15 @@ def identities_for_row(row): def migrate_sources(c): if c.execute("SELECT 1 FROM schema_migrations WHERE name='submission_source_baseline_v1'").fetchone():return with c: - c.execute('INSERT OR IGNORE INTO submission_sources SELECT id,source_record_id,source_submission_id FROM submissions') - c.execute("INSERT OR IGNORE INTO decision_sources SELECT d.id,s.source_record_id,'' FROM decisions d JOIN submissions s ON s.id=d.submission_id") + c.execute('INSERT INTO submission_sources SELECT id,source_record_id,source_submission_id FROM submissions WHERE 1=1 ON CONFLICT DO NOTHING') + c.execute("INSERT INTO decision_sources SELECT d.id,s.source_record_id,'' FROM decisions d JOIN submissions s ON s.id=d.submission_id WHERE 1=1 ON CONFLICT DO NOTHING") c.execute("INSERT INTO schema_migrations(name) VALUES('submission_source_baseline_v1')") def next_id(c,table,audit,former): # Merged IDs are kept in audit; never allocate them to a later unrelated event. + if getattr(c, 'postgres', False): + return c.next_id(table) return max(c.execute('SELECT COALESCE(max(id),0) FROM '+table).fetchone()[0],c.execute('SELECT COALESCE(max('+former+'),0) FROM '+audit).fetchone()[0])+1 @@ -72,7 +74,7 @@ def merge_submissions(c,target,other): if target==other:return editions=c.execute('SELECT id,edition_id FROM submissions WHERE id IN (?,?)',(target,other)).fetchall() if len(editions)!=2 or editions[0][1]!=editions[1][1]:raise ValueError('Cannot combine submissions from different editions') - c.execute('INSERT OR IGNORE INTO submission_sources SELECT ?,source_record_id,source_submission_id FROM submission_sources WHERE submission_id=?',(target,other)) + c.execute('INSERT INTO submission_sources SELECT ?,source_record_id,source_submission_id FROM submission_sources WHERE submission_id=? ON CONFLICT DO NOTHING',(target,other)) c.execute('DELETE FROM submission_sources WHERE submission_id=?',(other,)) for decision,status,raw,url,observed,date in c.execute('SELECT id,status,raw_value,evidence_url,observed_at,decision_date FROM decisions WHERE submission_id=?',(other,)).fetchall(): duplicate=c.execute('SELECT id,decision_date FROM decisions WHERE submission_id=? AND status=? AND raw_value=? AND evidence_url=? AND observed_at=?',(target,status,raw,url,observed)).fetchone() @@ -80,7 +82,7 @@ def merge_submissions(c,target,other): surviving=duplicate[0] if date and duplicate[1] and date!=duplicate[1]:raise ValueError('Conflicting dates for the same decision observation') if date and not duplicate[1]:c.execute('UPDATE decisions SET decision_date=? WHERE id=?',(date,surviving)) - c.execute('INSERT OR IGNORE INTO decision_sources SELECT ?,source_record_id,source_decision_id FROM decision_sources WHERE decision_id=?',(surviving,decision)) + c.execute('INSERT INTO decision_sources SELECT ?,source_record_id,source_decision_id FROM decision_sources WHERE decision_id=? ON CONFLICT DO NOTHING',(surviving,decision)) c.execute('DELETE FROM decision_sources WHERE decision_id=?',(decision,)) c.execute('UPDATE decision_merges SET surviving_decision_id=? WHERE surviving_decision_id=?',(surviving,decision)) c.execute("INSERT INTO decision_merges(former_decision_id,surviving_decision_id,reason) VALUES(?,?,'same decision observation on merged submission')",(decision,surviving)) @@ -114,9 +116,9 @@ def attach_submission(store,edition,paper,record,external_id,identities): else: submission=next_id(c,'submissions','submission_merges','former_submission_id') c.execute('INSERT INTO submissions(id,paper_id,edition_id,source_record_id,source_submission_id) VALUES(?,?,?,?,?)',(submission,paper,edition,record,external_id)) - c.execute('INSERT OR IGNORE INTO submission_sources VALUES(?,?,?)',(submission,record,external_id)) + c.execute('INSERT INTO submission_sources VALUES(?,?,?) ON CONFLICT DO NOTHING',(submission,record,external_id)) for scheme,value in identities: - c.execute('INSERT OR IGNORE INTO submission_identifiers VALUES(?,?,?,?)',(edition,scheme,value,submission)) + c.execute('INSERT INTO submission_identifiers VALUES(?,?,?,?) ON CONFLICT DO NOTHING',(edition,scheme,value,submission)) return submission @@ -130,5 +132,5 @@ def write_decision(store,submission,record,status,raw,url,observed_at=None,decis else: decision=next_id(c,'decisions','decision_merges','former_decision_id') c.execute('INSERT INTO decisions(id,submission_id,status,raw_value,evidence_url,observed_at,decision_date) VALUES(?,?,?,?,?,?,?)',(decision,submission,status,raw,url,observed,date)) - c.execute('INSERT OR IGNORE INTO decision_sources VALUES(?,?,?)',(decision,record,source_decision_id or '')) + c.execute('INSERT INTO decision_sources VALUES(?,?,?) ON CONFLICT DO NOTHING',(decision,record,source_decision_id or '')) return decision diff --git a/paper_atlas/ingestion/arxiv.py b/paper_atlas/ingestion/arxiv.py index 78d88a7..0e8ae01 100644 --- a/paper_atlas/ingestion/arxiv.py +++ b/paper_atlas/ingestion/arxiv.py @@ -1,4 +1,4 @@ -"""Stream the official metadata ZIP into SQLite. No PDFs, source files or builds.""" +"""Stream the official metadata ZIP into the configured metadata database.""" import argparse import json import time @@ -7,7 +7,8 @@ import zipfile from pathlib import Path -from paper_atlas.db.store import Store +from paper_atlas.db.backend import database_arguments, open_store, inserted_id +from paper_atlas.db.submissions import timestamp def refresh_categories(store,paper,record,categories): @@ -17,28 +18,29 @@ def refresh_categories(store,paper,record,categories): store.conn.execute("DELETE FROM paper_tags WHERE source_record_id=? AND evidence_key='arxiv:categories' AND tag_id IN (SELECT id FROM tags WHERE namespace='arxiv:category')",(record,)) -def ingest(db, archive): - store = Store(db) +def ingest(db, archive, *, pg_service=None, pg_schema=None, storage_path=None): + store = open_store(db,pg_service=pg_service,pg_schema=pg_schema) conn = store.conn - conn.execute("INSERT OR IGNORE INTO sources(name,homepage) VALUES('arxiv','https://arxiv.org')") - source = conn.execute("SELECT id FROM sources WHERE name='arxiv'").fetchone()[0] - previous = conn.execute('SELECT id,checkpoint_json FROM ingest_runs WHERE source_id=? AND status IN (\'running\',\'failed\') ORDER BY id DESC LIMIT 1', (source,)).fetchone() - fingerprint = {'path':str(Path(archive).resolve()), 'bytes':Path(archive).stat().st_size, 'mtime_ns':Path(archive).stat().st_mtime_ns} - start = 0 - if previous: - checkpoint = json.loads(previous[1] or '{}') - if checkpoint.get('archive') == fingerprint: - run = previous[0] - start = checkpoint.get('line', 0) - conn.execute("UPDATE ingest_runs SET status='running',error=NULL WHERE id=?", (run,)) - else: - raise RuntimeError('Pending run belongs to a different archive; resolve it explicitly') - else: - run = conn.execute('INSERT INTO ingest_runs(source_id,checkpoint_json) VALUES(?,?)', (source,json.dumps({'archive':fingerprint,'line':0}))).lastrowid - conn.commit() - line_number = start - last = time.monotonic() + run = None try: + source = store.source('arxiv','https://arxiv.org') + previous = conn.execute('SELECT id,checkpoint_json FROM ingest_runs WHERE source_id=? AND status IN (\'running\',\'failed\') ORDER BY id DESC LIMIT 1', (source,)).fetchone() + fingerprint = {'path':str(Path(archive).resolve()), 'bytes':Path(archive).stat().st_size, 'mtime_ns':Path(archive).stat().st_mtime_ns} + start = 0 + if previous: + checkpoint = json.loads(previous[1] or '{}') + if checkpoint.get('archive') == fingerprint: + run = previous[0] + start = checkpoint.get('line', 0) + conn.execute("UPDATE ingest_runs SET status='running',error=NULL WHERE id=?", (run,)) + else: + raise RuntimeError('Pending run belongs to a different archive; resolve it explicitly') + else: + run = inserted_id(conn,'INSERT INTO ingest_runs(source_id,checkpoint_json) VALUES(?,?)', (source,json.dumps({'archive':fingerprint,'line':0}))) + conn.commit() + line_number = start + last = time.monotonic() + monitored = Path(storage_path) if storage_path else (Path(db).parent if db else Path(archive).parent) with zipfile.ZipFile(archive) as zipped: # Reading the member to EOF also checks its CRC. with zipped.open('arxiv-metadata-oai-snapshot.json') as stream: @@ -54,20 +56,21 @@ def ingest(db, archive): paper,record = store.record('arxiv',external_id,'https://arxiv.org/abs/'+external_id,title,abstract,authors,{'arxiv':external_id,'doi':row.get('doi')},metadata,(row.get('versions') or [{}])[0].get('created')) refresh_categories(store,paper,record,row.get('categories') or '') if line_number % 500 == 0: - if shutil.disk_usage(Path(db).parent).free < 4*1024**3: + if shutil.disk_usage(monitored).free < 4*1024**3: raise RuntimeError('Stopped before exhausting disk: less than 4 GiB free') conn.execute('UPDATE ingest_runs SET records_seen=?,records_written=?,checkpoint_json=? WHERE id=?',(line_number,line_number,json.dumps({'archive':fingerprint,'line':line_number}),run)) conn.commit() if time.monotonic()-last >= 20: print(json.dumps({'run':run,'rows':line_number}),flush=True) last = time.monotonic() - conn.execute("UPDATE ingest_runs SET status='complete',finished_at=CURRENT_TIMESTAMP,records_seen=?,records_written=?,checkpoint_json=? WHERE id=?",(line_number,line_number,json.dumps({'archive':fingerprint,'line':line_number}),run)) + conn.execute("UPDATE ingest_runs SET status='complete',finished_at=?,records_seen=?,records_written=?,checkpoint_json=? WHERE id=?",(timestamp(),line_number,line_number,json.dumps({'archive':fingerprint,'line':line_number}),run)) conn.commit() print(json.dumps({'status':'complete','rows':line_number}),flush=True) except BaseException as exc: conn.rollback() - conn.execute("UPDATE ingest_runs SET status='failed',error=? WHERE id=?",(str(exc),run)) - conn.commit() + if run is not None: + conn.execute("UPDATE ingest_runs SET status='failed',error=? WHERE id=?",(str(exc),run)) + conn.commit() raise finally: conn.close() @@ -75,10 +78,11 @@ def ingest(db, archive): def main(argv=None, *, prog=None): parser = argparse.ArgumentParser(prog=prog) - parser.add_argument('--db', required=True) + database_arguments(parser) parser.add_argument('--archive', required=True) + parser.add_argument('--storage-path', help='Local filesystem to monitor for free space; specify the database volume when running beside PostgreSQL') args = parser.parse_args(argv) - ingest(args.db,args.archive) + ingest(args.db,args.archive,pg_service=args.pg_service,pg_schema=args.pg_schema,storage_path=args.storage_path) if __name__ == '__main__': diff --git a/paper_atlas/ingestion/backfill.py b/paper_atlas/ingestion/backfill.py index 7b973e6..35b296c 100644 --- a/paper_atlas/ingestion/backfill.py +++ b/paper_atlas/ingestion/backfill.py @@ -4,7 +4,7 @@ import hashlib import json from pathlib import Path -from paper_atlas.db.store import Store +from paper_atlas.db.backend import database_arguments, open_store from paper_atlas.db.submissions import attach_submission,identities_for_row from paper_atlas.collectors.pmlr_tracks import require_track_manifest, require_track_record @@ -19,15 +19,15 @@ def queue_files(path): return files -def backfill(db,files): +def backfill(db,files, *, pg_service=None, pg_schema=None): for filename in files: manifest=json.loads(Path(filename).with_suffix('.manifest.json').read_text()) require_track_manifest(manifest) if manifest['source'] in ('colm','neurips','iclr_program') and manifest.get('parser_version',0)<2: raise ValueError('Recollect old parser output before submission evidence backfill') - store=Store(db);c=store.conn;count=0 - before=c.execute('SELECT count(*) FROM submissions').fetchone()[0] + store=open_store(db,pg_service=pg_service,pg_schema=pg_schema);c=store.conn;count=0 try: + before=c.execute('SELECT count(*) FROM submissions').fetchone()[0] for filename in files: path=Path(filename);manifest=json.loads(path.with_suffix('.manifest.json').read_text());seen=set() with path.open() as stream: @@ -39,7 +39,7 @@ def backfill(db,files): raise ValueError('Evidence row differs from edition manifest') require_track_record(row, manifest) record=c.execute('SELECT r.id,r.paper_id,r.metadata_json FROM source_records r JOIN sources s ON s.id=r.source_id WHERE s.name=? AND r.external_id=?',(row['source'],row['external_id'])).fetchone() - edition=c.execute('SELECT e.id FROM editions e JOIN venues v ON v.id=e.venue_id WHERE v.acronym=? COLLATE NOCASE AND e.year=? AND e.track=?',(row['venue'],row['year'],row['track'])).fetchone() + edition=c.execute('SELECT e.id FROM editions e JOIN venues v ON v.id=e.venue_id WHERE lower(v.acronym)=lower(?) AND e.year=? AND e.track=?',(row['venue'],row['year'],row['track'])).fetchone() if not record or not edition:raise ValueError('Evidence backfill requires an already imported record and edition') if not c.execute('SELECT 1 FROM submission_sources ss JOIN submissions s ON s.id=ss.submission_id WHERE ss.source_record_id=? AND ss.source_submission_id=? AND s.edition_id=?',(record[0],row['external_id'],edition[0])).fetchone(): raise ValueError('No existing submission to backfill') @@ -62,8 +62,8 @@ def backfill(db,files): def main(argv=None, *, prog=None): - parser=argparse.ArgumentParser(prog=prog);parser.add_argument('--db',required=True);parser.add_argument('--queue',required=True) - args=parser.parse_args(argv);backfill(args.db,queue_files(args.queue)) + parser=argparse.ArgumentParser(prog=prog);database_arguments(parser);parser.add_argument('--queue',required=True) + args=parser.parse_args(argv);backfill(args.db,queue_files(args.queue),pg_service=args.pg_service,pg_schema=args.pg_schema) if __name__ == '__main__': diff --git a/paper_atlas/ingestion/records.py b/paper_atlas/ingestion/records.py index b71020d..930d2ed 100644 --- a/paper_atlas/ingestion/records.py +++ b/paper_atlas/ingestion/records.py @@ -4,12 +4,12 @@ import json import time from pathlib import Path -from paper_atlas.db.store import Store +from paper_atlas.db.backend import database_arguments, open_store, inserted_id from paper_atlas.db.submissions import identities_for_row,timestamp from paper_atlas.collectors.pmlr_tracks import require_track_manifest, require_track_record -def ingest_file(db,path): +def ingest_file(db,path, *, pg_service=None, pg_schema=None): path = Path(path) manifest_path = path.with_suffix('.manifest.json') manifest = json.loads(manifest_path.read_text()) if manifest_path.exists() else {} @@ -18,19 +18,20 @@ def ingest_file(db,path): raise ValueError('Recollect proceedings with parser_version>=2 before identity matching') if manifest.get('source')=='iclr_program' and manifest.get('parser_version',0)<2: raise ValueError('Recollect ICLR with parser_version>=2 before submission identity matching') - store = Store(db) - c = store.conn source_name = manifest.get('source') if not source_name: raise ValueError('A source manifest is required') - source_id = store.source(source_name,manifest['list_url']) - run = c.execute('INSERT INTO ingest_runs(source_id,checkpoint_json) VALUES(?,?)',(source_id,json.dumps({'file':str(path)}))).lastrowid - c.commit() + store = open_store(db, pg_service=pg_service, pg_schema=pg_schema) + c = store.conn + run = None count = 0 missing_abstracts = 0 seen = set() editions = set() try: + source_id = store.source(source_name,manifest['list_url']) + run = inserted_id(c,'INSERT INTO ingest_runs(source_id,checkpoint_json) VALUES(?,?)',(source_id,json.dumps({'file':str(path)}))) + c.commit() with path.open() as stream: for line in stream: row = json.loads(line) @@ -75,13 +76,14 @@ def ingest_file(db,path): raise ValueError(f'Manifest count {expected} differs from records {count}') complete=expected==count and manifest.get('status') not in ('partial','blocked') and not manifest.get('missing_abstract_count',0) and not missing_abstracts c.execute('INSERT INTO edition_sources(edition_id,source_id,list_url,status,expected_count,imported_count,last_success_at) VALUES(?,?,?,?,?,?,?) ON CONFLICT(edition_id,source_id,list_url) DO UPDATE SET status=excluded.status,expected_count=excluded.expected_count,imported_count=excluded.imported_count,last_success_at=excluded.last_success_at,last_error=NULL',(edition,source_id,manifest['list_url'],'complete' if complete else 'partial',expected,count,time.strftime('%Y-%m-%dT%H:%M:%SZ',time.gmtime()))) - c.execute("UPDATE ingest_runs SET status='complete',finished_at=CURRENT_TIMESTAMP,records_seen=?,records_written=? WHERE id=?",(count,count,run)) + c.execute("UPDATE ingest_runs SET status='complete',finished_at=?,records_seen=?,records_written=? WHERE id=?",(timestamp(),count,count,run)) c.commit() print(json.dumps({'file':str(path),'status':'imported','records':count,'missing_abstract_count':missing_abstracts}),flush=True) except BaseException as exc: c.rollback() - c.execute("UPDATE ingest_runs SET status='failed',finished_at=CURRENT_TIMESTAMP,error=? WHERE id=?",(str(exc),run)) - c.commit() + if run is not None: + c.execute("UPDATE ingest_runs SET status='failed',finished_at=?,records_seen=?,records_written=0,error=? WHERE id=?",(timestamp(),count,str(exc),run)) + c.commit() raise finally: c.close() @@ -89,11 +91,11 @@ def ingest_file(db,path): def main(argv=None, *, prog=None): parser = argparse.ArgumentParser(prog=prog) - parser.add_argument('--db',required=True) + database_arguments(parser) parser.add_argument('files',nargs='+') args = parser.parse_args(argv) for filename in args.files: - ingest_file(args.db,filename) + ingest_file(args.db,filename,pg_service=args.pg_service,pg_schema=args.pg_schema) if __name__ == '__main__': diff --git a/paper_atlas/postgres/writer.py b/paper_atlas/postgres/writer.py new file mode 100644 index 0000000..2556397 --- /dev/null +++ b/paper_atlas/postgres/writer.py @@ -0,0 +1,87 @@ +"""Native transactional writes reusing the shared identity/provenance policy.""" +import re +import zlib + +from paper_atlas.db.store import Store +from paper_atlas.postgres import text +from paper_atlas.postgres.client import connect +from paper_atlas.postgres.schema import TABLES, quote + + +class Result: + def __init__(self, cursor): + self.cursor = cursor + + def fetchone(self): + row = self.cursor.fetchone() + return tuple(map(text.decode, row)) if row is not None else None + + def fetchall(self): + return [tuple(map(text.decode, row)) for row in self.cursor.fetchall()] + + def __iter__(self): + return (tuple(map(text.decode, row)) for row in self.cursor) + + @property + def rowcount(self): + return self.cursor.rowcount + + +class WriteConnection: + postgres = True + + def __init__(self, service, schema=None): + self.connection = connect(service) + try: + if schema is None: + schema = self.connection.execute('SELECT current_schema()').fetchone()[0] + if not isinstance(schema, str) or not re.fullmatch(r'[a-z][a-z0-9_]{0,35}', schema): + raise ValueError('Invalid PostgreSQL write schema') + annotation = self.connection.execute("SELECT obj_description(oid,'pg_namespace') FROM pg_namespace WHERE nspname=%s", (schema,)).fetchone() + if not annotation or annotation[0] != 'paper-atlas:text-codec=' + text.CODEC: + raise RuntimeError('Writes require a published dataset with the supported text codec') + tables = {r[0] for r in self.connection.execute("SELECT table_name FROM information_schema.tables WHERE table_schema=%s AND table_type='BASE TABLE'", (schema,))} + if not set(TABLES) <= tables: + raise RuntimeError('Published business tables are incomplete or inaccessible') + if not self.connection.execute('SELECT pg_try_advisory_lock(%s,%s)', (615174, zlib.crc32(schema.encode()) & 0x7fffffff)).fetchone()[0]: + raise RuntimeError('Another writer holds this PostgreSQL dataset; retry after it closes') + self.schema = schema + self.connection.execute('SET search_path TO ' + quote(schema) + ',pg_catalog') + if not self.connection.execute("SELECT 1 FROM schema_migrations WHERE name='submission_source_baseline_v1'").fetchone(): + raise RuntimeError('Submission provenance baseline is missing') + self.connection.autocommit = False + except BaseException: + self.connection.close() + raise + + def execute(self, sql, params=()): + # The shared domain uses static SQL and bound parameters. It uses no + # SQLite-specific syntax; this adapter only binds/decodes text values. + return Result(self.connection.execute(sql.replace('?', '%s'), tuple(map(text.encode, params)))) + + def next_id(self, table): + if table not in ('submissions', 'decisions'): + raise ValueError('Unsupported explicit identity sequence') + return self.connection.execute("SELECT nextval(pg_get_serial_sequence(%s,'id'))", (quote(self.schema) + '.' + quote(table),)).fetchone()[0] + + def commit(self): + self.connection.commit() + + def rollback(self): + self.connection.rollback() + + def close(self): + # Ends the session and releases its advisory lock, including after an + # error/rollback. File/chunk commits deliberately retain the writer lock. + self.connection.close() + + def __enter__(self): + return self + + def __exit__(self, kind, error, traceback): + self.rollback() if kind else self.commit() + + +class PostgresStore(Store): + def __init__(self, service, schema=None): + self.conn = WriteConnection(service, schema) diff --git a/tests/test_postgres_ingestion.py b/tests/test_postgres_ingestion.py new file mode 100644 index 0000000..5527d61 --- /dev/null +++ b/tests/test_postgres_ingestion.py @@ -0,0 +1,220 @@ +from contextlib import closing, redirect_stdout +import io +import json +import os +from pathlib import Path +import sqlite3 +import tempfile +from types import SimpleNamespace +import unittest +from unittest.mock import patch +import uuid +import zipfile + +from paper_atlas.db.backend import open_store +from paper_atlas.db.store import Store +from paper_atlas.ingestion import arxiv +from paper_atlas.ingestion.backfill import backfill +from paper_atlas.ingestion.records import ingest_file +from paper_atlas.postgres.client import connect, ReadConnection +from paper_atlas.postgres.migration import migrate +from paper_atlas.postgres.schema import quote + + +class BackendSelectionTests(unittest.TestCase): + def test_ambiguous_targets_fail_before_opening_any_database(self): + for kwargs in ({}, {'db': 'example', 'pg_service': 'example'}, {'db': 'example', 'pg_schema': 'research'}): + with self.assertRaisesRegex(ValueError, 'Choose exactly one'): + open_store(**kwargs) + + +@unittest.skipUnless(os.environ.get('PAPER_ATLAS_TEST_PG_SERVICE'), 'Requires the approved remote PostgreSQL test service') +class PostgreSQLIngestionTests(unittest.TestCase): + def setUp(self): + self.service = os.environ['PAPER_ATLAS_TEST_PG_SERVICE'] + self.schema = 'pa_write_' + uuid.uuid4().hex[:12] + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + self.root = Path(self.temp.name) + source = self.root / 'seed.sqlite' + store = Store(source) + paper, record = store.record('arxiv', '2401.00001', 'https://arxiv.org/abs/2401.00001', + 'Existing reference paper', 'Existing abstract', []) + store.decision(paper, record, 'ICLR', 2024, 'main', 'seed', 'https://example.org/seed', 'rejected', 'Reject') + store.conn.execute("INSERT INTO paper_merges(former_paper_id,surviving_paper_id,reason) VALUES(900,?,'fixture')", (paper,)) + store.conn.execute("INSERT INTO submission_merges VALUES(5000,1,'fixture',CURRENT_TIMESTAMP)") + store.conn.execute("INSERT INTO decision_merges VALUES(6000,1,'fixture',CURRENT_TIMESTAMP)") + store.conn.execute("UPDATE sqlite_sequence SET seq=1000 WHERE name='papers'") + store.conn.commit(); store.conn.close() + self.addCleanup(self.cleanup) + with redirect_stdout(io.StringIO()): + migrate(source, self.service, self.root / 'migration.json', schema=self.schema, min_free_gib=0) + + def cleanup(self): + with connect(self.service) as pg: + names = pg.execute('SELECT nspname FROM pg_namespace WHERE starts_with(nspname,%s)', (self.schema,)).fetchall() + for name, in names: + pg.execute('DROP SCHEMA ' + quote(name) + ' CASCADE') + + def store(self): + return open_store(pg_service=self.service, pg_schema=self.schema) + + def row(self, source='openreview_snapshot', external='forumNew1', status='rejected', year=2025, venue='ICLR'): + return {'source': source, 'external_id': external, 'url': 'https://example.org/' + external, + 'venue': venue, 'year': year, 'track': 'main', 'status': status, + 'title': 'A newly observed shared paper', 'abstract': 'Before\x00survivor and literal\x01\x02', + 'authors': ['Jane Doe', 'John Smith'], 'identifiers': {'arxiv': '2501.99999'}, + 'observed_at': str(year) + '-01-20T00:00:00Z'} + + def file(self, name, rows, count=None): + path = self.root / (name + '.jsonl') + path.write_text(''.join(json.dumps(r, ensure_ascii=False) + '\n' for r in rows)) + first = rows[0] + manifest = {key: first[key] for key in ('source', 'venue', 'year', 'track')} + manifest.update(parser_version=2, list_url='https://example.org/' + name, + count=len(rows) if count is None else count, missing_abstract_count=0, status='complete') + path.with_suffix('.manifest.json').write_text(json.dumps(manifest)) + return path + + def ingest(self, path): + with redirect_stdout(io.StringIO()): + ingest_file(None, path, pg_service=self.service, pg_schema=self.schema) + + def test_native_writes_keep_nul_historical_ids_and_later_acceptance(self): + store = self.store() + try: + paper, record = store.record('public', 'new-forum', 'https://example.org/new', 'New research identity', + 'Left\x00right\x01\x02', [], {'arxiv': '2501.99999'}) + self.assertGreater(paper, 1000) + store.decision(paper, record, 'ICLR', 2025, 'main', 'new-forum', 'https://example.org/reject', 'rejected', 'Reject') + _, second = store.record('official', 'new-pub', 'https://example.org/pub', 'Changed title', None, [], {'arxiv': '2501.99999'}) + store.decision(paper, second, 'COLM', 2026, 'main', 'new-pub', 'https://example.org/accept', 'accepted', 'Accept') + store.conn.commit() + self.assertEqual(store.conn.execute('SELECT abstract FROM papers WHERE id=?', (paper,)).fetchone()[0], 'Left\x00right\x01\x02') + ids = store.conn.execute('SELECT s.id,d.id FROM submissions s JOIN decisions d ON d.submission_id=s.id WHERE s.paper_id=?', (paper,)).fetchall() + self.assertEqual(len(ids), 2) + self.assertTrue(all(sid > 5000 and did > 6000 for sid, did in ids)) + finally: + store.conn.close() + reader = ReadConnection(self.service, self.schema) + try: + from paper_atlas.querying import known_acceptances + self.assertEqual(known_acceptances(reader, paper)[0]['venue'], 'COLM') + finally: + reader.close() + + def test_file_import_is_idempotent_and_refreshes_only_own_tags(self): + row = self.row(); row['tags'] = [{'namespace': 'openreview_snapshot:area', 'value': 'theory'}] + path = self.file('reject', [row]); self.ingest(path) + store = self.store() + try: + paper = store.conn.execute("SELECT paper_id FROM identifiers WHERE scheme='arxiv' AND value='2501.99999'").fetchone()[0] + store.conn.execute("INSERT INTO tags(namespace,value) VALUES('custom:topic','manual') ON CONFLICT DO NOTHING") + tag = store.conn.execute("SELECT id FROM tags WHERE namespace='custom:topic'").fetchone()[0] + store.conn.execute('INSERT INTO paper_tags VALUES(?,?,NULL,?)', (paper, tag, 'manual')) + store.conn.commit() + finally: + store.conn.close() + row['tags'][0]['value'] = 'learning'; self.ingest(self.file('reject', [row])) + self.ingest(self.file('accept', [self.row('official', 'published', 'accepted', 2026, 'COLM')])) + store = self.store() + try: + self.assertEqual(store.conn.execute('SELECT count(*) FROM papers').fetchone()[0], 2) + self.assertEqual(store.conn.execute('SELECT count(*) FROM submissions WHERE paper_id=?', (paper,)).fetchone()[0], 2) + self.assertEqual(store.conn.execute('SELECT count(*) FROM decisions d JOIN submissions s ON s.id=d.submission_id WHERE s.paper_id=?', (paper,)).fetchone()[0], 2) + tags = set(store.conn.execute('SELECT t.namespace,t.value FROM paper_tags p JOIN tags t ON t.id=p.tag_id WHERE p.paper_id=?', (paper,))) + self.assertEqual(tags, {('custom:topic', 'manual'), ('openreview_snapshot:area', 'learning')}) + finally: + store.conn.close() + + def test_bad_manifest_rolls_back_whole_file_and_keeps_failed_run(self): + with self.assertRaisesRegex(ValueError, 'Manifest count'): + self.ingest(self.file('bad', [self.row()], count=2)) + store = self.store() + try: + self.assertEqual(store.conn.execute('SELECT count(*) FROM papers').fetchone()[0], 1) + self.assertEqual(store.conn.execute('SELECT count(*) FROM source_records').fetchone()[0], 1) + self.assertEqual(store.conn.execute('SELECT status,records_seen,records_written FROM ingest_runs').fetchone(), ('failed', 1, 0)) + finally: + store.conn.close() + + def test_writer_lock_survives_commit_and_rollback_but_readers_continue(self): + first = self.store() + try: + first.conn.commit() + with self.assertRaisesRegex(RuntimeError, 'Another writer'): + self.store() + first.conn.rollback() + with self.assertRaisesRegex(RuntimeError, 'Another writer'): + self.store() + reader = ReadConnection(self.service, self.schema) + try: + self.assertEqual(reader.execute('SELECT count(*) FROM papers').fetchone()[0], 1) + finally: + reader.close() + finally: + first.conn.close() + reopened = self.store(); reopened.conn.close() + + def test_unpublished_schema_is_rejected_without_bootstrapping_tables(self): + staging = self.schema + '_unfinished' + with connect(self.service) as pg: + pg.execute('CREATE SCHEMA ' + quote(staging)) + with self.assertRaisesRegex(RuntimeError, 'published dataset'): + open_store(pg_service=self.service, pg_schema=staging) + with connect(self.service) as pg: + self.assertEqual(pg.execute('SELECT count(*) FROM information_schema.tables WHERE table_schema=%s', (staging,)).fetchone()[0], 0) + + def test_backfill_merges_evidence_without_reimporting_or_losing_history(self): + files = [] + for source, external, title in [('official', 'legacy', 'First bibliographic name'), ('openreview_snapshot', 'forum', 'Renamed submission')]: + row = self.row(source, external, 'accepted') + row.update(title=title, authors=[], identifiers={}, evidence_url='https://example.org/same-decision') + path = self.file(source, [row]); self.ingest(path) + row['identifiers'] = {'openreview': 'sharedForum'} + self.file(source, [row]); files.append(path) + with redirect_stdout(io.StringIO()): + result = backfill(None, files, pg_service=self.service, pg_schema=self.schema) + self.assertEqual((result['submissions_before'], result['submissions_after']), (3, 2)) + store = self.store() + try: + self.assertEqual(store.conn.execute('SELECT count(*) FROM papers').fetchone()[0], 2) + self.assertEqual(store.conn.execute('SELECT count(*) FROM decision_sources').fetchone()[0], 3) + self.assertEqual(store.conn.execute('SELECT count(*) FROM decision_merges').fetchone()[0], 2) + finally: + store.conn.close() + + def test_arxiv_chunk_resume_preserves_committed_ids_and_category_evidence(self): + archive = self.root / 'arxiv.zip' + rows = [{'id': f'2501.{i:05d}', 'title': f'Checkpoint paper {i}', 'abstract': 'Metadata abstract', + 'authors': [], 'categories': 'cs.AI'} for i in range(1, 502)] + with zipfile.ZipFile(archive, 'w') as zipped: + zipped.writestr('arxiv-metadata-oai-snapshot.json', ''.join(json.dumps(r) + '\n' for r in rows)) + calls = 0 + original = arxiv.refresh_categories + def interrupt(*args): + nonlocal calls + calls += 1 + if calls == 501: + raise RuntimeError('controlled interruption') + return original(*args) + with patch('paper_atlas.ingestion.arxiv.shutil.disk_usage', return_value=SimpleNamespace(free=8 * 1024**3)), \ + patch('paper_atlas.ingestion.arxiv.refresh_categories', side_effect=interrupt), redirect_stdout(io.StringIO()): + with self.assertRaisesRegex(RuntimeError, 'controlled interruption'): + arxiv.ingest(None, archive, pg_service=self.service, pg_schema=self.schema) + store = self.store() + try: + self.assertEqual(store.conn.execute('SELECT count(*) FROM papers').fetchone()[0], 501) + preserved = store.conn.execute("SELECT paper_id FROM identifiers WHERE scheme='arxiv' AND value='2501.00001'").fetchone()[0] + finally: + store.conn.close() + with redirect_stdout(io.StringIO()): + arxiv.ingest(None, archive, pg_service=self.service, pg_schema=self.schema) + store = self.store() + try: + self.assertEqual(store.conn.execute('SELECT count(*) FROM papers').fetchone()[0], 502) + self.assertEqual(store.conn.execute("SELECT paper_id FROM identifiers WHERE scheme='arxiv' AND value='2501.00001'").fetchone()[0], preserved) + self.assertEqual(store.conn.execute('SELECT count(*) FROM paper_tags').fetchone()[0], 501) + self.assertEqual(store.conn.execute('SELECT status,records_written FROM ingest_runs').fetchone(), ('complete', 501)) + finally: + store.conn.close()