from __future__ import annotations import asyncio import logging from datetime import UTC, datetime from typing import Any from sqlalchemy import select, text from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from app.db.models import ( CrawlRun, RunStatus, RunTriggerType, ScholarProfile, ) from app.logging_utils import structured_log from app.services.ingestion import queue as queue_service from app.services.ingestion import safety as run_safety_service from app.services.ingestion.constants import RUN_LOCK_NAMESPACE from app.services.ingestion.enrichment import EnrichmentRunner from app.services.ingestion.pagination import PaginationEngine from app.services.ingestion.preflight import check_scholar_reachable from app.services.ingestion.run_completion import ( complete_run_for_user, int_or_default, run_execution_summary, ) from app.services.ingestion.scholar_processing import run_scholar_iteration from app.services.ingestion.types import ( RunAlertSummary, RunAlreadyInProgressError, RunBlockedBySafetyPolicyError, RunFailureSummary, RunProgress, ) from app.services.scholar.source import ScholarSource from app.services.settings import application as user_settings_service from app.settings import settings logger = logging.getLogger(__name__) ACTIVE_RUN_INDEX_NAME = "uq_crawl_runs_user_active" _background_tasks: set[asyncio.Task[Any]] = set() def _is_active_run_integrity_error(exc: IntegrityError) -> bool: original_error = getattr(exc, "orig", None) if ACTIVE_RUN_INDEX_NAME in str(exc): return True if original_error is None: return False if ACTIVE_RUN_INDEX_NAME in str(original_error): return True diagnostics = getattr(original_error, "diag", None) if diagnostics is None: return False return getattr(diagnostics, "constraint_name", None) == ACTIVE_RUN_INDEX_NAME def _resolve_paging_kwargs( *, request_delay_seconds: int, network_error_retries: int, retry_backoff_seconds: float, rate_limit_retries: int | None, rate_limit_backoff_seconds: float | None, max_pages_per_scholar: int, page_size: int, ) -> dict[str, Any]: return { "request_delay_seconds": request_delay_seconds, "network_error_retries": network_error_retries, "retry_backoff_seconds": retry_backoff_seconds, "rate_limit_retries": rate_limit_retries if rate_limit_retries is not None else settings.ingestion_rate_limit_retries, "rate_limit_backoff_seconds": rate_limit_backoff_seconds if rate_limit_backoff_seconds is not None else settings.ingestion_rate_limit_backoff_seconds, "max_pages_per_scholar": max_pages_per_scholar, "page_size": page_size, } def _threshold_kwargs( *, alert_blocked_failure_threshold: int, alert_network_failure_threshold: int, alert_retry_scheduled_threshold: int, ) -> dict[str, Any]: return { "alert_blocked_failure_threshold": alert_blocked_failure_threshold, "alert_network_failure_threshold": alert_network_failure_threshold, "alert_retry_scheduled_threshold": alert_retry_scheduled_threshold, } def _log_run_completed( *, run: CrawlRun, user_id: int, scholars: list[ScholarProfile], progress: RunProgress, failure_summary: RunFailureSummary, alert_summary: RunAlertSummary, ) -> None: structured_log( logger, "info", "ingestion.run_completed", user_id=user_id, crawl_run_id=run.id, status=run.status.value, scholar_count=len(scholars), succeeded_count=progress.succeeded_count, failed_count=progress.failed_count, partial_count=progress.partial_count, new_publication_count=run.new_pub_count, blocked_failure_count=alert_summary.blocked_failure_count, network_failure_count=alert_summary.network_failure_count, retries_scheduled_count=failure_summary.retries_scheduled_count, alert_flags=alert_summary.alert_flags, ) class ScholarIngestionService: def __init__(self, *, source: ScholarSource) -> None: self._source = source self._pagination = PaginationEngine(source=source) self._enrichment = EnrichmentRunner() @staticmethod def _effective_request_delay_seconds(value: int) -> int: policy_minimum = user_settings_service.resolve_request_delay_minimum( settings.ingestion_min_request_delay_seconds ) return max(policy_minimum, int_or_default(value, policy_minimum)) async def _load_user_settings_for_run( self, db_session: AsyncSession, *, user_id: int, trigger_type: RunTriggerType, ): user_settings = await user_settings_service.get_or_create_settings(db_session, user_id=user_id) await self._enforce_safety_gate( db_session, user_settings=user_settings, user_id=user_id, trigger_type=trigger_type ) return user_settings async def _enforce_safety_gate( self, db_session: AsyncSession, *, user_settings, user_id: int, trigger_type: RunTriggerType, ) -> None: now_utc = datetime.now(UTC) previous = run_safety_service.get_safety_event_context(user_settings, now_utc=now_utc) if run_safety_service.clear_expired_cooldown(user_settings, now_utc=now_utc): await db_session.commit() await db_session.refresh(user_settings) structured_log( logger, "info", "ingestion.cooldown_cleared", user_id=user_id, reason=previous.get("cooldown_reason"), cooldown_until=previous.get("cooldown_until"), ) now_utc = datetime.now(UTC) if run_safety_service.is_cooldown_active(user_settings, now_utc=now_utc): await self._raise_safety_blocked_start( db_session, user_settings=user_settings, user_id=user_id, trigger_type=trigger_type, now_utc=now_utc, ) async def _raise_safety_blocked_start( self, db_session: AsyncSession, *, user_settings, user_id: int, trigger_type: RunTriggerType, now_utc: datetime, ) -> None: safety_state = run_safety_service.register_cooldown_blocked_start(user_settings, now_utc=now_utc) await db_session.commit() structured_log( logger, "warning", "ingestion.safety_policy_blocked_run_start", user_id=user_id, trigger_type=trigger_type.value, reason=safety_state.get("cooldown_reason"), cooldown_until=safety_state.get("cooldown_until"), cooldown_remaining_seconds=safety_state.get("cooldown_remaining_seconds"), blocked_start_count=((safety_state.get("counters") or {}).get("blocked_start_count")), ) raise RunBlockedBySafetyPolicyError( code="scrape_cooldown_active", message="Scrape safety cooldown is active; run start is temporarily blocked.", safety_state=safety_state, ) async def _load_target_scholars( self, db_session: AsyncSession, *, user_id: int, filtered_scholar_ids: set[int] | None, ) -> list[ScholarProfile]: scholars_stmt = ( select(ScholarProfile) .where(ScholarProfile.user_id == user_id, ScholarProfile.is_enabled.is_(True)) .order_by(ScholarProfile.created_at.asc(), ScholarProfile.id.asc()) ) if filtered_scholar_ids is not None: scholars_stmt = scholars_stmt.where(ScholarProfile.id.in_(filtered_scholar_ids)) scholars_result = await db_session.execute(scholars_stmt) scholars = list(scholars_result.scalars().all()) if filtered_scholar_ids is not None: found_ids = {int(s.id) for s in scholars} for sid in filtered_scholar_ids - found_ids: await queue_service.clear_job_for_scholar(db_session, user_id=user_id, scholar_profile_id=sid) return scholars async def _run_preflight_guard( self, db_session: AsyncSession, *, user_settings: Any, user_id: int, scholars: list[ScholarProfile], ) -> None: if not scholars: return result = await check_scholar_reachable( self._source, scholar_id=scholars[0].scholar_id, ) if result.passed: return now_utc = datetime.now(UTC) safety_state, _ = run_safety_service.apply_run_safety_outcome( user_settings, run_id=0, blocked_failure_count=1, network_failure_count=0, blocked_failure_threshold=1, network_failure_threshold=1, blocked_cooldown_seconds=settings.ingestion_safety_cooldown_blocked_seconds, network_cooldown_seconds=settings.ingestion_safety_cooldown_network_seconds, now_utc=now_utc, ) await db_session.commit() structured_log( logger, "warning", "ingestion.cooldown_entered_preflight", user_id=user_id, block_reason=result.block_reason, cooldown_until=safety_state.get("cooldown_until"), ) raise RunBlockedBySafetyPolicyError( code="scrape_cooldown_active", message=f"Preflight detected Scholar block ({result.block_reason}). Cooldown activated.", safety_state=safety_state, ) async def _initialize_run_for_user( self, db_session: AsyncSession, *, user_id: int, trigger_type: RunTriggerType, scholar_profile_ids: set[int] | None, start_cstart_by_scholar_id: dict[int, int] | None, paging_kwargs: dict[str, Any], threshold_kwargs: dict[str, Any], idempotency_key: str | None, ) -> tuple[Any, CrawlRun, list[ScholarProfile], dict[int, int]]: user_settings = await self._load_user_settings_for_run( db_session, user_id=user_id, trigger_type=trigger_type, ) if not await self._try_acquire_user_lock(db_session, user_id=user_id): raise RunAlreadyInProgressError(f"Run already in progress for user_id={user_id}.") filtered_scholar_ids = {int(v) for v in scholar_profile_ids} if scholar_profile_ids is not None else None start_cstart_map = {int(k): max(0, int(v)) for k, v in (start_cstart_by_scholar_id or {}).items()} scholars = await self._load_target_scholars( db_session, user_id=user_id, filtered_scholar_ids=filtered_scholar_ids, ) await self._run_preflight_guard( db_session, user_settings=user_settings, user_id=user_id, scholars=scholars, ) run = await self._start_run_record( db_session, user_id=user_id, trigger_type=trigger_type, scholars=scholars, filtered=filtered_scholar_ids is not None, idempotency_key=idempotency_key, paging_kwargs=paging_kwargs, threshold_kwargs=threshold_kwargs, ) return user_settings, run, scholars, start_cstart_map async def _start_run_record( self, db_session: AsyncSession, *, user_id: int, trigger_type: RunTriggerType, scholars: list[ScholarProfile], filtered: bool, idempotency_key: str | None, paging_kwargs: dict[str, Any], threshold_kwargs: dict[str, Any], ) -> CrawlRun: structured_log( logger, "info", "ingestion.run_started", user_id=user_id, trigger_type=trigger_type.value, scholar_count=len(scholars), is_filtered_run=filtered, idempotency_key=idempotency_key, **paging_kwargs, **threshold_kwargs, ) run = CrawlRun( user_id=user_id, trigger_type=trigger_type, status=RunStatus.RUNNING, scholar_count=len(scholars), new_pub_count=0, idempotency_key=idempotency_key, error_log={}, ) db_session.add(run) try: await db_session.flush() except IntegrityError as exc: if _is_active_run_integrity_error(exc): await db_session.rollback() raise RunAlreadyInProgressError(f"Run already in progress for user_id={user_id}.") from exc raise return run async def initialize_run( self, db_session: AsyncSession, *, user_id: int, trigger_type: RunTriggerType, request_delay_seconds: int, network_error_retries: int = 1, retry_backoff_seconds: float = 1.0, rate_limit_retries: int | None = None, rate_limit_backoff_seconds: float | None = None, max_pages_per_scholar: int = 30, page_size: int = 100, scholar_profile_ids: set[int] | None = None, start_cstart_by_scholar_id: dict[int, int] | None = None, idempotency_key: str | None = None, alert_blocked_failure_threshold: int = 1, alert_network_failure_threshold: int = 2, alert_retry_scheduled_threshold: int = 3, ) -> tuple[CrawlRun, list[ScholarProfile], dict[int, int]]: effective_delay = self._effective_request_delay_seconds(request_delay_seconds) if effective_delay != int_or_default(request_delay_seconds, effective_delay): structured_log( logger, "warning", "ingestion.delay_coerced", user_id=user_id, requested_request_delay_seconds=int_or_default(request_delay_seconds, 0), effective_request_delay_seconds=effective_delay, policy_minimum_request_delay_seconds=user_settings_service.resolve_request_delay_minimum( settings.ingestion_min_request_delay_seconds ), ) paging = _resolve_paging_kwargs( request_delay_seconds=effective_delay, 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, ) thresholds = _threshold_kwargs( alert_blocked_failure_threshold=alert_blocked_failure_threshold, alert_network_failure_threshold=alert_network_failure_threshold, alert_retry_scheduled_threshold=alert_retry_scheduled_threshold, ) _, run, scholars, start_cstart_map = await self._initialize_run_for_user( db_session, user_id=user_id, trigger_type=trigger_type, scholar_profile_ids=scholar_profile_ids, start_cstart_by_scholar_id=start_cstart_by_scholar_id, idempotency_key=idempotency_key, paging_kwargs=paging, threshold_kwargs=thresholds, ) return run, scholars, start_cstart_map async def execute_run( self, session_factory: Any, *, run_id: int, user_id: int, scholars: list[ScholarProfile], start_cstart_map: dict[int, int], request_delay_seconds: int, network_error_retries: int = 1, retry_backoff_seconds: float = 1.0, rate_limit_retries: int | None = None, rate_limit_backoff_seconds: float | None = None, max_pages_per_scholar: int = 30, page_size: int = 100, auto_queue_continuations: bool = True, queue_delay_seconds: int = 60, alert_blocked_failure_threshold: int = 1, alert_network_failure_threshold: int = 2, alert_retry_scheduled_threshold: int = 3, idempotency_key: str | None = None, ) -> None: paging = _resolve_paging_kwargs( 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, ) thresholds = _threshold_kwargs( alert_blocked_failure_threshold=alert_blocked_failure_threshold, alert_network_failure_threshold=alert_network_failure_threshold, alert_retry_scheduled_threshold=alert_retry_scheduled_threshold, ) async with session_factory() as db_session: try: run, user_settings, attached_scholars = await self._prepare_execute_run( db_session, run_id=run_id, user_id=user_id, scholars=scholars ) progress = await run_scholar_iteration( db_session, pagination=self._pagination, run=run, scholars=attached_scholars, user_id=user_id, start_cstart_map=start_cstart_map, auto_queue_continuations=auto_queue_continuations, queue_delay_seconds=queue_delay_seconds, **paging, ) failure_summary, alert_summary = complete_run_for_user( user_settings=user_settings, run=run, scholars=attached_scholars, user_id=user_id, progress=progress, idempotency_key=idempotency_key, **thresholds, ) intended_final_status = run.status if intended_final_status not in (RunStatus.CANCELED,): run.status = RunStatus.RESOLVING await db_session.commit() _log_run_completed( run=run, user_id=user_id, scholars=attached_scholars, progress=progress, failure_summary=failure_summary, alert_summary=alert_summary, ) if intended_final_status not in (RunStatus.CANCELED,): task = asyncio.create_task( self._background_enrich( session_factory, run_id=run.id, intended_final_status=intended_final_status, openalex_api_key=getattr(user_settings, "openalex_api_key", None), ) ) _background_tasks.add(task) task.add_done_callback(_background_tasks.discard) except Exception as exc: await db_session.rollback() logger.exception("ingestion.background_run_failed", extra={"run_id": run_id, "user_id": user_id}) await self._fail_run_in_background(session_factory, run_id, exc) async def _prepare_execute_run( self, db_session: AsyncSession, *, run_id: int, user_id: int, scholars: list[ScholarProfile], ) -> tuple[CrawlRun, Any, list[ScholarProfile]]: run_result = await db_session.execute(select(CrawlRun).where(CrawlRun.id == run_id)) run = run_result.scalar_one() user_settings = await user_settings_service.get_or_create_settings(db_session, user_id=user_id) scholar_ids = [s.id for s in scholars] scholars_result = await db_session.execute( select(ScholarProfile) .where(ScholarProfile.id.in_(scholar_ids)) .order_by(ScholarProfile.created_at.asc(), ScholarProfile.id.asc()) ) return run, user_settings, list(scholars_result.scalars().all()) async def _fail_run_in_background(self, session_factory: Any, run_id: int, exc: Exception) -> None: async with session_factory() as cleanup_session: run_to_fail = await cleanup_session.get(CrawlRun, run_id) if run_to_fail: run_to_fail.status = RunStatus.FAILED run_to_fail.end_dt = datetime.now(UTC) run_to_fail.error_log["terminal_exception"] = str(exc) await cleanup_session.commit() async def _background_enrich( self, session_factory: Any, *, run_id: int, intended_final_status: RunStatus, openalex_api_key: str | None = None, ) -> None: try: async with session_factory() as session: await self._enrichment.enrich_pending_publications( session, run_id=run_id, openalex_api_key=openalex_api_key, ) run = await session.get(CrawlRun, run_id) if run is not None and run.status == RunStatus.RESOLVING: run.status = intended_final_status await session.commit() logger.info( "ingestion.background_enrichment_complete", extra={"run_id": run_id, "final_status": str(intended_final_status)}, ) except Exception: logger.exception("ingestion.background_enrichment_failed", extra={"run_id": run_id}) try: async with session_factory() as fallback_session: run = await fallback_session.get(CrawlRun, run_id) if run is not None and run.status == RunStatus.RESOLVING: run.status = intended_final_status await fallback_session.commit() except Exception: logger.exception("ingestion.background_enrichment_fallback_failed", extra={"run_id": run_id}) async def run_for_user( self, db_session: AsyncSession, *, user_id: int, trigger_type: RunTriggerType, request_delay_seconds: int, network_error_retries: int = 1, retry_backoff_seconds: float = 1.0, rate_limit_retries: int | None = None, rate_limit_backoff_seconds: float | None = None, max_pages_per_scholar: int = 30, page_size: int = 100, scholar_profile_ids: set[int] | None = None, start_cstart_by_scholar_id: dict[int, int] | None = None, auto_queue_continuations: bool = True, queue_delay_seconds: int = 60, idempotency_key: str | None = None, alert_blocked_failure_threshold: int = 1, alert_network_failure_threshold: int = 2, alert_retry_scheduled_threshold: int = 3, ): run, scholars, start_cstart_map = await self.initialize_run( db_session, user_id=user_id, trigger_type=trigger_type, 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, scholar_profile_ids=scholar_profile_ids, start_cstart_by_scholar_id=start_cstart_by_scholar_id, idempotency_key=idempotency_key, alert_blocked_failure_threshold=alert_blocked_failure_threshold, alert_network_failure_threshold=alert_network_failure_threshold, alert_retry_scheduled_threshold=alert_retry_scheduled_threshold, ) paging = _resolve_paging_kwargs( request_delay_seconds=self._effective_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, ) thresholds = _threshold_kwargs( alert_blocked_failure_threshold=alert_blocked_failure_threshold, 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, run=run, scholars=scholars, user_id=user_id, start_cstart_map=start_cstart_map, auto_queue_continuations=auto_queue_continuations, queue_delay_seconds=queue_delay_seconds, **paging, ) user_settings = await user_settings_service.get_or_create_settings(db_session, user_id=user_id) failure_summary, alert_summary = complete_run_for_user( user_settings=user_settings, run=run, scholars=scholars, user_id=user_id, progress=progress, idempotency_key=idempotency_key, **thresholds, ) intended_final_status = run.status 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, run_id=run.id, openalex_api_key=getattr(user_settings, "openalex_api_key", None), ) except Exception: logger.exception("ingestion.enrichment_failed", extra={"run_id": run.id}) if run.status == RunStatus.RESOLVING: run.status = intended_final_status await db_session.commit() async def _try_acquire_user_lock( self, db_session: AsyncSession, *, user_id: int, ) -> bool: result = await db_session.execute( text("SELECT pg_try_advisory_xact_lock(:namespace, :user_key)"), {"namespace": RUN_LOCK_NAMESPACE, "user_key": int(user_id)}, ) return bool(result.scalar_one())