J621/backend/apps/library/matching.py

"""Server-side matching of library items against e621.

MD5 lookups, manual post links, metadata refresh (which is also how deleted
posts are detected) and the background library-wide scan.
"""

import logging
import threading
from datetime import timedelta

from django.db import connection
from django.utils import timezone

from . import e621
from .models import MatchTask, MediaItem
from .services import trim_e621_post

logger = logging.getLogger(__name__)

STALE_AFTER = timedelta(minutes=30)
RATING_VALUES = {"s", "q", "e"}


# --- single-item operations -------------------------------------------------


def apply_post_to_item(item, post):
    """Store trimmed e621 metadata on an item and mark it matched."""
    trimmed = trim_e621_post(post)
    if trimmed is None or not trimmed.get("id"):
        raise e621.E621Error("e621 returned an unexpected post payload.")
    item.e621_post_id = int(trimmed["id"])
    item.e621_data = trimmed
    item.e621_match_status = MediaItem.E621_MATCHED
    item.e621_checked_at = timezone.now()
    update_fields = [
        "e621_post_id",
        "e621_data",
        "e621_match_status",
        "e621_checked_at",
    ]
    rating = trimmed.get("rating")
    if not item.rating and rating in RATING_VALUES:
        item.rating = rating
        update_fields.append("rating")
    item.save(update_fields=update_fields + ["updated_at"])
    return trimmed


def mark_not_found(item):
    item.e621_match_status = MediaItem.E621_NOT_FOUND
    item.e621_checked_at = timezone.now()
    item.save(update_fields=["e621_match_status", "e621_checked_at", "updated_at"])


def mark_deleted(item):
    item.e621_match_status = MediaItem.E621_DELETED
    item.e621_checked_at = timezone.now()
    item.save(update_fields=["e621_match_status", "e621_checked_at", "updated_at"])


def match_item_by_md5(item, user):
    """Look the item's MD5 up on e621. Returns the resulting match status."""
    post = e621.find_post_by_md5(user, item.md5)
    if post is None:
        mark_not_found(item)
        return item.e621_match_status
    apply_post_to_item(item, post)
    return item.e621_match_status


def refresh_item(item, user):
    """Re-fetch a matched post: refresh metadata or flag it as deleted."""
    if not item.e621_post_id:
        return match_item_by_md5(item, user)
    try:
        post = e621.fetch_post(user, item.e621_post_id)
    except e621.E621NotFound:
        mark_deleted(item)
        return item.e621_match_status
    apply_post_to_item(item, post)
    return item.e621_match_status


def link_item_to_post(item, post_id, user):
    """Manually link an item to a post id. Returns True if the MD5s match."""
    post = e621.fetch_post(user, post_id)
    trimmed = apply_post_to_item(item, post)
    post_md5 = str((trimmed.get("file") or {}).get("md5") or "").strip().lower()
    return bool(post_md5) and post_md5 == item.md5.lower()


def check_item(item, user):
    """Best check for this item: refresh a known link, else MD5 lookup."""
    if item.e621_post_id and item.e621_match_status == MediaItem.E621_MATCHED:
        return refresh_item(item, user)
    return match_item_by_md5(item, user)


def unlink_item(item):
    """Forget the e621 link so the item can be checked again."""
    item.e621_post_id = None
    item.e621_data = None
    item.e621_match_status = MediaItem.E621_UNKNOWN
    item.e621_checked_at = None
    item.save(
        update_fields=[
            "e621_post_id",
            "e621_data",
            "e621_match_status",
            "e621_checked_at",
            "updated_at",
        ]
    )
    return item


# --- batch scan -------------------------------------------------------------


def reap_stale_match_tasks():
    """Surface scans left hanging by a recycled worker as errors."""
    cutoff = timezone.now() - STALE_AFTER
    return MatchTask.objects.filter(
        status__in=[MatchTask.STATUS_PENDING, MatchTask.STATUS_RUNNING],
        updated_at__lt=cutoff,
    ).update(
        status=MatchTask.STATUS_ERROR,
        error="The worker restarted before this scan finished.",
        updated_at=timezone.now(),
    )


def scope_queryset(scope):
    queryset = MediaItem.objects.order_by("id")
    if scope == MatchTask.SCOPE_ALL:
        return queryset
    return queryset.filter(e621_match_status=MediaItem.E621_UNKNOWN)


def start_match_task(task_id):
    thread = threading.Thread(target=run_match_task, args=(task_id,), daemon=True)
    thread.start()


def run_match_task(task_id):
    task = MatchTask.objects.filter(id=task_id).first()
    if task is None:
        return

    item_ids = list(scope_queryset(task.scope).values_list("id", flat=True))
    counters = {"matched": 0, "not_found": 0, "deleted": 0}
    failed = 0
    MatchTask.objects.filter(id=task_id).update(
        status=MatchTask.STATUS_RUNNING,
        total=len(item_ids),
        processed=0,
        updated_at=timezone.now(),
    )

    def save_progress(processed):
        MatchTask.objects.filter(id=task_id).update(
            processed=processed,
            failed=failed,
            updated_at=timezone.now(),
            **counters,
        )

    try:
        processed = 0
        for item_id in item_ids:
            if MatchTask.objects.filter(id=task_id, cancelled=True).exists():
                MatchTask.objects.filter(id=task_id).update(
                    status=MatchTask.STATUS_CANCELLED, updated_at=timezone.now()
                )
                return
            item = MediaItem.objects.filter(id=item_id).first()
            if item is None:
                processed += 1
                continue
            try:
                status = check_item(item, task.user)
                counters[status] = counters.get(status, 0) + 1
            except e621.E621Error as exc:
                # A broken credential or an e621 outage would fail every
                # remaining item; abort instead of hammering the API.
                logger.warning("Match scan %s aborted: %s", task_id, exc)
                failed += 1
                MatchTask.objects.filter(id=task_id).update(
                    status=MatchTask.STATUS_ERROR,
                    processed=processed,
                    failed=failed,
                    error=str(exc),
                    updated_at=timezone.now(),
                    **counters,
                )
                return
            processed += 1
            save_progress(processed)
        MatchTask.objects.filter(id=task_id).update(
            status=MatchTask.STATUS_COMPLETE,
            processed=len(item_ids),
            updated_at=timezone.now(),
            **counters,
        )
    finally:
        # Background threads hold their own DB connection; release it so
        # Gunicorn workers do not leak connections when threads finish.
        connection.close()