From 2a6e573a523433ef3b29a6228bc5347f7594a46a Mon Sep 17 00:00:00 2001 From: Justin Visser Date: Fri, 27 Feb 2026 16:18:21 +0100 Subject: [PATCH] refactor: extract helpers from long functions in ingestion core - Extract _load_unenriched_publications, _enrich_batch, _flush_and_sweep_duplicates from enrich_pending_publications - Extract _sanitize_titles shared helper for title cleaning - Extract _run_iteration_and_complete and _inline_enrich_and_finalize from run_for_user - Extract _run_first_pass and _run_depth_pass from run_scholar_iteration - Extract _finalize_successful_queue_job and _finalize_failed_queue_job from _finalize_queue_job_after_run Co-Authored-By: Claude Opus 4.6 --- app/services/ingestion/application.py | 59 ++++- app/services/ingestion/enrichment.py | 223 ++++++++++--------- app/services/ingestion/queue_runner.py | 138 ++++++------ app/services/ingestion/scholar_processing.py | 160 ++++++++----- 4 files changed, 345 insertions(+), 235 deletions(-) diff --git a/app/services/ingestion/application.py b/app/services/ingestion/application.py index 01e4dd1..3f0ad48 100644 --- a/app/services/ingestion/application.py +++ b/app/services/ingestion/application.py @@ -576,6 +576,46 @@ class ScholarIngestionService: alert_network_failure_threshold=alert_network_failure_threshold, alert_retry_scheduled_threshold=alert_retry_scheduled_threshold, ) + progress, failure_summary, alert_summary = await self._run_iteration_and_complete( + db_session, + run=run, + scholars=scholars, + user_id=user_id, + start_cstart_map=start_cstart_map, + paging=paging, + thresholds=thresholds, + auto_queue_continuations=auto_queue_continuations, + queue_delay_seconds=queue_delay_seconds, + idempotency_key=idempotency_key, + ) + user_settings = await user_settings_service.get_or_create_settings(db_session, user_id=user_id) + await self._inline_enrich_and_finalize( + db_session, run=run, user_settings=user_settings, intended_final_status=run.status + ) + _log_run_completed( + run=run, + user_id=user_id, + scholars=scholars, + progress=progress, + failure_summary=failure_summary, + alert_summary=alert_summary, + ) + return run_execution_summary(run=run, scholars=scholars, progress=progress) + + async def _run_iteration_and_complete( + self, + db_session: AsyncSession, + *, + run: CrawlRun, + scholars: list[ScholarProfile], + user_id: int, + start_cstart_map: dict[int, int], + paging: dict[str, Any], + thresholds: dict[str, Any], + auto_queue_continuations: bool, + queue_delay_seconds: int, + idempotency_key: str | None, + ) -> tuple[RunProgress, RunFailureSummary, RunAlertSummary]: progress = await run_scholar_iteration( db_session, pagination=self._pagination, @@ -601,6 +641,16 @@ class ScholarIngestionService: if intended_final_status not in (RunStatus.CANCELED,): run.status = RunStatus.RESOLVING await db_session.commit() + return progress, failure_summary, alert_summary + + async def _inline_enrich_and_finalize( + self, + db_session: AsyncSession, + *, + run: CrawlRun, + user_settings: Any, + intended_final_status: RunStatus, + ) -> None: try: await self._enrichment.enrich_pending_publications( db_session, @@ -612,15 +662,6 @@ class ScholarIngestionService: if run.status == RunStatus.RESOLVING: run.status = intended_final_status await db_session.commit() - _log_run_completed( - run=run, - user_id=user_id, - scholars=scholars, - progress=progress, - failure_summary=failure_summary, - alert_summary=alert_summary, - ) - return run_execution_summary(run=run, scholars=scholars, progress=progress) async def _try_acquire_user_lock( self, diff --git a/app/services/ingestion/enrichment.py b/app/services/ingestion/enrichment.py index e63a6aa..9232efb 100644 --- a/app/services/ingestion/enrichment.py +++ b/app/services/ingestion/enrichment.py @@ -25,6 +25,18 @@ from app.settings import settings logger = logging.getLogger(__name__) +def _sanitize_titles(publications: list) -> list[str]: + titles = [] + for p in publications: + raw = getattr(p, "title_raw", None) or getattr(p, "title", None) + if not raw or not raw.strip(): + continue + safe = " ".join(re.sub(r"[^\w\s]", " ", raw).split()) + if safe: + titles.append(safe) + return titles + + class EnrichmentRunner: """Post-run OpenAlex enrichment logic. @@ -44,31 +56,15 @@ class EnrichmentRunner: raise RuntimeError(f"Missing crawl_run for run_id={run_id}.") return status == RunStatus.CANCELED - async def enrich_pending_publications( + async def _load_unenriched_publications( self, db_session: AsyncSession, *, run_id: int, - openalex_api_key: str | None = None, - ) -> None: - """Enrich unenriched publications with OpenAlex data. - - Stops immediately on budget exhaustion (429 with $0 remaining). - Sleeps 60s and continues on transient rate limits. - """ - from app.services.openalex.client import ( - OpenAlexBudgetExhaustedError, - OpenAlexClient, - OpenAlexRateLimitError, - ) - from app.services.openalex.matching import find_best_match - + ) -> tuple[int, list[Publication]]: run_result = await db_session.execute(select(CrawlRun.user_id).where(CrawlRun.id == run_id)) user_id = run_result.scalar_one() - - now = datetime.now(UTC) - cooldown_threshold = now - timedelta(days=7) - + cooldown_threshold = datetime.now(UTC) - timedelta(days=7) stmt = ( select(Publication) .join(ScholarPublication) @@ -84,88 +80,51 @@ class EnrichmentRunner: .distinct() ) result = await db_session.execute(stmt) - publications = list(result.scalars().all()) + return user_id, list(result.scalars().all()) - if not publications: - return + async def _enrich_batch( + self, + db_session: AsyncSession, + *, + batch: list[Publication], + run_id: int, + openalex_works: list, + now: datetime, + arxiv_lookup_allowed: bool, + ) -> tuple[bool, bool]: + from app.services.openalex.matching import find_best_match - resolved_key = openalex_api_key or settings.openalex_api_key - client = OpenAlexClient(api_key=resolved_key, mailto=settings.crossref_api_mailto) - batch_size = 25 - arxiv_lookup_allowed = True - - for i in range(0, len(publications), batch_size): + for p in batch: if await self.run_is_canceled(db_session, run_id=run_id): logger.info("ingestion.enrichment_aborted", extra={"run_id": run_id}) - return - batch = publications[i : i + batch_size] - titles = [ - " ".join(re.sub(r"[^\w\s]", " ", p.title_raw).split()) - for p in batch - if p.title_raw and p.title_raw.strip() - ] - - if not titles: - continue - - try: - openalex_works = await client.get_works_by_filter( - {"title.search": "|".join(titles)}, limit=batch_size * 3 - ) - except OpenAlexBudgetExhaustedError: - structured_log( - logger, - "warning", - "ingestion.openalex_budget_exhausted", - run_id=run_id, - ) - break - except OpenAlexRateLimitError: - structured_log( - logger, - "warning", - "ingestion.openalex_rate_limited", - run_id=run_id, - ) - await asyncio.sleep(60) - continue - except Exception as e: - structured_log( - logger, - "warning", - "ingestion.openalex_enrichment_failed", - error=str(e), - run_id=run_id, - ) - continue - - for p in batch: - if await self.run_is_canceled(db_session, run_id=run_id): - logger.info("ingestion.enrichment_aborted", extra={"run_id": run_id}) - return - - p.openalex_last_attempt_at = now - arxiv_lookup_allowed = await self._discover_identifiers_for_enrichment( - db_session, - publication=p, - run_id=run_id, - allow_arxiv_lookup=arxiv_lookup_allowed, - ) - - match = find_best_match( - target_title=p.title_raw, - target_year=p.year, - target_authors=p.author_text or "", - candidates=openalex_works, - ) - if match: - p.year = match.publication_year or p.year - p.citation_count = match.cited_by_count or p.citation_count - p.pdf_url = match.oa_url or p.pdf_url - p.openalex_enriched = True + return False, arxiv_lookup_allowed + p.openalex_last_attempt_at = now + arxiv_lookup_allowed = await self._discover_identifiers_for_enrichment( + db_session, + publication=p, + run_id=run_id, + allow_arxiv_lookup=arxiv_lookup_allowed, + ) + match = find_best_match( + target_title=p.title_raw, + target_year=p.year, + target_authors=p.author_text or "", + candidates=openalex_works, + ) + if match: + p.year = match.publication_year or p.year + p.citation_count = match.cited_by_count or p.citation_count + p.pdf_url = match.oa_url or p.pdf_url + p.openalex_enriched = True + return True, arxiv_lookup_allowed + async def _flush_and_sweep_duplicates( + self, + db_session: AsyncSession, + *, + run_id: int, + ) -> None: await db_session.flush() - from app.services.publications.dedup import sweep_identifier_duplicates merge_count = await sweep_identifier_duplicates(db_session) @@ -178,6 +137,64 @@ class EnrichmentRunner: run_id=run_id, ) + async def enrich_pending_publications( + self, + db_session: AsyncSession, + *, + run_id: int, + openalex_api_key: str | None = None, + ) -> None: + from app.services.openalex.client import ( + OpenAlexBudgetExhaustedError, + OpenAlexClient, + OpenAlexRateLimitError, + ) + + _, publications = await self._load_unenriched_publications(db_session, run_id=run_id) + if not publications: + return + + resolved_key = openalex_api_key or settings.openalex_api_key + client = OpenAlexClient(api_key=resolved_key, mailto=settings.crossref_api_mailto) + batch_size = 25 + now = datetime.now(UTC) + arxiv_lookup_allowed = True + + for i in range(0, len(publications), batch_size): + if await self.run_is_canceled(db_session, run_id=run_id): + logger.info("ingestion.enrichment_aborted", extra={"run_id": run_id}) + return + batch = publications[i : i + batch_size] + titles = _sanitize_titles(batch) + if not titles: + continue + try: + openalex_works = await client.get_works_by_filter( + {"title.search": "|".join(titles)}, limit=batch_size * 3 + ) + except OpenAlexBudgetExhaustedError: + structured_log(logger, "warning", "ingestion.openalex_budget_exhausted", run_id=run_id) + break + except OpenAlexRateLimitError: + structured_log(logger, "warning", "ingestion.openalex_rate_limited", run_id=run_id) + await asyncio.sleep(60) + continue + except Exception as e: + structured_log(logger, "warning", "ingestion.openalex_enrichment_failed", error=str(e), run_id=run_id) + continue + should_continue, arxiv_lookup_allowed = await self._enrich_batch( + db_session, + batch=batch, + run_id=run_id, + openalex_works=openalex_works, + now=now, + arxiv_lookup_allowed=arxiv_lookup_allowed, + ) + if not should_continue: + return + + await self._flush_and_sweep_duplicates(db_session, run_id=run_id) + async def _discover_identifiers_for_enrichment( self, db_session: AsyncSession, @@ -275,23 +292,15 @@ class EnrichmentRunner: for i in range(0, len(publications), batch_size): batch = publications[i : i + batch_size] - - titles = [] - for p in batch: - if not p.title: - continue - safe_title = re.sub(r"[^\w\s]", " ", p.title) - safe_title = " ".join(safe_title.split()) - if safe_title: - titles.append(safe_title) - + titles = _sanitize_titles(batch) if not titles: enriched.extend(batch) continue - query = "|".join(t for t in titles) try: - openalex_works = await client.get_works_by_filter({"title.search": query}, limit=batch_size * 3) + openalex_works = await client.get_works_by_filter( + {"title.search": "|".join(titles)}, limit=batch_size * 3 + ) except Exception as e: logger.warning( "ingestion.openalex_enrichment_failed", diff --git a/app/services/ingestion/queue_runner.py b/app/services/ingestion/queue_runner.py index 4eaaa98..45d9dd7 100644 --- a/app/services/ingestion/queue_runner.py +++ b/app/services/ingestion/queue_runner.py @@ -281,74 +281,72 @@ class QueueJobRunner: await self._reschedule_queue_job_after_exception(job, exc=exc) return None - async def _finalize_queue_job_after_run(self, job: queue_service.ContinuationQueueJob, run_summary) -> None: - session_factory = get_session_factory() - async with session_factory() as session: - if int(run_summary.failed_count) <= 0: - queue_item = await queue_service.reset_attempt_count(session, job_id=job.id) - await session.commit() - if queue_item is None: - structured_log( - logger, - "info", - "scheduler.queue_item_resolved", - queue_item_id=job.id, - user_id=job.user_id, - run_id=run_summary.crawl_run_id, - status=run_summary.status.value, - ) - return - structured_log( - logger, - "info", - "scheduler.queue_item_progressed", - queue_item_id=job.id, - user_id=job.user_id, - run_id=run_summary.crawl_run_id, - status=run_summary.status.value, - attempt_count=int(queue_item.attempt_count), - ) - return - queue_item = await queue_service.increment_attempt_count(session, job_id=job.id) - if queue_item is None: - await session.commit() - structured_log( - logger, - "info", - "scheduler.queue_item_resolved", - queue_item_id=job.id, - user_id=job.user_id, - run_id=run_summary.crawl_run_id, - status=run_summary.status.value, - ) - return - if int(queue_item.attempt_count) >= self._continuation_max_attempts: - await queue_service.mark_dropped(session, job_id=job.id, reason="max_attempts_after_run") - await session.commit() - structured_log( - logger, - "warning", - "scheduler.queue_item_dropped_max_attempts_after_run", - queue_item_id=job.id, - user_id=job.user_id, - attempt_count=queue_item.attempt_count, - run_id=run_summary.crawl_run_id, - status=run_summary.status.value, - ) - return - delay_seconds = queue_service.compute_backoff_seconds( - base_seconds=self._continuation_base_delay_seconds, - attempt_count=int(queue_item.attempt_count), - max_seconds=self._continuation_max_delay_seconds, - ) - await queue_service.reschedule_job( - session, - job_id=job.id, - delay_seconds=delay_seconds, - reason=queue_item.reason, - error=queue_item.last_error, + async def _finalize_successful_queue_job(self, session, job, run_summary) -> None: + queue_item = await queue_service.reset_attempt_count(session, job_id=job.id) + await session.commit() + if queue_item is None: + structured_log( + logger, + "info", + "scheduler.queue_item_resolved", + queue_item_id=job.id, + user_id=job.user_id, + run_id=run_summary.crawl_run_id, + status=run_summary.status.value, ) + return + structured_log( + logger, + "info", + "scheduler.queue_item_progressed", + queue_item_id=job.id, + user_id=job.user_id, + run_id=run_summary.crawl_run_id, + status=run_summary.status.value, + attempt_count=int(queue_item.attempt_count), + ) + + async def _finalize_failed_queue_job(self, session, job, run_summary) -> None: + queue_item = await queue_service.increment_attempt_count(session, job_id=job.id) + if queue_item is None: await session.commit() + structured_log( + logger, + "info", + "scheduler.queue_item_resolved", + queue_item_id=job.id, + user_id=job.user_id, + run_id=run_summary.crawl_run_id, + status=run_summary.status.value, + ) + return + if int(queue_item.attempt_count) >= self._continuation_max_attempts: + await queue_service.mark_dropped(session, job_id=job.id, reason="max_attempts_after_run") + await session.commit() + structured_log( + logger, + "warning", + "scheduler.queue_item_dropped_max_attempts_after_run", + queue_item_id=job.id, + user_id=job.user_id, + attempt_count=queue_item.attempt_count, + run_id=run_summary.crawl_run_id, + status=run_summary.status.value, + ) + return + delay_seconds = queue_service.compute_backoff_seconds( + base_seconds=self._continuation_base_delay_seconds, + attempt_count=int(queue_item.attempt_count), + max_seconds=self._continuation_max_delay_seconds, + ) + await queue_service.reschedule_job( + session, + job_id=job.id, + delay_seconds=delay_seconds, + reason=queue_item.reason, + error=queue_item.last_error, + ) + await session.commit() structured_log( logger, "info", @@ -361,6 +359,14 @@ class QueueJobRunner: delay_seconds=delay_seconds, ) + async def _finalize_queue_job_after_run(self, job: queue_service.ContinuationQueueJob, run_summary) -> None: + session_factory = get_session_factory() + async with session_factory() as session: + if int(run_summary.failed_count) <= 0: + await self._finalize_successful_queue_job(session, job, run_summary) + else: + await self._finalize_failed_queue_job(session, job, run_summary) + async def _run_queue_job(self, job: queue_service.ContinuationQueueJob) -> None: if await self._drop_queue_job_if_max_attempts(job): return diff --git a/app/services/ingestion/scholar_processing.py b/app/services/ingestion/scholar_processing.py index 3096dee..2bbddfd 100644 --- a/app/services/ingestion/scholar_processing.py +++ b/app/services/ingestion/scholar_processing.py @@ -528,6 +528,87 @@ def unexpected_scholar_exception_outcome( ) +async def _run_first_pass( + db_session: AsyncSession, + *, + scholars: list[ScholarProfile], + pagination: PaginationEngine, + run: CrawlRun, + user_id: int, + start_cstart_map: dict[int, int], + scholar_kwargs: dict[str, Any], + request_delay_seconds: int, + queue_delay_seconds: int, + progress: RunProgress, +) -> dict[int, int]: + first_pass_cstarts: dict[int, int] = {} + for index, scholar in enumerate(scholars): + await db_session.refresh(run) + if run.status == RunStatus.CANCELED: + structured_log(logger, "info", "ingestion.run_canceled", run_id=run.id, user_id=user_id) + return first_pass_cstarts + if index > 0 and request_delay_seconds > 0: + await asyncio.sleep(float(request_delay_seconds)) + start_cstart = int(start_cstart_map.get(int(scholar.id), 0)) + outcome = await process_scholar( + db_session, + pagination=pagination, + run=run, + scholar=scholar, + user_id=user_id, + start_cstart=start_cstart, + max_pages_per_scholar=1, + auto_queue_continuations=False, + queue_delay_seconds=queue_delay_seconds, + **scholar_kwargs, + ) + apply_outcome_to_progress(progress=progress, outcome=outcome) + resume_cstart = outcome.result_entry.get("continuation_cstart") + if resume_cstart is not None and int(resume_cstart) > start_cstart: + first_pass_cstarts[int(scholar.id)] = int(resume_cstart) + return first_pass_cstarts + + +async def _run_depth_pass( + db_session: AsyncSession, + *, + scholars: list[ScholarProfile], + first_pass_cstarts: dict[int, int], + pagination: PaginationEngine, + run: CrawlRun, + user_id: int, + scholar_kwargs: dict[str, Any], + request_delay_seconds: int, + remaining_max: int, + auto_queue_continuations: bool, + queue_delay_seconds: int, + progress: RunProgress, +) -> None: + for index, scholar in enumerate(scholars): + resume_cstart = first_pass_cstarts.get(int(scholar.id)) + if resume_cstart is None: + continue + await db_session.refresh(run) + if run.status == RunStatus.CANCELED: + structured_log(logger, "info", "ingestion.run_canceled", run_id=run.id, user_id=user_id) + break + if index > 0 and request_delay_seconds > 0: + await asyncio.sleep(float(request_delay_seconds)) + outcome = await process_scholar( + db_session, + pagination=pagination, + run=run, + scholar=scholar, + user_id=user_id, + start_cstart=resume_cstart, + max_pages_per_scholar=remaining_max, + auto_queue_continuations=auto_queue_continuations, + queue_delay_seconds=queue_delay_seconds, + **scholar_kwargs, + ) + apply_outcome_to_progress(progress=progress, outcome=outcome) + + async def run_scholar_iteration( db_session: AsyncSession, *, @@ -555,60 +636,33 @@ async def run_scholar_iteration( "rate_limit_backoff_seconds": rate_limit_backoff_seconds, "page_size": page_size, } - - # ── Pass 1: first page of every scholar (breadth-first) ────────── - first_pass_cstarts: dict[int, int] = {} - for index, scholar in enumerate(scholars): - await db_session.refresh(run) - if run.status == RunStatus.CANCELED: - structured_log(logger, "info", "ingestion.run_canceled", run_id=run.id, user_id=user_id) - return progress - if index > 0 and request_delay_seconds > 0: - await asyncio.sleep(float(request_delay_seconds)) - start_cstart = int(start_cstart_map.get(int(scholar.id), 0)) - outcome = await process_scholar( - db_session, - pagination=pagination, - run=run, - scholar=scholar, - user_id=user_id, - start_cstart=start_cstart, - max_pages_per_scholar=1, - auto_queue_continuations=False, - queue_delay_seconds=queue_delay_seconds, - **scholar_kwargs, - ) - apply_outcome_to_progress(progress=progress, outcome=outcome) - resume_cstart = outcome.result_entry.get("continuation_cstart") - if resume_cstart is not None and int(resume_cstart) > start_cstart: - first_pass_cstarts[int(scholar.id)] = int(resume_cstart) - - # ── Pass 2: remaining pages for each scholar (depth) ───────────── + first_pass_cstarts = await _run_first_pass( + db_session, + scholars=scholars, + pagination=pagination, + run=run, + user_id=user_id, + start_cstart_map=start_cstart_map, + scholar_kwargs=scholar_kwargs, + request_delay_seconds=request_delay_seconds, + queue_delay_seconds=queue_delay_seconds, + progress=progress, + ) remaining_max = max(max_pages_per_scholar - 1, 0) if remaining_max <= 0: return progress - - for index, scholar in enumerate(scholars): - resume_cstart = first_pass_cstarts.get(int(scholar.id)) - if resume_cstart is None: - continue - await db_session.refresh(run) - if run.status == RunStatus.CANCELED: - structured_log(logger, "info", "ingestion.run_canceled", run_id=run.id, user_id=user_id) - break - if index > 0 and request_delay_seconds > 0: - await asyncio.sleep(float(request_delay_seconds)) - outcome = await process_scholar( - db_session, - pagination=pagination, - run=run, - scholar=scholar, - user_id=user_id, - start_cstart=resume_cstart, - max_pages_per_scholar=remaining_max, - auto_queue_continuations=auto_queue_continuations, - queue_delay_seconds=queue_delay_seconds, - **scholar_kwargs, - ) - apply_outcome_to_progress(progress=progress, outcome=outcome) + await _run_depth_pass( + db_session, + scholars=scholars, + first_pass_cstarts=first_pass_cstarts, + pagination=pagination, + run=run, + user_id=user_id, + scholar_kwargs=scholar_kwargs, + request_delay_seconds=request_delay_seconds, + remaining_max=remaining_max, + auto_queue_continuations=auto_queue_continuations, + queue_delay_seconds=queue_delay_seconds, + progress=progress, + ) return progress