from __future__ import annotations import logging from dataclasses import dataclass from datetime import datetime from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.db.models import ( PublicationPdfJob, User, ) from app.services.publications.pdf_queue_common import ( PDF_STATUS_FAILED, PDF_STATUS_QUEUED, PDF_STATUS_RESOLVED, PDF_STATUS_RUNNING, event_row, queued_job, utcnow, ) from app.services.publications.pdf_queue_queries import ( missing_pdf_candidates, retry_item_for_publication_id, ) from app.services.publications.pdf_queue_resolution import schedule_rows from app.services.publications.types import PublicationListItem from app.settings import settings PDF_STATUS_UNTRACKED = "untracked" PDF_EVENT_QUEUED = "queued" logger = logging.getLogger(__name__) @dataclass(frozen=True) class PdfRequeueResult: publication_exists: bool queued: bool @dataclass(frozen=True) class PdfBulkQueueResult: requested_count: int queued_count: int def _publication_ids(rows: list[PublicationListItem]) -> list[int]: return sorted({row.publication_id for row in rows}) def _status_from_job(row: PublicationListItem, job: PublicationPdfJob | None) -> str: if row.pdf_url: return PDF_STATUS_RESOLVED if job is None: return PDF_STATUS_UNTRACKED return job.status def _item_from_row_and_job( row: PublicationListItem, job: PublicationPdfJob | None, ) -> PublicationListItem: return PublicationListItem( publication_id=row.publication_id, scholar_profile_id=row.scholar_profile_id, scholar_label=row.scholar_label, title=row.title, year=row.year, citation_count=row.citation_count, venue_text=row.venue_text, pub_url=row.pub_url, pdf_url=row.pdf_url, is_read=row.is_read, is_favorite=row.is_favorite, first_seen_at=row.first_seen_at, is_new_in_latest_run=row.is_new_in_latest_run, pdf_status=_status_from_job(row, job), pdf_attempt_count=int(job.attempt_count) if job is not None else 0, pdf_failure_reason=job.last_failure_reason if job is not None else None, pdf_failure_detail=job.last_failure_detail if job is not None else None, display_identifier=row.display_identifier, ) def _queueable_rows( rows: list[PublicationListItem], *, max_items: int, ) -> list[PublicationListItem]: bounded = max(0, int(max_items)) if bounded == 0: return [] candidates = [row for row in rows if not row.pdf_url] return candidates[:bounded] def _auto_retry_interval_seconds() -> int: return max(int(settings.pdf_auto_retry_interval_seconds), 1) def _auto_retry_first_interval_seconds() -> int: return max(int(settings.pdf_auto_retry_first_interval_seconds), 1) def _auto_retry_max_attempts() -> int: return max(int(settings.pdf_auto_retry_max_attempts), 1) def _retry_interval_seconds_for_attempt_count(attempt_count: int) -> int: if int(attempt_count) <= 1: return _auto_retry_first_interval_seconds() return _auto_retry_interval_seconds() def _cooldown_active( *, last_attempt_at: datetime | None, attempt_count: int, ) -> bool: if last_attempt_at is None: return False elapsed = (utcnow() - last_attempt_at).total_seconds() return elapsed < _retry_interval_seconds_for_attempt_count(int(attempt_count)) def _can_enqueue_job( job: PublicationPdfJob | None, *, force_retry: bool, ) -> bool: if job is None: return True if job.status in {PDF_STATUS_QUEUED, PDF_STATUS_RUNNING}: return False if force_retry: return job.status in {PDF_STATUS_FAILED, PDF_STATUS_RESOLVED, PDF_STATUS_UNTRACKED} if job.status == PDF_STATUS_RESOLVED: return False if int(job.attempt_count) >= _auto_retry_max_attempts(): return False return not _cooldown_active( last_attempt_at=job.last_attempt_at, attempt_count=int(job.attempt_count), ) def _mark_job_queued(job: PublicationPdfJob, *, user_id: int) -> None: now = utcnow() job.status = PDF_STATUS_QUEUED job.queued_at = now job.last_requested_by_user_id = user_id job.last_failure_reason = None job.last_failure_detail = None job.last_source = None def _state_map(jobs: list[PublicationPdfJob]) -> dict[int, PublicationPdfJob]: return {int(job.publication_id): job for job in jobs} async def _jobs_for_publication_ids( db_session: AsyncSession, *, publication_ids: list[int], ) -> dict[int, PublicationPdfJob]: if not publication_ids: return {} result = await db_session.execute( select(PublicationPdfJob).where(PublicationPdfJob.publication_id.in_(publication_ids)) ) return _state_map(list(result.scalars())) async def overlay_pdf_job_state( db_session: AsyncSession, *, rows: list[PublicationListItem], ) -> list[PublicationListItem]: if not rows: return [] jobs = await _jobs_for_publication_ids( db_session, publication_ids=_publication_ids(rows), ) return [_item_from_row_and_job(row, jobs.get(row.publication_id)) for row in rows] async def _enqueue_rows( db_session: AsyncSession, *, user_id: int, rows: list[PublicationListItem], force_retry: bool, ) -> list[PublicationListItem]: if not rows: return [] queued: list[PublicationListItem] = [] jobs = await _jobs_for_publication_ids( db_session, publication_ids=_publication_ids(rows), ) for row in rows: job = jobs.get(row.publication_id) if not _can_enqueue_job(job, force_retry=force_retry): continue if job is None: job = queued_job(publication_id=row.publication_id, user_id=user_id) jobs[row.publication_id] = job db_session.add(job) else: _mark_job_queued(job, user_id=user_id) db_session.add( event_row( publication_id=row.publication_id, user_id=user_id, event_type=PDF_EVENT_QUEUED, status=PDF_STATUS_QUEUED, ) ) queued.append(row) if queued: await db_session.commit() return queued # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- async def enqueue_missing_pdf_jobs( db_session: AsyncSession, *, user_id: int, request_email: str | None, rows: list[PublicationListItem], max_items: int, ) -> list[int]: queueable = _queueable_rows(rows, max_items=max_items) queued_rows = await _enqueue_rows( db_session, user_id=user_id, rows=queueable, force_retry=False, ) schedule_rows(user_id=user_id, request_email=request_email, rows=queued_rows) return [row.publication_id for row in queued_rows] async def enqueue_retry_pdf_job( db_session: AsyncSession, *, user_id: int, request_email: str | None, row: PublicationListItem, ) -> bool: queued_rows = await _enqueue_rows( db_session, user_id=user_id, rows=[row], force_retry=True, ) schedule_rows(user_id=user_id, request_email=request_email, rows=queued_rows) return bool(queued_rows) async def enqueue_retry_pdf_job_for_publication_id( db_session: AsyncSession, *, user_id: int, request_email: str | None, publication_id: int, ) -> PdfRequeueResult: row = await retry_item_for_publication_id( db_session, publication_id=publication_id, ) if row is None: return PdfRequeueResult(publication_exists=False, queued=False) queued = await enqueue_retry_pdf_job( db_session, user_id=user_id, request_email=request_email, row=row, ) return PdfRequeueResult(publication_exists=True, queued=queued) async def enqueue_all_missing_pdf_jobs( db_session: AsyncSession, *, user_id: int, request_email: str | None, limit: int = 1000, ) -> PdfBulkQueueResult: candidates = await missing_pdf_candidates(db_session, limit=limit) queued_rows = await _enqueue_rows( db_session, user_id=user_id, rows=candidates, force_retry=True, ) schedule_rows(user_id=user_id, request_email=request_email, rows=queued_rows) return PdfBulkQueueResult( requested_count=len(candidates), queued_count=len(queued_rows), ) async def drain_ready_jobs( db_session: AsyncSession, *, limit: int, max_attempts: int, ) -> int: result = await db_session.execute(select(User.id).where(User.is_active.is_(True)).order_by(User.id.asc()).limit(1)) system_user_id = result.scalar_one_or_none() if system_user_id is None: return 0 bulk_result = await enqueue_all_missing_pdf_jobs( db_session, user_id=system_user_id, request_email=settings.unpaywall_email, limit=limit, ) return bulk_result.queued_count