"""FC-3i: task_run lifecycle via Celery signals. Subscribes to task_prerun / task_postrun / task_failure / task_retry and persists one task_run row per task attempt. Drop-in for every existing and future Celery task — no per-task instrumentation. Signal handlers run inside the worker process (sync context); DB writes go through the existing shared sync engine (backend.app.tasks._sync_engine.sync_session_factory) — one engine per worker process, not per-task, so we don't blow Postgres max_connections under load (the reason FC-3g shared-engine fix existed). Failure-mode discipline (operator-pressed point): every handler is wrapped in try/except that swallows + logs. If the DB is down or the handler has a bug, the real task still runs — the dashboard goes dark for that interval. Monitoring NEVER breaks the thing it's monitoring. """ import logging from datetime import UTC, datetime from celery.exceptions import SoftTimeLimitExceeded from celery.signals import task_failure, task_postrun, task_prerun, task_retry from .models import TaskRun from .tasks._sync_engine import sync_session_factory log = logging.getLogger(__name__) # Celery-internal tasks that would generate dashboard noise without # operational value. Conservative list; extend only when a specific # task proves noisy. _UNTRACKED_TASK_NAMES = frozenset({ "celery.chord_unlock", "celery.backend_cleanup", "celery.chunks", }) _MAX_ERROR_MESSAGE_LEN = 2000 _MAX_ARGS_SUMMARY_LEN = 255 _MAX_WORKER_HOSTNAME_LEN = 128 # PostgreSQL Integer is signed 32-bit. Tasks called with a first-arg # int outside this range (e.g. an absurdly large mock value, or a # string-of-digits coercible to int but bigger than 2^31-1) would crash # the INSERT with NumericValueOutOfRange. Bound the recorded value to # the column's range; values outside become None (target_id is # nullable, so this is safe). _INT32_MAX = 2_147_483_647 _INT32_MIN = -2_147_483_648 def _queue_for(task) -> str: """Reverse the task→queue routing from celery_app.task_routes. Keep in sync if task_routes is reordered.""" name = getattr(task, "name", "") or "" if name.startswith("backend.app.tasks.import_file."): return "import" if name.startswith("backend.app.tasks.ml."): return "ml" if name.startswith("backend.app.tasks.thumbnail."): return "thumbnail" if name.startswith("backend.app.tasks.download."): return "download" if name.startswith("backend.app.tasks.scan."): return "scan" if name.startswith(( "backend.app.tasks.maintenance.", "backend.app.tasks.migration.", )): return "maintenance" return "default" def _target_id_from_args(args) -> int | None: """Best-effort: if the first positional arg parses as int AND fits in the column's signed-32-bit range, record it as target_id (image_id, source_id, etc.). Never raises.""" if not args: return None try: value = int(args[0]) except (TypeError, ValueError): return None if value < _INT32_MIN or value > _INT32_MAX: return None return value def _truncate(s, limit: int) -> str | None: if s is None: return None text = str(s) return text if len(text) <= limit else text[:limit] def _is_tracked(task_name: str | None) -> bool: return bool(task_name) and task_name not in _UNTRACKED_TASK_NAMES @task_prerun.connect def _on_prerun(sender=None, task_id=None, task=None, args=None, kwargs=None, **_): if not _is_tracked(getattr(task, "name", None)): return try: Session = sync_session_factory() with Session() as session: session.add(TaskRun( celery_task_id=task_id or "", queue=_queue_for(task), task_name=task.name, target_id=_target_id_from_args(args), started_at=datetime.now(UTC), status="running", args_summary=_truncate(repr(args), _MAX_ARGS_SUMMARY_LEN), worker_hostname=_truncate( getattr(sender, "hostname", None), _MAX_WORKER_HOSTNAME_LEN, ), )) session.commit() except Exception: # noqa: BLE001 — never break the worker log.exception("task_run prerun insert failed (task=%s)", getattr(task, "name", "?")) def _finalize(task_id: str, *, status: str, error_type: str | None = None, error_message: str | None = None, retry_count: int | None = None) -> None: """Shared write path for postrun/failure/retry. Picks the most- recent task_run row for this celery_task_id that's still 'running' (retries reuse the same celery_task_id; each new attempt's prerun inserts a fresh row, so finalize targets the latest running row).""" try: from sqlalchemy import select Session = sync_session_factory() now = datetime.now(UTC) with Session() as session: row = session.execute( select(TaskRun) .where(TaskRun.celery_task_id == task_id) .where(TaskRun.status == "running") .order_by(TaskRun.id.desc()) .limit(1) ).scalar_one_or_none() if row is None: return # no prerun row (untracked, insert failed, or already finalized) row.finished_at = now row.duration_ms = int( (now - row.started_at).total_seconds() * 1000 ) row.status = status if error_type is not None: row.error_type = _truncate(error_type, 128) if error_message is not None: row.error_message = _truncate(error_message, _MAX_ERROR_MESSAGE_LEN) if retry_count is not None: row.retry_count = retry_count session.commit() except Exception: # noqa: BLE001 log.exception("task_run finalize failed (task_id=%s)", task_id) @task_postrun.connect def _on_postrun(sender=None, task_id=None, task=None, args=None, kwargs=None, retval=None, state=None, **_): if not _is_tracked(getattr(task, "name", None)): return # state is one of SUCCESS/FAILURE/RETRY/etc. Only handle SUCCESS; # task_failure handles FAILURE explicitly (with the exception). if state != "SUCCESS": return _finalize(task_id, status="ok") @task_failure.connect def _on_failure(sender=None, task_id=None, exception=None, args=None, kwargs=None, einfo=None, **_): if not _is_tracked(getattr(sender, "name", None)): return status = ("timeout" if isinstance(exception, SoftTimeLimitExceeded) else "error") _finalize( task_id, status=status, error_type=type(exception).__name__ if exception else "Unknown", error_message=str(exception) if exception else None, ) @task_retry.connect def _on_retry(sender=None, request=None, reason=None, einfo=None, **_): if not _is_tracked(getattr(sender, "name", None)): return task_id = getattr(request, "id", None) if not task_id: return # Mark current attempt's row as 'retry' (terminal for this row). # The next attempt's task_prerun inserts a fresh row. _finalize( task_id, status="retry", error_type=type(reason).__name__ if reason else "Retry", error_message=str(reason) if reason else None, retry_count=getattr(request, "retries", 0), )