"""Task Scheduler for running collection jobs.""" import asyncio from datetime import UTC, datetime, timedelta from typing import Any, Dict, Optional from apscheduler.schedulers.asyncio import AsyncIOScheduler from apscheduler.triggers.interval import IntervalTrigger from sqlalchemy import select from app.core.enums import JobStatus from app.core.logging import get_logger from app.db.session import async_session_factory from app.core.time import to_iso8601_utc from app.models.datasource import DataSource from app.models.task import CollectionTask from app.services.business_logs import emit_business_log, emit_business_log_background, exception_context from app.services.collectors.registry import collector_registry from app.services.datasource_connectivity import ( build_builtin_connectivity_checksum, get_builtin_effective_candidate, save_connectivity_success, ) logger = get_logger(__name__) scheduler = AsyncIOScheduler() RUNNING_TASK_GUARD_TIMEOUT_MINUTES = 90 RUNNING_COLLECTOR_TASKS: dict[str, asyncio.Task[Any]] = {} def _collector_task_name(collector_name: str) -> str: return f"collector:{collector_name}" def get_running_collector_task(collector_name: str) -> asyncio.Task[Any] | None: task = RUNNING_COLLECTOR_TASKS.get(collector_name) if task is not None and not task.done(): return task if task is not None and task.done(): RUNNING_COLLECTOR_TASKS.pop(collector_name, None) target_name = _collector_task_name(collector_name) for candidate in asyncio.all_tasks(): if candidate.done(): continue if candidate.get_name() == target_name: RUNNING_COLLECTOR_TASKS[collector_name] = candidate return candidate return None async def _update_next_run_at(datasource: DataSource, session) -> None: job = scheduler.get_job(datasource.source) datasource.next_run_at = job.next_run_time if job else None await session.commit() async def _apply_datasource_schedule(datasource: DataSource, session) -> None: collector = collector_registry.get(datasource.source) if not collector: logger.warning_event( "Collector not found for datasource", event="collector.schedule.collector_missing", context={"collector_name": datasource.source}, ) return collector_registry.set_active(datasource.source, datasource.is_active) existing_job = scheduler.get_job(datasource.source) if existing_job: scheduler.remove_job(datasource.source) if datasource.is_active: scheduler.add_job( run_collector_task, trigger=IntervalTrigger(minutes=max(1, datasource.frequency_minutes)), id=datasource.source, name=datasource.name, replace_existing=True, kwargs={"collector_name": datasource.source}, ) logger.info_event( "Scheduled collector", event="collector.schedule.updated", context={"collector_name": datasource.source, "frequency_minutes": datasource.frequency_minutes}, ) else: logger.info_event( "Collector disabled", event="collector.schedule.disabled", context={"collector_name": datasource.source}, ) await _update_next_run_at(datasource, session) async def run_collector_task(collector_name: str): """Run a single collector task.""" collector = collector_registry.get(collector_name) if not collector: logger.error_event( "Collector not found", event="collector.run.collector_missing", context={"collector_name": collector_name}, ) return async with async_session_factory() as db: result = await db.execute(select(DataSource).where(DataSource.source == collector_name)) datasource = result.scalar_one_or_none() if not datasource: logger.error_event( "Datasource not found for collector", event="collector.run.datasource_missing", context={"collector_name": collector_name}, ) return if not datasource.is_active: logger.info_event( "Skipping disabled collector", event="collector.run.skipped_disabled", context={"collector_name": collector_name}, ) await emit_business_log( logger, event="collector.run.skipped_disabled", message="Skipping disabled collector", category="collector", service="scheduler", module=__name__, context={"collector_name": collector_name, "datasource_id": datasource.id, "status": "skipped"}, ) return running_result = await db.execute( select(CollectionTask) .where( CollectionTask.datasource_id == datasource.id, CollectionTask.status == JobStatus.RUNNING.value, ) .order_by(CollectionTask.started_at.desc(), CollectionTask.id.desc()) .limit(1) ) existing_running = running_result.scalar_one_or_none() if existing_running is not None: now = datetime.now(UTC) started_at = existing_running.started_at if started_at is not None and started_at.tzinfo is None: started_at = started_at.replace(tzinfo=UTC) is_stale = ( started_at is not None and (now - started_at) > timedelta(minutes=RUNNING_TASK_GUARD_TIMEOUT_MINUTES) ) if not is_stale: logger.warning_event( "Skipping collector trigger because task is already running", event="collector.run.skipped_already_running", context={"collector_name": collector_name, "task_id": existing_running.id}, ) await emit_business_log( logger, event="collector.run.skipped_already_running", message="Skipping collector trigger because task is already running", category="collector", level="warning", service="scheduler", module=__name__, context={ "collector_name": collector_name, "datasource_id": datasource.id, "task_id": existing_running.id, "status": "skipped", }, ) return existing_error = (existing_running.error_message or "").strip() stale_reason = ( f"Marked failed automatically after stale running timeout " f"({RUNNING_TASK_GUARD_TIMEOUT_MINUTES}m) in scheduler guard" ) existing_running.status = JobStatus.FAILED.value existing_running.phase = "failed" existing_running.completed_at = now existing_running.error_message = ( f"{existing_error}\n{stale_reason}".strip() if existing_error else stale_reason ) await db.commit() logger.warning_event( "Marked stale running task as failed before rerun", event="collector.run.stale_task_failed", context={"collector_name": collector_name, "task_id": existing_running.id}, ) await emit_business_log( logger, event="collector.run.stale_task_failed", message="Marked stale running task as failed before rerun", category="collector", level="warning", service="scheduler", module=__name__, context={ "collector_name": collector_name, "datasource_id": datasource.id, "task_id": existing_running.id, "status": "failed", }, ) try: datasource_id = datasource.id datasource_source = datasource.source collector._datasource_id = datasource_id logger.info_event( "Running collector", event="collector.run.started", context={"collector_name": collector_name, "datasource_id": datasource_id}, ) await emit_business_log( logger, event="collector.run.scheduled_started", message="Scheduler started collector run", category="collector", service="scheduler", module=__name__, context={"collector_name": collector_name, "datasource_id": datasource_id, "status": "running"}, ) task_result = await collector.run(db) datasource = await db.get(DataSource, datasource_id) if datasource is None: logger.error_event( "Datasource disappeared after collector run", event="collector.run.datasource_missing_after_run", context={"collector_name": collector_name, "datasource_id": datasource_id}, ) return datasource.last_run_at = datetime.now(UTC) datasource.last_status = task_result.get("status") if datasource.last_status == JobStatus.SUCCESS.value: effective_candidate = await get_builtin_effective_candidate(db, datasource_source) checksum, _credential_context = await build_builtin_connectivity_checksum( datasource_source, effective_candidate["endpoint"], effective_candidate["auth_type"], effective_candidate["headers"], effective_candidate["config"], db, ) await save_connectivity_success( db, datasource_source, checksum, {"status_code": None}, connected_by="collection", ) await _update_next_run_at(datasource, db) logger.info_event( "Collector completed", event="collector.run.completed", context={"collector_name": collector_name, "datasource_id": datasource_id, "result": task_result}, ) await emit_business_log( logger, event="collector.run.scheduled_completed", message="Scheduler completed collector run", category="collector", service="scheduler", module=__name__, context={ "collector_name": collector_name, "datasource_id": datasource_id, "status": task_result.get("status"), "result": task_result, }, ) except asyncio.CancelledError: await db.rollback() datasource = await db.get(DataSource, datasource_id) datasource.last_run_at = datetime.now(UTC) datasource.last_status = JobStatus.CANCELLED.value await db.commit() logger.warning_event( "Collector cancelled by operator", event="collector.run.cancelled", context={"collector_name": collector_name, "datasource_id": datasource.id}, ) await emit_business_log( logger, event="collector.run.cancelled", message="Collector cancelled by operator", category="collector", level="warning", service="scheduler", module=__name__, context={"collector_name": collector_name, "datasource_id": datasource.id, "status": "cancelled"}, ) raise except Exception as exc: await db.rollback() datasource = await db.get(DataSource, datasource_id) datasource.last_run_at = datetime.now(UTC) datasource.last_status = JobStatus.FAILED.value await db.commit() logger.exception_event( "Collector failed", event="collector.run.failed", context={"collector_name": collector_name, "datasource_id": datasource.id, "error": str(exc)}, ) await emit_business_log( logger, event="collector.run.failed", message="Collector failed", category="collector", level="error", service="scheduler", module=__name__, context=exception_context( exc, {"collector_name": collector_name, "datasource_id": datasource.id, "status": "failed"}, ), ) async def cleanup_stale_running_tasks(max_age_hours: int = 2) -> int: """Mark stale running tasks as failed after restarts or collector hangs.""" cutoff = datetime.now(UTC) - timedelta(hours=max_age_hours) async with async_session_factory() as db: result = await db.execute( select(CollectionTask).where( CollectionTask.status == JobStatus.RUNNING.value, CollectionTask.started_at.is_not(None), CollectionTask.started_at < cutoff, ) ) stale_tasks = result.scalars().all() for task in stale_tasks: task.status = JobStatus.FAILED.value task.phase = "failed" task.completed_at = datetime.now(UTC) existing_error = (task.error_message or "").strip() cleanup_error = "Marked failed automatically after stale running task cleanup" task.error_message = f"{existing_error}\n{cleanup_error}".strip() if existing_error else cleanup_error if stale_tasks: await db.commit() logger.warning_event( "Cleaned up stale running collection tasks", event="collector.cleanup.stale_tasks_cleaned", context={"count": len(stale_tasks)}, ) return len(stale_tasks) def start_scheduler() -> None: """Start the scheduler.""" if not scheduler.running: scheduler.start() logger.info_event("Scheduler started", event="scheduler.started") def stop_scheduler() -> None: """Stop the scheduler.""" if scheduler.running: scheduler.shutdown(wait=False) logger.info_event("Scheduler stopped", event="scheduler.stopped") async def sync_scheduler_with_datasources() -> None: """Synchronize scheduler jobs with datasource table.""" async with async_session_factory() as db: result = await db.execute(select(DataSource).order_by(DataSource.id)) datasources = result.scalars().all() configured_sources = {datasource.source for datasource in datasources} for job in list(scheduler.get_jobs()): if job.id not in configured_sources: scheduler.remove_job(job.id) for datasource in datasources: await _apply_datasource_schedule(datasource, db) async def sync_datasource_job(datasource_id: int) -> bool: """Synchronize a single datasource job after settings changes.""" async with async_session_factory() as db: datasource = await db.get(DataSource, datasource_id) if not datasource: return False await _apply_datasource_schedule(datasource, db) return True def get_scheduler_jobs() -> list[Dict[str, Any]]: """Get all scheduled jobs.""" jobs = [] for job in scheduler.get_jobs(): jobs.append( { "id": job.id, "name": job.name, "next_run_time": to_iso8601_utc(job.next_run_time), "trigger": str(job.trigger), } ) return jobs async def get_latest_task_id_for_datasource(datasource_id: int) -> Optional[int]: from app.models.task import CollectionTask async with async_session_factory() as db: result = await db.execute( select(CollectionTask.id) .where(CollectionTask.datasource_id == datasource_id) .order_by(CollectionTask.created_at.desc(), CollectionTask.id.desc()) .limit(1) ) return result.scalar_one_or_none() def run_collector_now(collector_name: str) -> bool: """Run a collector immediately (not scheduled).""" collector = collector_registry.get(collector_name) if not collector: logger.error_event( "Collector not found", event="collector.trigger.collector_missing", context={"collector_name": collector_name}, ) return False existing_task = get_running_collector_task(collector_name) if existing_task is not None and not existing_task.done(): logger.warning_event( "Collector is already running in-memory; skipping duplicate trigger", event="collector.trigger.skipped_already_running", context={"collector_name": collector_name}, ) emit_business_log_background( logger, event="collector.trigger.skipped_already_running", message="Collector is already running in-memory; skipping duplicate trigger", category="collector", level="warning", service="scheduler", module=__name__, context={"collector_name": collector_name, "status": "skipped"}, ) return False try: task = asyncio.create_task(run_collector_task(collector_name), name=_collector_task_name(collector_name)) RUNNING_COLLECTOR_TASKS[collector_name] = task def _cleanup_task(done_task: asyncio.Task[Any]) -> None: current = RUNNING_COLLECTOR_TASKS.get(collector_name) if current is done_task: RUNNING_COLLECTOR_TASKS.pop(collector_name, None) task.add_done_callback(_cleanup_task) logger.info_event( "Triggered collector", event="collector.trigger.started", context={"collector_name": collector_name}, ) emit_business_log_background( logger, event="collector.trigger.started", message="Triggered collector", category="collector", service="scheduler", module=__name__, context={"collector_name": collector_name, "status": "queued"}, ) return True except Exception as exc: logger.error_event( "Failed to trigger collector", event="collector.trigger.failed", context={"collector_name": collector_name, "error": str(exc)}, ) emit_business_log_background( logger, event="collector.trigger.failed", message="Failed to trigger collector", category="collector", level="error", service="scheduler", module=__name__, context=exception_context(exc, {"collector_name": collector_name, "status": "failed"}), ) return False def is_collector_running(collector_name: str) -> bool: task = get_running_collector_task(collector_name) return bool(task is not None and not task.done()) async def cancel_running_collector_now(collector_name: str) -> bool: task = get_running_collector_task(collector_name) if task is None or task.done(): RUNNING_COLLECTOR_TASKS.pop(collector_name, None) return False task.cancel() try: await task except asyncio.CancelledError: return True return task.cancelled()