Bound exact event migration counts with primary-key pagination

docs/event-identity.md

@@ -17,7 +17,7 @@ PGSERVICEFILE=/etc/paper-atlas/pg_service.conf paper-atlas maintain event-identi
   --pg-service foundry-research --pg-schema research --out /data/paper-index/event-plan.json
 ```
 
-检查计划的数据库身份、原 edition IDs、业务表数量、目标来源与历史指纹。已有未标记 WWW1995 edition 必须逐投稿都有固定 DBLP Fall 卷证据，才会只为原 edition 填入 `www-1995-fall`；缺证或混合卷会拒绝计划。本文命令不表示生产已经迁移。
+检查计划的数据库身份、原 edition IDs、业务表数量、目标来源与历史指纹。精确计数在同一快照中按主键分页，每页最多5000行，只读取主键；不使用估计行数。CLI向stderr持续报告计数进度，避免大库长时间无输出。已有未标记 WWW1995 edition 必须逐投稿都有固定 DBLP Fall 卷证据，才会只为原 edition 填入 `www-1995-fall`；缺证或混合卷会拒绝计划。本文命令不表示生产已经迁移。
 
 应用时传入同一文件及命令输出的文件 SHA256：
 

paper_atlas/ingestion/event_identity.py

@@ -5,6 +5,7 @@ import hashlib
 import json
 from pathlib import Path
 import sqlite3
+import sys
 
 from paper_atlas.db.backend import database_arguments, open_read_connection, require_sqlite_writable
 from paper_atlas.db.events import has_event_key, pinned_www_fall
@@ -21,7 +22,39 @@ def values(rows):
     return [[r[i] for i in range(len(r))] for r in rows]
 
 
-def inspect(c):
+def exact_counts(c, *, page_size=5000, progress=None):
+    """Count every row in the same snapshot through bounded primary-key pages."""
+    if not isinstance(page_size, int) or page_size < 1:
+        raise ValueError('Positive count page size required')
+    result = {}
+    for table in TABLES:
+        if table == 'schema_migrations':
+            continue
+        if getattr(c, 'postgres', False):
+            columns = [r[0] for r in c.execute("SELECT a.attname FROM pg_constraint pc JOIN pg_attribute a ON a.attrelid=pc.conrelid AND a.attnum=ANY(pc.conkey) WHERE pc.conrelid=?::regclass AND pc.contype='p' ORDER BY array_position(pc.conkey,a.attnum)", (table,))]
+        else:
+            columns = [r[1] for r in sorted(c.execute('PRAGMA table_info('+quote(table)+')'), key=lambda r:r[5]) if r[5]]
+        if not columns:
+            raise ValueError('Primary key required for exact count: '+table)
+        fields = ','.join(quote(column) for column in columns)
+        last = None; count = 0; pages = 0
+        if progress:
+            progress({'table':table, 'rows_counted':0, 'complete':False})
+        while True:
+            where = (' WHERE ('+fields+') > ('+','.join('?' for _ in columns)+')') if last is not None else ''
+            batch = values(c.execute('SELECT '+fields+' FROM '+quote(table)+where+' ORDER BY '+fields+' LIMIT ?', (*last, page_size) if last is not None else (page_size,)).fetchall())
+            count += len(batch); pages += 1
+            complete = len(batch) < page_size
+            if progress and (complete or pages % 20 == 0):
+                progress({'table':table, 'rows_counted':count, 'complete':complete})
+            if complete:
+                break
+            last = batch[-1]
+        result[table] = count
+    return result
+
+
+def inspect(c, *, progress=None):
     enabled = has_event_key(c)
     if getattr(c, 'postgres', False):
         columns = [r[0] for r in c.execute("SELECT column_name FROM information_schema.columns WHERE table_schema=current_schema() AND table_name='editions' ORDER BY ordinal_position")]
@@ -38,7 +71,7 @@ def inspect(c):
         raise ValueError('Unexpected edition schema; require a reviewed migration')
     rows = values(c.execute('SELECT id,venue_id,year,track FROM editions ORDER BY id'))
     keys = values(c.execute('SELECT id,event_key FROM editions ORDER BY id')) if enabled else [[r[0], ''] for r in rows]
-    counts = {table:c.execute('SELECT count(*) FROM '+table).fetchone()[0] for table in TABLES if table != 'schema_migrations'}
+    counts = exact_counts(c, progress=progress)
     targets = values(c.execute("SELECT e.id FROM editions e JOIN venues v ON v.id=e.venue_id WHERE v.acronym='WWW' AND e.year=1995 AND e.track='main' ORDER BY e.id"))
     assignments = []
     for (eid,) in targets:
@@ -77,18 +110,18 @@ def database_identity(c, db=None):
     return {'engine':'sqlite', 'path':str(Path(db).resolve())}
 
 
-def plan(db=None, *, pg_service=None, pg_schema=None):
+def plan(db=None, *, pg_service=None, pg_schema=None, progress=None):
     with closing(open_read_connection(db, pg_service=pg_service, pg_schema=pg_schema)) as c:
-        return {'format':1, 'database':database_identity(c, db), 'before':inspect(c)}
+        return {'format':1, 'database':database_identity(c, db), 'before':inspect(c, progress=progress)}
 
 
-def apply_plan(c, proposal):
+def apply_plan(c, proposal, *, progress=None):
     before = proposal['before']
     if proposal.get('format') != 1:
         raise ValueError('Unsupported event identity plan')
     if getattr(c, 'postgres', False):
         c.execute('LOCK TABLE editions IN ACCESS EXCLUSIVE MODE')
-    current = inspect(c)
+    current = inspect(c, progress=progress)
     signature = digest(proposal)
     marker = MARKER + signature
     if c.execute('SELECT 1 FROM schema_migrations WHERE name=?', (marker,)).fetchone():
@@ -124,7 +157,7 @@ def apply_plan(c, proposal):
         if result.rowcount != 1:
             raise ValueError('Edition assignment changed during upgrade')
     c.execute('INSERT INTO schema_migrations(name) VALUES(?)', (marker,))
-    after = inspect(c)
+    after = inspect(c, progress=progress)
     expected_keys = dict(before['event_keys'])
     expected_keys.update({a['edition_id']:a['event_key'] for a in before['assignments']})
     if (after['editions_sha256'] != before['editions_sha256'] or after['protected'] != before['protected']
@@ -135,14 +168,14 @@ def apply_plan(c, proposal):
     return {'status':'applied', 'updated_editions':len(before['assignments']), 'after':after}
 
 
-def upgrade(proposal, db=None, *, pg_service=None, pg_schema=None):
+def upgrade(proposal, db=None, *, pg_service=None, pg_schema=None, progress=None):
     if pg_service:
         from paper_atlas.postgres.writer import WriteConnection
         with closing(WriteConnection(pg_service, pg_schema)) as c:
             if proposal['database'] != database_identity(c):
                 raise ValueError('Plan targets a different database')
             try:
-                result = apply_plan(c, proposal); c.commit(); return result
+                result = apply_plan(c, proposal, progress=progress); c.commit(); return result
             except BaseException:
                 c.rollback(); raise
     require_sqlite_writable(db)
@@ -154,7 +187,7 @@ def upgrade(proposal, db=None, *, pg_service=None, pg_schema=None):
         c.execute('PRAGMA foreign_keys=OFF')  # Outside the transaction; no cascading deletes.
         try:
             c.execute('BEGIN IMMEDIATE')
-            result = apply_plan(c, proposal); c.commit(); return result
+            result = apply_plan(c, proposal, progress=progress); c.commit(); return result
         except BaseException:
             c.rollback(); raise
         finally:
@@ -171,13 +204,14 @@ def main(argv=None, *, prog=None):
     target = Path(args.out)
     if target.exists():
         raise ValueError('Refusing to overwrite an existing plan/report')
+    progress = lambda event: print(json.dumps({'phase':'exact_row_count', **event}), file=sys.stderr, flush=True)
     if args.apply_plan:
         raw = Path(args.apply_plan).read_bytes()
         if hashlib.sha256(raw).hexdigest() != args.plan_sha256:
             raise ValueError('Plan file SHA256 differs from approved bytes')
-        result = upgrade(json.loads(raw), args.db, pg_service=args.pg_service, pg_schema=args.pg_schema)
+        result = upgrade(json.loads(raw), args.db, pg_service=args.pg_service, pg_schema=args.pg_schema, progress=progress)
     else:
-        result = plan(args.db, pg_service=args.pg_service, pg_schema=args.pg_schema)
+        result = plan(args.db, pg_service=args.pg_service, pg_schema=args.pg_schema, progress=progress)
     target.write_text(json.dumps(result, ensure_ascii=False, indent=2)+'\n')
     print(json.dumps({'out':str(target), 'sha256':hashlib.sha256(target.read_bytes()).hexdigest(),
                       'status':result.get('status','planned')}), flush=True)

tests/test_event_identity.py

@@ -13,7 +13,7 @@ from paper_atlas.db.store import Store
 from paper_atlas.db.backend import open_store
 from paper_atlas.db.events import has_event_key
 from paper_atlas.collectors.dblp import XML_SHA
-from paper_atlas.ingestion.event_identity import plan, upgrade
+from paper_atlas.ingestion.event_identity import plan, upgrade, exact_counts
 from paper_atlas.ingestion.records import ingest_file
 from paper_atlas.ingestion.backfill import backfill
 from paper_atlas.postgres.client import connect
@@ -129,6 +129,15 @@ class EventIdentityTests(unittest.TestCase):
         with sqlite3.connect(self.db)as c:
             self.assertEqual(c.execute("SELECT e.event_key FROM submission_identifiers si JOIN submissions s ON s.id=si.submission_id JOIN editions e ON e.id=s.edition_id WHERE si.scheme='openreview' AND si.value='forumFall'").fetchone()[0],'www-1995-fall')
 
+    def test_exact_paged_counts_cover_composite_key_boundaries(self):
+        s=Store(self.db);self.addCleanup(s.conn.close)
+        for i in range(7):add(s,'entry'+str(i),'www-1995-fall')
+        s.conn.commit();events=[];counts=exact_counts(s.conn,page_size=2,progress=events.append)
+        self.assertEqual(counts['source_records'],7)
+        self.assertEqual(counts['identifiers'],8)
+        self.assertEqual(counts['submission_sources'],7)
+        self.assertEqual([e['rows_counted']for e in events if e['table']=='identifiers'and e['complete']],[8])
+
     def test_legacy_single_event_writes_and_reads_remain_available(self):
         legacy(self.db);s=Store(self.db);self.addCleanup(s.conn.close)
         add(s,'ordinary',venue='COLM',year=2025);s.conn.commit()
@@ -178,6 +187,14 @@ class PostgreSQLEventIdentityTests(unittest.TestCase):
         with redirect_stdout(io.StringIO())as out:querying.main(['--pg-service',self.service,'--pg-schema',self.schema,'--paper',str(paper)])
         row=json.loads(out.getvalue())[0];self.assertEqual(len(row['submissions']),2);self.assertEqual(len(row['known_acceptances']),1)
 
+    def test_native_paged_counts_preserve_encoded_composite_key_order(self):
+        with opened_store(pg_service=self.service,pg_schema=self.schema)as s:
+            for value in ['a\x00b','a\x01b','a\x02b','z']:
+                s.conn.execute('INSERT INTO identifiers(scheme,value,paper_id) VALUES(?,?,?)',('custom',value,1))
+            s.conn.commit();counts=exact_counts(s.conn,page_size=2)
+            self.assertEqual(counts['identifiers'],6)
+            self.assertEqual(counts['submissions'],1)
+
     def test_stale_native_plan_is_rejected_and_lock_released(self):
         proposal=plan(pg_service=self.service,pg_schema=self.schema)
         with opened_store(pg_service=self.service,pg_schema=self.schema)as s:s.conn.execute("UPDATE editions SET track='unspecified'");s.conn.commit()