56 lines
2.3 KiB
Python
56 lines
2.3 KiB
Python
|
"""Run FULL citation refresh with incremental progress logging"""
|
||
|
|
import asyncio, sys, os
|
||
|
|
from datetime import datetime
|
||
|
|
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
|
||
|
|
from sqlalchemy.orm import sessionmaker
|
||
|
|
from sqlalchemy import select
|
||
|
|
from app.models.literature import GlobalLiterature
|
||
|
|
from app.services.pubmed_api import fetch_citedby_counts, update_citation_counts
|
||
|
|
|
||
|
|
DATABASE_URL = 'postgresql+asyncpg://scilit:scilit_prod_2026@postgres:5432/scilit'
|
||
|
|
|
||
|
|
async def main():
|
||
|
|
engine = create_async_engine(DATABASE_URL, pool_size=4)
|
||
|
|
async_session = sessionmaker(engine, class_=AsyncSession)
|
||
|
|
|
||
|
|
print(f"[{datetime.now().isoformat()}] Starting full citation refresh...", flush=True)
|
||
|
|
|
||
|
|
async with async_session() as db:
|
||
|
|
# Get all PMIDs
|
||
|
|
result = await db.execute(select(GlobalLiterature.pmid).order_by(GlobalLiterature.created_at.desc()))
|
||
|
|
all_pmids = [r for (r,) in result]
|
||
|
|
total = len(all_pmids)
|
||
|
|
print(f"[{datetime.now().isoformat()}] Found {total} PMIDs to process", flush=True)
|
||
|
|
|
||
|
|
# Process in chunks of 5000 (matching the admin API limit=5000 behavior)
|
||
|
|
BATCH = 5000
|
||
|
|
total_updated = 0
|
||
|
|
total_errors = 0
|
||
|
|
|
||
|
|
for i in range(0, len(all_pmids), BATCH):
|
||
|
|
batch = all_pmids[i:i+BATCH]
|
||
|
|
counts = await fetch_citedby_counts(batch)
|
||
|
|
|
||
|
|
async with async_session()() as s:
|
||
|
|
updated = 0
|
||
|
|
for pmid, cited in counts.items():
|
||
|
|
from sqlalchemy import update as sql_update
|
||
|
|
await s.execute(
|
||
|
|
sql_update(GlobalLiterature)
|
||
|
|
.where(GlobalLiterature.pmid == pmid)
|
||
|
|
.values(cited_by_count=cited, updated_at=datetime.now())
|
||
|
|
)
|
||
|
|
updated += 1
|
||
|
|
await s.commit()
|
||
|
|
|
||
|
|
total_updated += updated
|
||
|
|
errors_in_batch = len(batch) - len(counts)
|
||
|
|
total_errors += errors_in_batch
|
||
|
|
pct = min(100, (i + BATCH) / total * 100)
|
||
|
|
print(f"[{datetime.now().isoformat()}] batch {i//BATCH+1}/{(total+BATCH-1)//BATCH}: +{updated} updated ({pct:.0f}%)", flush=True)
|
||
|
|
|
||
|
|
print(f"[{datetime.now().isoformat()}] Done. Updated: {total_updated}, Errors: {total_errors}", flush=True)
|
||
|
|
await engine.dispose()
|
||
|
|
|
||
|
|
asyncio.run(main())
|