326 lines
19 KiB
Python
326 lines
19 KiB
Python
import contextlib
|
||
from contextlib import closing
|
||
import io
|
||
import itertools
|
||
import json
|
||
import os
|
||
from pathlib import Path
|
||
import sqlite3
|
||
import tempfile
|
||
import unittest
|
||
from unittest.mock import patch
|
||
import uuid
|
||
|
||
from paper_atlas.db.store import Store
|
||
from paper_atlas.postgres import schema, text
|
||
from paper_atlas.postgres.client import ReadConnection, Row, connect
|
||
from paper_atlas.postgres.migration import migrate
|
||
from paper_atlas import querying
|
||
|
||
|
||
class SchemaTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.temporary = tempfile.TemporaryDirectory()
|
||
self.addCleanup(self.temporary.cleanup)
|
||
self.path = Path(self.temporary.name) / 'source.sqlite'
|
||
self.store = Store(self.path)
|
||
self.addCleanup(self.store.conn.close)
|
||
|
||
def test_every_business_table_is_included_and_engine_tables_are_excluded(self):
|
||
plan = schema.inspect(self.store.conn)
|
||
self.assertEqual(set(plan), set(schema.TABLES))
|
||
self.assertEqual(len(plan), 21)
|
||
self.assertNotIn('papers_fts', plan)
|
||
self.assertEqual(plan['paper_tags']['pk'], ['paper_id', 'tag_id', 'evidence_key'])
|
||
self.assertEqual(len(plan['source_records']['foreign_keys']), 2)
|
||
sql = schema.table_sql('stage', 'papers', plan['papers'])
|
||
self.assertIn('GENERATED BY DEFAULT AS IDENTITY', sql)
|
||
self.assertIn('TEXT COLLATE "C"', sql)
|
||
|
||
def test_unknown_tables_and_check_constraints_fail_before_copy(self):
|
||
self.store.conn.execute('CREATE TABLE extra(id INTEGER PRIMARY KEY)')
|
||
with self.assertRaisesRegex(ValueError, 'Unexpected business'):
|
||
schema.inspect(self.store.conn)
|
||
self.store.conn.execute('DROP TABLE extra')
|
||
self.store.conn.execute('ALTER TABLE papers ADD COLUMN extra TEXT CHECK(extra IS NULL)')
|
||
with self.assertRaisesRegex(ValueError, 'CHECK'):
|
||
schema.inspect(self.store.conn)
|
||
|
||
def test_cutover_stops_legacy_writers_but_keeps_read_access(self):
|
||
self.store.conn.commit(); self.store.conn.close()
|
||
Path(str(self.path) + '.postgresql-cutover.json').write_text('{"status":"active"}')
|
||
with self.assertRaisesRegex(RuntimeError, 'cut over'):
|
||
Store(self.path)
|
||
with sqlite3.connect(self.path.as_uri() + '?mode=ro', uri=True) as c:
|
||
self.assertEqual(c.execute('SELECT count(*) FROM papers').fetchone()[0], 0)
|
||
|
||
def test_rows_support_mapping_and_positional_query_access(self):
|
||
row = Row(['id', 'title'], (7, 'A'))
|
||
self.assertEqual(row[0], row['id'])
|
||
self.assertEqual(dict(row), {'id': 7, 'title': 'A'})
|
||
|
||
def test_connection_services_cannot_be_password_urls(self):
|
||
with self.assertRaisesRegex(ValueError, 'service name'):
|
||
connect('postgresql://user:secret@host/db')
|
||
|
||
def test_checks_compare_literal_values_not_only_constraint_counts(self):
|
||
clauses = schema.checks_in("CREATE TABLE t(x TEXT CHECK(x IN ('a','b)')),y TEXT CHECK(y IS NULL))")
|
||
self.assertEqual(clauses, ["x IN ('a','b)')", 'y IS NULL'])
|
||
self.assertEqual(schema.normalize_check("x IN ( 'a' )"), schema.normalize_check("x in('a')"))
|
||
self.assertNotEqual(schema.normalize_check("x IN ('A')"), schema.normalize_check("x IN ('a')"))
|
||
|
||
def test_text_codec_is_bijective_and_preserves_key_order_and_prefixes(self):
|
||
alphabet = ['\x00', '\x01', '\x02', 'a', 'é', '🙂']
|
||
values = [''.join(v) for n in range(4) for v in itertools.product(alphabet, repeat=n)]
|
||
for value in values:
|
||
encoded = text.encode(value)
|
||
self.assertNotIn('\x00', encoded)
|
||
self.assertEqual(text.decode(encoded), value)
|
||
self.assertEqual(sorted(map(text.encode, values)), list(map(text.encode, sorted(values))))
|
||
for value in [None, 7, 0.5]:
|
||
self.assertEqual(text.encode(value), value)
|
||
self.assertEqual(text.decode(value), value)
|
||
for value in ['\x01', 'prefix\x01x']:
|
||
with self.assertRaisesRegex(ValueError, 'Invalid'):
|
||
text.decode(value)
|
||
|
||
|
||
@unittest.skipUnless(os.environ.get('PAPER_ATLAS_TEST_PG_SERVICE'), 'Requires an explicitly configured remote PostgreSQL test service')
|
||
class PostgreSQLIntegrationTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.service = os.environ['PAPER_ATLAS_TEST_PG_SERVICE']
|
||
self.target = 'pa_test_' + uuid.uuid4().hex[:14]
|
||
self.temporary = tempfile.TemporaryDirectory()
|
||
self.addCleanup(self.temporary.cleanup)
|
||
self.db = Path(self.temporary.name) / 'source.sqlite'
|
||
self.out = Path(self.temporary.name) / 'report.json'
|
||
store = Store(self.db)
|
||
paper, record = store.record('arxiv', '2401.00001', 'https://arxiv.org/abs/2401.00001',
|
||
'Quiet graph models — 研究', 'Graph models learn reinforcement policies.\nQuoted "text", tabs\tand Unicode α.',
|
||
['Jane Doe', 'John Smith'], {'arxiv': '2401.00001'}, {'nested': {'a': ['α', None, 1]}, 'raw': '\\n'})
|
||
store.record_tags(paper, record, 'arxiv', [{'namespace': 'arxiv:category', 'value': 'cs.AI'}])
|
||
store.decision(paper, record, 'ICLR', 2024, 'main', 'forum-reject', 'https://example.org/reject',
|
||
'rejected', 'Reject', identities=[{'scheme': 'openreview', 'value': 'forum-reject'}], observed_at='2024-01-01T00:00:00Z')
|
||
same, second = store.record('official', '2025/paper', 'https://example.org/accepted',
|
||
'Quiet graph models — 研究', 'Graph models learn reinforcement policies.', ['Jane Doe', 'John Smith'], {'arxiv': '2401.00001'})
|
||
self.assertEqual(same, paper)
|
||
store.decision(paper, second, 'COLM', 2025, 'main', 'forum-accept', 'https://example.org/accept',
|
||
'accepted', 'Accept', ['Conference on Language Modeling'], observed_at='2025-01-01T00:00:00Z')
|
||
store.record('arxiv', '2401.00002', 'https://arxiv.org/abs/2401.00002', 'Empty abstract', None, [])
|
||
tag = store.conn.execute("INSERT INTO tags(namespace,value) VALUES('custom:topic','manual')").lastrowid
|
||
store.conn.execute('INSERT INTO paper_tags VALUES(?,?,NULL,?)', (paper, tag, 'manual:one'))
|
||
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_schemas)
|
||
|
||
def cleanup_schemas(self):
|
||
with connect(self.service) as pg:
|
||
names = pg.execute('SELECT nspname FROM pg_namespace WHERE starts_with(nspname,%s)', (self.target,)).fetchall()
|
||
for name, in names:
|
||
pg.execute('DROP SCHEMA ' + schema.quote(name) + ' CASCADE')
|
||
|
||
def migrate(self):
|
||
with contextlib.redirect_stdout(io.StringIO()):
|
||
return migrate(self.db, self.service, self.out, schema=self.target, min_free_gib=0)
|
||
|
||
def published(self):
|
||
with connect(self.service) as pg:
|
||
return bool(pg.execute('SELECT 1 FROM pg_namespace WHERE nspname=%s', (self.target,)).fetchone())
|
||
|
||
def test_full_content_history_tags_sequences_and_readonly_queries(self):
|
||
report = self.migrate()
|
||
self.assertEqual(report['status'], 'complete')
|
||
self.assertEqual(len(report['tables']), 21)
|
||
for entry in report['tables'].values():
|
||
self.assertEqual(entry['source_rows'], entry['target_rows'])
|
||
self.assertEqual(entry['source_sha256'], entry['target_sha256'])
|
||
with connect(self.service) as pg, pg.transaction(force_rollback=True):
|
||
prefix = schema.quote(self.target) + '.'
|
||
value = pg.execute('INSERT INTO ' + prefix + 'papers(title,title_key) VALUES(%s,%s) RETURNING id', ('new', 'new')).fetchone()[0]
|
||
self.assertGreater(value, 1000)
|
||
record = pg.execute('INSERT INTO ' + prefix + "source_records(source_id,external_id,paper_id,url,title,abstract,authors_json,metadata_json) SELECT source_id,'new',%s,url,title,abstract,authors_json,metadata_json FROM " + prefix + 'source_records LIMIT 1 RETURNING id', (value,)).fetchone()[0]
|
||
edition = pg.execute('SELECT id FROM ' + prefix + 'editions LIMIT 1').fetchone()[0]
|
||
submission = pg.execute('INSERT INTO ' + prefix + 'submissions(paper_id,edition_id,source_record_id,source_submission_id) VALUES(%s,%s,%s,%s) RETURNING id', (value, edition, record, 'new')).fetchone()[0]
|
||
decision = pg.execute('INSERT INTO ' + prefix + "decisions(submission_id,status,raw_value,evidence_url,observed_at) VALUES(%s,'unknown','fixture','https://example.org/new','2026-01-01T00:00:00Z') RETURNING id", (submission,)).fetchone()[0]
|
||
self.assertGreater(submission, 5000)
|
||
self.assertGreater(decision, 6000)
|
||
real = ReadConnection
|
||
with patch('paper_atlas.postgres.client.ReadConnection', lambda service: real(service, self.target)):
|
||
output = io.StringIO()
|
||
with contextlib.redirect_stdout(output):
|
||
querying.main(['--pg-service', self.service, '--query', 'graph models', '--status', 'rejected', '--outcome', 'accepted'])
|
||
results = json.loads(output.getvalue())
|
||
self.assertEqual(len(results), 1)
|
||
self.assertEqual(results[0]['known_acceptances'][0]['venue'], 'COLM')
|
||
output = io.StringIO()
|
||
with contextlib.redirect_stdout(output):
|
||
querying.main(['--pg-service', self.service, '--venue', 'conference on language modeling', '--tag', 'custom:topic=manual', '--paper', '1'])
|
||
detail = json.loads(output.getvalue())[0]
|
||
self.assertEqual(len(detail['submissions']), 2)
|
||
self.assertEqual({t['namespace'] for t in detail['tags']}, {'arxiv:category', 'custom:topic'})
|
||
output = io.StringIO()
|
||
with contextlib.redirect_stdout(output):
|
||
querying.main(['--pg-service', self.service, '--stats'])
|
||
stats = json.loads(output.getvalue())
|
||
self.assertEqual(stats['papers'], 2)
|
||
self.assertEqual(stats['papers_with_known_acceptance'], 1)
|
||
reader = real(self.service, self.target)
|
||
try:
|
||
import psycopg
|
||
with self.assertRaises(psycopg.errors.ReadOnlySqlTransaction):
|
||
reader.execute('DELETE FROM papers')
|
||
finally:
|
||
reader.close()
|
||
|
||
def test_history_audit_matches_sqlite_for_actual_postgresql_snapshot(self):
|
||
from paper_atlas.audits.histories import report
|
||
from tests.test_histories import add_history_fixture
|
||
store = Store(self.db)
|
||
try:
|
||
add_history_fixture(store)
|
||
finally:
|
||
store.conn.close()
|
||
expected = report(self.db)
|
||
self.assertEqual(self.migrate()['status'], 'complete')
|
||
actual = report(pg_service=self.service, pg_schema=self.target)
|
||
self.assertEqual(actual['engine'], 'postgresql')
|
||
self.assertEqual(actual['schema'], self.target)
|
||
for key in ('counts', 'latest_attempt_status_counts', 'evidence_relationship_checks', 'examples'):
|
||
self.assertEqual(actual[key], expected[key], key)
|
||
|
||
def interrupt_after_first_table(self):
|
||
calls = 0
|
||
def interrupt(*args):
|
||
nonlocal calls
|
||
calls += 1
|
||
if calls == 3:
|
||
raise RuntimeError('controlled interruption')
|
||
with patch('paper_atlas.postgres.migration.free_space', side_effect=interrupt):
|
||
with self.assertRaisesRegex(RuntimeError, 'controlled interruption'):
|
||
self.migrate()
|
||
self.assertFalse(self.published())
|
||
|
||
def test_interrupted_copy_resumes_without_duplicate_rows(self):
|
||
self.interrupt_after_first_table()
|
||
self.assertEqual(self.migrate()['status'], 'complete')
|
||
|
||
def test_resume_rejects_a_changed_source_snapshot(self):
|
||
self.interrupt_after_first_table()
|
||
with closing(sqlite3.connect(self.db)) as source, source:
|
||
source.execute("UPDATE papers SET abstract='changed' WHERE id=1")
|
||
with self.assertRaisesRegex(ValueError, 'exact same'):
|
||
self.migrate()
|
||
self.assertFalse(self.published())
|
||
|
||
def test_full_hash_detects_tampering_in_a_completed_staging_table(self):
|
||
self.interrupt_after_first_table()
|
||
with connect(self.service) as pg:
|
||
stage = pg.execute('SELECT stage FROM ' + schema.quote(self.target + '_migration') + '.run WHERE id=1').fetchone()[0]
|
||
pg.execute('UPDATE ' + schema.quote(stage) + ".sources SET homepage='tampered'")
|
||
with self.assertRaisesRegex(ValueError, 'content verification'):
|
||
self.migrate()
|
||
self.assertFalse(self.published())
|
||
|
||
def test_foreign_key_violation_never_publishes_target(self):
|
||
with closing(sqlite3.connect(self.db)) as source, source:
|
||
source.execute('PRAGMA foreign_keys=OFF')
|
||
source.execute("INSERT INTO identifiers VALUES('orphan','bad',999999)")
|
||
import psycopg
|
||
with self.assertRaises(psycopg.errors.ForeignKeyViolation):
|
||
self.migrate()
|
||
self.assertFalse(self.published())
|
||
|
||
def test_published_target_cannot_be_overwritten(self):
|
||
self.migrate()
|
||
self.out.unlink()
|
||
with self.assertRaisesRegex(ValueError, 'never be overwritten'):
|
||
self.migrate()
|
||
|
||
def test_nul_and_literal_escape_survive_full_copy_keys_query_and_hash(self):
|
||
original = 'before\x00survivor; literal\x01\x01 and \x01\x02 — 研究'
|
||
keys = ['a\x00z', 'a\x01z', 'a\x01\x02z', 'a\x02z']
|
||
with closing(sqlite3.connect(self.db)) as source, source:
|
||
source.execute('UPDATE papers SET abstract=? WHERE id=1', (original,))
|
||
source.execute('UPDATE source_records SET abstract=? WHERE id=1', (original,))
|
||
source.executemany("INSERT INTO identifiers VALUES('control',?,1)", [(k,) for k in keys])
|
||
source.execute("INSERT INTO tags(namespace,value) VALUES('custom:control',?)", (original,))
|
||
tag = source.execute("SELECT id FROM tags WHERE namespace='custom:control'").fetchone()[0]
|
||
source.executemany('INSERT INTO paper_tags VALUES(1,?,NULL,?)', [(tag, k) for k in keys])
|
||
report = self.migrate()
|
||
self.assertEqual(report['text_codec'], text.CODEC)
|
||
for entry in report['tables'].values():
|
||
self.assertEqual(entry['source_sha256'], entry['target_sha256'])
|
||
self.assertNotEqual(report['tables']['source_records']['source_sha256'],
|
||
report['tables']['source_records']['target_storage_sha256'])
|
||
with connect(self.service) as pg:
|
||
self.assertEqual(pg.execute("SELECT obj_description(oid,'pg_namespace') FROM pg_namespace WHERE nspname=%s", (self.target,)).fetchone()[0], 'paper-atlas:text-codec=' + text.CODEC)
|
||
raw = pg.execute('SELECT abstract FROM ' + schema.quote(self.target) + '.source_records WHERE id=1').fetchone()[0]
|
||
self.assertNotIn('\x00', raw)
|
||
self.assertEqual(text.decode(raw), original)
|
||
reader = ReadConnection(self.service, self.target)
|
||
try:
|
||
self.assertEqual(reader.execute('SELECT abstract FROM source_records WHERE id=?', (1,)).fetchone()[0], original)
|
||
self.assertEqual([r[0] for r in reader.execute("SELECT value FROM identifiers WHERE scheme='control' ORDER BY value")], sorted(keys))
|
||
self.assertEqual(reader.execute("SELECT value FROM identifiers WHERE scheme='control' AND value=?", (keys[0],)).fetchone()[0], keys[0])
|
||
finally:
|
||
reader.close()
|
||
real = ReadConnection
|
||
with patch('paper_atlas.postgres.client.ReadConnection', lambda service: real(service, self.target)):
|
||
output = io.StringIO()
|
||
with contextlib.redirect_stdout(output):
|
||
querying.main(['--pg-service', self.service, '--paper', '1', '--tag', 'custom:control=' + original])
|
||
result = json.loads(output.getvalue())[0]
|
||
self.assertEqual(result['abstract'], original)
|
||
self.assertEqual(next(s['abstract'] for s in result['sources'] if s['external_id'] == '2401.00001'), original)
|
||
output = io.StringIO()
|
||
with contextlib.redirect_stdout(output):
|
||
querying.main(['--pg-service', self.service, '--query', 'survivor'])
|
||
self.assertEqual([r['id'] for r in json.loads(output.getvalue())], [1])
|
||
|
||
def test_old_completed_staging_table_is_upgraded_atomically_on_resume(self):
|
||
original = 'https://example.org/literal\x01\x02'
|
||
with closing(sqlite3.connect(self.db)) as source, source:
|
||
source.execute('UPDATE sources SET homepage=? WHERE id=1', (original,))
|
||
self.interrupt_after_first_table()
|
||
with connect(self.service) as pg:
|
||
control = schema.quote(self.target + '_migration')
|
||
stage = pg.execute('SELECT stage FROM ' + control + '.run WHERE id=1').fetchone()[0]
|
||
# Reproduce the old implementation: raw SOH text and no codec metadata.
|
||
pg.execute('UPDATE ' + schema.quote(stage) + '.sources SET homepage=%s WHERE id=1', (original,))
|
||
pg.execute('ALTER TABLE ' + control + '.run DROP COLUMN text_codec')
|
||
report = self.migrate()
|
||
self.assertEqual(report['legacy_text_upgrade_rows']['sources'], 1)
|
||
self.assertEqual(report['tables']['sources']['source_sha256'], report['tables']['sources']['target_sha256'])
|
||
reader = ReadConnection(self.service, self.target)
|
||
try:
|
||
self.assertEqual(reader.execute('SELECT homepage FROM sources WHERE id=1').fetchone()[0], original)
|
||
finally:
|
||
reader.close()
|
||
|
||
def test_legacy_unique_keys_can_upgrade_after_constraints_were_built(self):
|
||
keys = ['x\x01', 'x\x01\x02']
|
||
with closing(sqlite3.connect(self.db)) as source, source:
|
||
source.executemany("INSERT INTO identifiers VALUES('control',?,1)", [(k,) for k in keys])
|
||
# Emulate an old run interrupted after constraints: no storage encoding.
|
||
with patch('paper_atlas.postgres.text.encode', side_effect=lambda value: value), \
|
||
patch('paper_atlas.postgres.migration.sequences', side_effect=RuntimeError('after constraints')):
|
||
with self.assertRaisesRegex(RuntimeError, 'after constraints'):
|
||
self.migrate()
|
||
with connect(self.service) as pg:
|
||
pg.execute('ALTER TABLE ' + schema.quote(self.target + '_migration') + '.run DROP COLUMN text_codec')
|
||
report = self.migrate()
|
||
self.assertEqual(report['status'], 'complete')
|
||
self.assertEqual(report['tables']['identifiers']['source_sha256'], report['tables']['identifiers']['target_sha256'])
|
||
reader = ReadConnection(self.service, self.target)
|
||
try:
|
||
self.assertEqual([r[0] for r in reader.execute("SELECT value FROM identifiers WHERE scheme='control' ORDER BY value")], keys)
|
||
finally:
|
||
reader.close()
|
||
|
||
|
||
if __name__ == '__main__':
|
||
unittest.main()
|