221 lines
12 KiB
Python
221 lines
12 KiB
Python
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()
|