from __future__ import annotations import asyncio import logging import random from collections.abc import Callable, Coroutine from datetime import UTC, datetime from typing import Any from sqlalchemy.exc import InterfaceError, OperationalError from sqlalchemy.ext.asyncio import AsyncSession from app.db.models import ( CrawlRun, RunStatus, ScholarProfile, ) from app.logging_utils import structured_log from app.services.ingestion import queue as queue_service from app.services.ingestion.constants import ( RESUMABLE_PARTIAL_REASON_PREFIXES, RESUMABLE_PARTIAL_REASONS, ) from app.services.ingestion.pagination import PaginationEngine from app.services.ingestion.publication_upsert import upsert_profile_publications from app.services.ingestion.run_completion import apply_outcome_to_progress from app.services.ingestion.scholar_outcomes import ( apply_first_page_profile_metadata, assert_valid_paged_parse_result, build_result_entry, skipped_no_change_outcome, unexpected_scholar_exception_outcome, upsert_publications_outcome, ) from app.services.ingestion.types import ( PagedParseResult, RunProgress, ScholarProcessingOutcome, ) from app.services.scholar.parser import ParseState from app.services.scholar.state_detection import is_hard_challenge_reason logger = logging.getLogger(__name__) async def sync_continuation_queue( db_session: AsyncSession, *, user_id: int, scholar: ScholarProfile, run: CrawlRun, start_cstart: int, result_entry: dict[str, Any], paged_parse_result: PagedParseResult, auto_queue_continuations: bool, queue_delay_seconds: int, ) -> None: queue_reason, queue_cstart = resolve_continuation_queue_target( outcome=str(result_entry.get("outcome", "")), state=str(result_entry.get("state", "")), pagination_truncated_reason=paged_parse_result.pagination_truncated_reason, continuation_cstart=paged_parse_result.continuation_cstart, fallback_cstart=start_cstart, ) if auto_queue_continuations and queue_reason is not None: await queue_service.upsert_job( db_session, user_id=user_id, scholar_profile_id=scholar.id, resume_cstart=queue_cstart, reason=queue_reason, run_id=run.id, delay_seconds=queue_delay_seconds, ) result_entry["continuation_enqueued"] = True result_entry["continuation_reason"] = queue_reason result_entry["continuation_cstart"] = queue_cstart return if await queue_service.clear_job_for_scholar(db_session, user_id=user_id, scholar_profile_id=scholar.id): result_entry["continuation_cleared"] = True def resolve_continuation_queue_target( *, outcome: str, state: str, pagination_truncated_reason: str | None, continuation_cstart: int | None, fallback_cstart: int, ) -> tuple[str | None, int]: if outcome == "partial": reason = (pagination_truncated_reason or "").strip() if reason in RESUMABLE_PARTIAL_REASONS or reason.startswith(RESUMABLE_PARTIAL_REASON_PREFIXES): return reason, queue_service.normalize_cstart( continuation_cstart if continuation_cstart is not None else fallback_cstart ) return None, queue_service.normalize_cstart(fallback_cstart) if outcome == "failed" and state == ParseState.NETWORK_ERROR.value: return "network_error_retry", queue_service.normalize_cstart( continuation_cstart if continuation_cstart is not None else fallback_cstart ) return None, queue_service.normalize_cstart(fallback_cstart) async def process_scholar( db_session: AsyncSession, *, pagination: PaginationEngine, run: CrawlRun, scholar: ScholarProfile, user_id: int, request_delay_seconds: int, network_error_retries: int, retry_backoff_seconds: float, rate_limit_retries: int, rate_limit_backoff_seconds: float, max_pages_per_scholar: int, page_size: int, start_cstart: int, auto_queue_continuations: bool, queue_delay_seconds: int, ) -> ScholarProcessingOutcome: try: run_dt, paged_parse_result, result_entry = await _fetch_and_prepare_scholar_result( db_session, pagination=pagination, run=run, scholar=scholar, user_id=user_id, start_cstart=start_cstart, request_delay_seconds=request_delay_seconds, network_error_retries=network_error_retries, retry_backoff_seconds=retry_backoff_seconds, rate_limit_retries=rate_limit_retries, rate_limit_backoff_seconds=rate_limit_backoff_seconds, max_pages_per_scholar=max_pages_per_scholar, page_size=page_size, ) outcome = await _resolve_scholar_outcome( db_session, run=run, scholar=scholar, run_dt=run_dt, paged_parse_result=paged_parse_result, result_entry=result_entry, ) await sync_continuation_queue( db_session, user_id=user_id, scholar=scholar, run=run, start_cstart=start_cstart, result_entry=outcome.result_entry, paged_parse_result=paged_parse_result, auto_queue_continuations=auto_queue_continuations, queue_delay_seconds=queue_delay_seconds, ) return outcome except (OperationalError, InterfaceError): raise except Exception as exc: return unexpected_scholar_exception_outcome( run=run, scholar=scholar, start_cstart=start_cstart, exc=exc, ) async def _fetch_and_prepare_scholar_result( db_session: AsyncSession, *, pagination: PaginationEngine, run: CrawlRun, scholar: ScholarProfile, user_id: int, start_cstart: int, request_delay_seconds: int, network_error_retries: int, retry_backoff_seconds: float, rate_limit_retries: int, rate_limit_backoff_seconds: float, max_pages_per_scholar: int, page_size: int, ) -> tuple[datetime, PagedParseResult, dict[str, Any]]: run_dt = datetime.now(UTC) paged_parse_result = await pagination.fetch_and_parse_all_pages( scholar=scholar, run=run, db_session=db_session, start_cstart=start_cstart, request_delay_seconds=request_delay_seconds, network_error_retries=network_error_retries, retry_backoff_seconds=retry_backoff_seconds, rate_limit_retries=rate_limit_retries, rate_limit_backoff_seconds=rate_limit_backoff_seconds, max_pages=max_pages_per_scholar, page_size=page_size, previous_initial_page_fingerprint_sha256=scholar.last_initial_page_fingerprint_sha256, upsert_publications_fn=upsert_profile_publications, ) assert_valid_paged_parse_result(scholar_id=scholar.scholar_id, paged_parse_result=paged_parse_result) apply_first_page_profile_metadata(scholar=scholar, paged_parse_result=paged_parse_result, run_dt=run_dt) parsed_page = paged_parse_result.parsed_page structured_log( logger, "info", "ingestion.scholar_parsed", user_id=user_id, crawl_run_id=run.id, scholar_profile_id=scholar.id, scholar_id=scholar.scholar_id, state=parsed_page.state.value, publication_count=len(paged_parse_result.publications), has_show_more_button=parsed_page.has_show_more_button, pages_fetched=paged_parse_result.pages_fetched, pages_attempted=paged_parse_result.pages_attempted, has_more_remaining=paged_parse_result.has_more_remaining, pagination_truncated_reason=paged_parse_result.pagination_truncated_reason, warning_count=len(parsed_page.warnings), skipped_no_change=paged_parse_result.skipped_no_change, ) result_entry = build_result_entry( scholar=scholar, start_cstart=start_cstart, paged_parse_result=paged_parse_result, ) return run_dt, paged_parse_result, result_entry async def _resolve_scholar_outcome( db_session: AsyncSession, *, run: CrawlRun, scholar: ScholarProfile, run_dt: datetime, paged_parse_result: PagedParseResult, result_entry: dict[str, Any], ) -> ScholarProcessingOutcome: if paged_parse_result.skipped_no_change: return skipped_no_change_outcome( scholar=scholar, run_dt=run_dt, paged_parse_result=paged_parse_result, result_entry=result_entry, ) return await upsert_publications_outcome( db_session, run=run, scholar=scholar, run_dt=run_dt, paged_parse_result=paged_parse_result, result_entry=result_entry, ) def _is_hard_challenge_outcome(outcome: ScholarProcessingOutcome) -> bool: entry = outcome.result_entry return str(entry.get("state", "")) == ParseState.BLOCKED_OR_CAPTCHA.value and is_hard_challenge_reason( str(entry.get("state_reason", "")) ) 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, on_progress: Callable[[int, int], Coroutine[Any, Any, None]] | None = None, ) -> 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: jitter = random.uniform(0.0, min(float(request_delay_seconds), 2.0)) await asyncio.sleep(float(request_delay_seconds) + jitter) 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) if _is_hard_challenge_outcome(outcome): structured_log( logger, "warning", "ingestion.run_aborted_hard_challenge", run_id=run.id, user_id=user_id, scholar_id=scholar.scholar_id, state_reason=outcome.result_entry.get("state_reason"), scholars_remaining=len(scholars) - index - 1, ) return first_pass_cstarts if on_progress is not None: await on_progress(1, 0) 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, on_progress: Callable[[int, int], Coroutine[Any, Any, None]] | None = None, ) -> 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: jitter = random.uniform(0.0, min(float(request_delay_seconds), 2.0)) await asyncio.sleep(float(request_delay_seconds) + jitter) 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) if _is_hard_challenge_outcome(outcome): structured_log( logger, "warning", "ingestion.run_aborted_hard_challenge", run_id=run.id, user_id=user_id, scholar_id=scholar.scholar_id, state_reason=outcome.result_entry.get("state_reason"), scholars_remaining=len(scholars) - index - 1, ) break if on_progress is not None: await on_progress(0, 1) async def run_scholar_iteration( db_session: AsyncSession, *, pagination: PaginationEngine, run: CrawlRun, scholars: list[ScholarProfile], user_id: int, start_cstart_map: dict[int, int], request_delay_seconds: int, network_error_retries: int, retry_backoff_seconds: float, rate_limit_retries: int, rate_limit_backoff_seconds: float, max_pages_per_scholar: int, page_size: int, auto_queue_continuations: bool, queue_delay_seconds: int, ) -> RunProgress: from app.services.runs.events import run_events progress = RunProgress() scholar_kwargs: dict[str, Any] = { "request_delay_seconds": request_delay_seconds, "network_error_retries": network_error_retries, "retry_backoff_seconds": retry_backoff_seconds, "rate_limit_retries": rate_limit_retries, "rate_limit_backoff_seconds": rate_limit_backoff_seconds, "page_size": page_size, } visited = 0 finished = 0 total = len(scholars) async def _emit(v: int = 0, f: int = 0) -> None: nonlocal visited, finished visited += v finished += f await run_events.publish( run.id, "scholar_progress", {"visited": visited, "finished": finished, "total": total}, ) await _emit() 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, on_progress=_emit, ) remaining_max = max(max_pages_per_scholar - 1, 0) scholars_finished_in_first_pass = len(scholars) - len(first_pass_cstarts) if scholars_finished_in_first_pass > 0: await _emit(f=scholars_finished_in_first_pass) if remaining_max <= 0: return progress 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, on_progress=_emit, ) return progress