FabledCurator/backend/app/celery_signals.py

"""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.

    Audit 2026-06-02: backup/admin/library_audit prefixes were
    missing here even though task_routes sent all three to
    'maintenance'. The TaskRun.queue column then lied for those
    rows (claimed 'default') so per-queue dashboard filters and
    per-queue threshold overrides silently missed them.
    """
    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.backup.",
        "backend.app.tasks.admin.",
        "backend.app.tasks.library_audit.",
    )):
        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),
    )