Preserve NUL text through PostgreSQL migration and resume

tests/test_postgres.py

@@ -1,6 +1,7 @@
 import contextlib
 from contextlib import closing
 import io
+import itertools
 import json
 import os
 from pathlib import Path
@@ -11,7 +12,7 @@ from unittest.mock import patch
 import uuid
 
 from paper_atlas.db.store import Store
-from paper_atlas.postgres import schema
+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
@@ -68,6 +69,21 @@ class SchemaTests(unittest.TestCase):
         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):
@@ -207,6 +223,87 @@ class PostgreSQLIntegrationTests(unittest.TestCase):
         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()