"""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()