FabledCurator/backend/app/tasks/library_audit.py

"""scan_library_for_rule Celery task — iterates image_record in keyset-
paginated batches, evaluates the audit rule per image, populates
LibraryAuditRun.matched_ids. Runs on the maintenance queue with a 2h soft
time limit (plenty of margin for 100k+ image libraries at ~100ms PIL
decode + histogram per image).

State machine:
  start: status='running'
  end success: status='ready'
  end error: status='error', error=traceback
  oversize: status='error', error='matched too many images; tighten threshold'
  external cancel: scan sees status='cancelled' between batches, exits.
"""

import logging
import traceback
from datetime import UTC, datetime

from celery.exceptions import SoftTimeLimitExceeded
from PIL import Image
from sqlalchemy import select, update
from sqlalchemy.exc import DBAPIError, OperationalError

from ..celery_app import celery
from ..models import ImageRecord, LibraryAuditRun
from ..services.audits import single_color, transparency
from ._sync_engine import sync_session_factory as _sync_session_factory

log = logging.getLogger(__name__)

_BATCH = 500
_PROGRESS_TICK = 100
_MAX_MATCHED = 50_000

_RULES = {
    "transparency": transparency.evaluate,
    "single_color": single_color.evaluate,
}


@celery.task(
    name="backend.app.tasks.library_audit.scan_library_for_rule",
    bind=True,
    autoretry_for=(OperationalError, DBAPIError),
    retry_backoff=5,
    retry_backoff_max=60,
    retry_jitter=True,
    max_retries=3,
    soft_time_limit=7200,
    time_limit=7500,
)
def scan_library_for_rule(self, audit_id: int) -> dict:
    """See module docstring. Returns a small summary dict for eager-mode
    test assertions (real workers ignore the return value)."""
    SessionLocal = _sync_session_factory()
    try:
        with SessionLocal() as session:
            audit = session.get(LibraryAuditRun, audit_id)
            if audit is None:
                return {"audit_id": audit_id, "status": "missing"}
            evaluate = _RULES.get(audit.rule)
            if evaluate is None:
                _mark_error(session, audit_id, f"unknown rule {audit.rule!r}")
                return {"audit_id": audit_id, "status": "error"}
            params = dict(audit.params or {})
            matched: list[int] = []
            scanned = 0
            last_id = 0
            while True:
                # Cancellation check between batches.
                current_status = session.execute(
                    select(LibraryAuditRun.status)
                    .where(LibraryAuditRun.id == audit_id)
                ).scalar_one()
                if current_status == "cancelled":
                    return {"audit_id": audit_id, "status": "cancelled"}
                rows = session.execute(
                    select(ImageRecord.id, ImageRecord.path)
                    .where(ImageRecord.id > last_id)
                    .where(ImageRecord.mime.like("image/%"))
                    .order_by(ImageRecord.id.asc())
                    .limit(_BATCH)
                ).all()
                if not rows:
                    break
                for image_id, image_path in rows:
                    last_id = image_id
                    scanned += 1
                    try:
                        with Image.open(image_path) as im:
                            try:
                                if evaluate(im, **params):
                                    matched.append(image_id)
                            except Exception as exc:  # noqa: BLE001
                                log.warning(
                                    "audit %s: rule evaluate failed on %s: %s",
                                    audit_id, image_path, exc,
                                )
                    except FileNotFoundError:
                        log.warning(
                            "audit %s: image_record %s file missing at %s; skipping",
                            audit_id, image_id, image_path,
                        )
                    except OSError as exc:
                        log.warning(
                            "audit %s: PIL load failed for %s: %s",
                            audit_id, image_path, exc,
                        )
                    if len(matched) > _MAX_MATCHED:
                        _mark_error(
                            session, audit_id,
                            f"matched > {_MAX_MATCHED} images; "
                            "tighten threshold and re-run",
                        )
                        return {"audit_id": audit_id, "status": "error"}
                    if scanned % _PROGRESS_TICK == 0:
                        session.execute(
                            update(LibraryAuditRun)
                            .where(LibraryAuditRun.id == audit_id)
                            .values(scanned_count=scanned)
                        )
                        session.commit()
            # Final state.
            session.execute(
                update(LibraryAuditRun)
                .where(LibraryAuditRun.id == audit_id)
                .values(
                    scanned_count=scanned,
                    matched_count=len(matched),
                    matched_ids=matched,
                    status="ready",
                    finished_at=datetime.now(UTC),
                )
            )
            session.commit()
            return {
                "audit_id": audit_id,
                "status": "ready",
                "scanned": scanned,
                "matched": len(matched),
            }
    except SoftTimeLimitExceeded:
        with SessionLocal() as session:
            _mark_error(session, audit_id, "soft_time_limit exceeded (>7200s)")
        raise
    except (OperationalError, DBAPIError):
        # Retryable per the decorator; leave row in 'running' and let
        # autoretry try again. Recovery sweep catches if all retries fail.
        raise
    except Exception:  # noqa: BLE001
        tb = traceback.format_exc()
        with SessionLocal() as session:
            _mark_error(session, audit_id, tb)
        raise


def _mark_error(session, audit_id: int, error_msg: str) -> None:
    session.execute(
        update(LibraryAuditRun)
        .where(LibraryAuditRun.id == audit_id)
        .values(
            status="error",
            error=error_msg,
            finished_at=datetime.now(UTC),
        )
    )
    session.commit()