79 lines
4.4 KiB
Python
79 lines
4.4 KiB
Python
"""Stream the official metadata ZIP into SQLite. No PDFs, source files or builds."""
|
|
import argparse
|
|
import json
|
|
import time
|
|
import shutil
|
|
import zipfile
|
|
from pathlib import Path
|
|
|
|
|
|
from index_core import Store
|
|
|
|
|
|
def ingest(db, archive):
|
|
store = Store(db)
|
|
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()
|
|
try:
|
|
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:
|
|
for line_number, line in enumerate(stream, 1):
|
|
if line_number <= start:
|
|
continue
|
|
row = json.loads(line)
|
|
external_id = row['id']
|
|
title = ' '.join(row['title'].split())
|
|
abstract = row.get('abstract', '').strip()
|
|
authors = row.get('authors_parsed') or row.get('authors') or []
|
|
metadata = {k:v for k,v in row.items() if k not in ('title','abstract','authors','authors_parsed')}
|
|
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'))
|
|
conn.execute("DELETE FROM paper_tags WHERE source_record_id=? AND evidence_key='arxiv:categories'", (record,))
|
|
for category in (row.get('categories') or '').split():
|
|
conn.execute("INSERT OR IGNORE INTO tags(namespace,value) VALUES('arxiv:category',?)",(category,))
|
|
tag = conn.execute("SELECT id FROM tags WHERE namespace='arxiv:category' AND value=?",(category,)).fetchone()[0]
|
|
conn.execute('INSERT OR IGNORE INTO paper_tags VALUES(?,?,?,?)',(paper,tag,record,'arxiv:categories'))
|
|
if line_number % 500 == 0:
|
|
if shutil.disk_usage(Path(db).parent).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.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()
|
|
raise
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--db', required=True)
|
|
parser.add_argument('--archive', required=True)
|
|
args = parser.parse_args()
|
|
ingest(args.db,args.archive)
|