"""Staged uploads: files land in a temp folder until they are resolved. Flow: - POST /api/uploads/ stage a file (computes MD5, detects duplicates) - GET /api/uploads/ list my staged uploads - GET /api/uploads//file/ serve the staged file (previews, IQDB) - POST /api/uploads//iqdb/ store IQDB candidates -> visual_match - POST /api/uploads//resolve/ link to a post or apply custom metadata - DELETE /api/uploads// discard a staged upload / dismiss a record """ import logging import shutil from pathlib import Path from urllib.parse import urlparse from django.conf import settings from django.contrib.auth import get_user_model from django.core import signing from django.core.exceptions import ValidationError from django.http import Http404 from django.utils import timezone from rest_framework import mixins, status, viewsets from rest_framework.decorators import action from rest_framework.parsers import FormParser, JSONParser, MultiPartParser from rest_framework.permissions import AllowAny from rest_framework.response import Response from . import services from .models import MediaItem, TempUpload from .permissions import CanUpload from .serializers import TempUploadListSerializer, TempUploadSerializer from .tools import HASH_FIELDS, hashed_items, hashes_similarity logger = logging.getLogger(__name__) def build_hash_index(): """Library hash mappings for similarity scans, loaded once per batch. Only items that actually carry perceptual hashes are included; the old per-file scan walked every row (including videos and unchecked items). """ algorithms = list(HASH_FIELDS) return [ (item, {field: getattr(item, field, "") for field in algorithms}) for item in hashed_items(algorithms) ] def match_hashes(hashes, index, limit=10, user=None, request=None): """Library items whose perceptual hashes are close to ``hashes``.""" algorithms = list(HASH_FIELDS) threshold = settings.VISUAL_MATCH_THRESHOLD matches = [] for item, item_hashes in index: similarity = hashes_similarity(hashes, item_hashes, algorithms, threshold) if similarity is None: continue location = item.locations.first() matches.append( { "j_id": f"J-{item.id}", "filename": Path(location.rel_path).name if location else item.md5, "similarity": round(similarity * 100, 1), "thumbnail_url": services.signed_media_url( item, user, "thumbnail", request=request ), } ) matches.sort(key=lambda entry: entry["similarity"], reverse=True) return matches[:limit] def find_library_matches(path, limit=10, user=None, request=None, index=None): """Library items visually similar to a staged file. Pass a prebuilt ``index`` (see build_hash_index) to reuse it across a whole batch instead of rescanning the library per file. """ hashes = services.compute_visual_hashes(path) if not hashes: return [] if index is None: index = build_hash_index() return match_hashes(hashes, index, limit=limit, user=user, request=request) def complete_temp_upload(temp, download_url=None): """Index the upload into the library. With a ``download_url`` the e621 original is fetched and the staged copy is dropped; otherwise the staged file itself moves into the library. """ folder = Path(settings.WATCHED_FOLDER) if download_url: file_data = {} if isinstance(temp.e621_data, dict): file_data = temp.e621_data.get("file") or {} extension = Path(urlparse(download_url).path).suffix if not extension: extension = Path(temp.original_filename).suffix md5 = file_data.get("md5") name = f"{md5}{extension}" if md5 else (temp.original_filename or "download") destination = services.unique_destination(folder, name) try: services.download_file(download_url, destination) except Exception: destination.unlink(missing_ok=True) raise item, _, location, _ = services.index_file(destination, folder) services.rename_location_to_j_id(item, location) services.ensure_visual_hashes(item) if temp.file: temp.file.delete(save=False) else: if not temp.file: raise ValueError("staged file is missing") destination = services.unique_destination( folder, temp.original_filename or "upload" ) with temp.file.open("rb") as source, open(destination, "wb") as target: shutil.copyfileobj(source, target) item, _, location, _ = services.index_file(destination, folder) services.rename_location_to_j_id(item, location) services.ensure_visual_hashes(item) temp.file.delete(save=False) temp.library_item = item temp.status = TempUpload.STATUS_COMPLETED update_fields = [] if item.uploaded_by_id is None: item.uploaded_by = temp.user update_fields.append("uploaded_by") if temp.e621_post_id is not None: item.e621_post_id = temp.e621_post_id item.e621_data = temp.e621_data 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", ] if temp.custom_rating in {"s", "q", "e"} and not item.rating: item.rating = temp.custom_rating update_fields.append("rating") if temp.custom_tags: item.tags = temp.custom_tags update_fields.append("tags") if temp.custom_notes: item.notes = temp.custom_notes update_fields.append("notes") if update_fields: item.save(update_fields=update_fields + ["updated_at"]) temp.save() return item class TempUploadViewSet( mixins.ListModelMixin, mixins.RetrieveModelMixin, mixins.DestroyModelMixin, viewsets.GenericViewSet, ): serializer_class = TempUploadSerializer permission_classes = [CanUpload] parser_classes = [MultiPartParser, FormParser, JSONParser] # Unpaginated: the board shows every staged upload (69-file batches were # silently cut to the API's 48-item page). pagination_class = None http_method_names = ["get", "post", "delete", "head", "options"] def get_serializer_class(self): # The board polls the list, so its payload stays small; the metadata # modal fetches the full row from the detail endpoint. if self.action == "list": return TempUploadListSerializer return TempUploadSerializer def get_queryset(self): # Staged uploads are private: everyone, staff included, only sees # their own board. (The file action still lets staff read bytes by id # for support purposes.) return TempUpload.objects.select_related("library_item").filter( user=self.request.user ) def create(self, request): upload = request.FILES.get("file") if upload is None: return Response( {"detail": "A file is required."}, status=status.HTTP_400_BAD_REQUEST ) extension = Path(upload.name).suffix.lower() if extension not in services.ALLOWED_EXTENSIONS: return Response( {"detail": f"Unsupported file type: {extension or 'unknown'}"}, status=status.HTTP_400_BAD_REQUEST, ) temp = TempUpload.objects.create( user=request.user, file=upload, original_filename=upload.name, size=upload.size, ) temp.md5 = services.compute_md5(temp.file.path) existing = MediaItem.objects.filter(md5=temp.md5).first() if existing is not None: temp.status = TempUpload.STATUS_COMPLETED temp.resolution = TempUpload.RESOLUTION_DUPLICATE temp.library_item = existing temp.file.delete(save=False) # Visual similarity and IQDB run in the background pipeline so a large # batch uploads at full speed and the work survives the browser. temp.save() # Kick the server-side pipeline; staging no longer waits on e621 and # the work continues even if the browser navigates away. from .upload_pipeline import start_pipeline start_pipeline(request.user) return Response( self.get_serializer(temp).data, status=status.HTTP_201_CREATED ) @action(detail=True, methods=["post"], url_path="visual-match") def visual_match(self, request, pk=None): """Run the local visual-similarity pass for one staged upload.""" temp = self.get_object() if temp.status == TempUpload.STATUS_COMPLETED: return Response( {"detail": "This upload is already in the library."}, status=status.HTTP_400_BAD_REQUEST, ) if not temp.file: return Response( {"detail": "The staged file is missing."}, status=status.HTTP_400_BAD_REQUEST, ) matches = find_library_matches( temp.file.path, user=request.user, request=request ) temp.visual_matches = matches if matches and temp.status == TempUpload.STATUS_PENDING: temp.status = TempUpload.STATUS_VISUAL_MATCH temp.save(update_fields=["visual_matches", "status", "updated_at"]) return Response(self.get_serializer(temp).data) @action(detail=False, methods=["get"]) def status(self, request): """Cheap pipeline state for the shell indicator and the upload page.""" from .upload_pipeline import status_payload return Response(status_payload(request.user)) @action(detail=False, methods=["post"]) def process(self, request): """Start (or resume) the pipeline for the caller's staged uploads. Idempotent: the client calls this after staging files, on page load and when a paused run should be retried. """ from .upload_pipeline import start_pipeline, status_payload start_pipeline(request.user) return Response(status_payload(request.user)) @action(detail=True, methods=["post"]) def retry(self, request, pk=None): """Queue one staged upload for another pipeline pass.""" from .upload_pipeline import start_pipeline, status_payload temp = self.get_object() if temp.status == TempUpload.STATUS_COMPLETED: return Response( {"detail": "This upload is already in the library."}, status=status.HTTP_400_BAD_REQUEST, ) if not temp.file: return Response( {"detail": "The staged file is missing."}, status=status.HTTP_400_BAD_REQUEST, ) update = { "claimed_at": None, "attempts": 0, "pipeline_error": "", "updated_at": timezone.now(), } # An explicit phase re-runs that one check even if it already ran. phase = str(request.data.get("phase") or "").strip() if phase == "md5": update["e621_checked_at"] = None elif phase == "visual": update["visual_matches"] = None update["visual_checked_at"] = None elif phase == "iqdb": update["iqdb_data"] = None if temp.status == TempUpload.STATUS_ERROR: # A failed import has to go through the MD5 phase again so the # completion is retried; other errors only re-run missing phases. update["status"] = TempUpload.STATUS_PENDING update["e621_checked_at"] = None elif temp.status not in ( TempUpload.STATUS_PENDING, TempUpload.STATUS_VISUAL_MATCH, ): update["status"] = TempUpload.STATUS_PENDING TempUpload.objects.filter(pk=temp.pk).update(**update) start_pipeline(request.user) return Response(status_payload(request.user)) @action(detail=False, methods=["post"], url_path="retry-all") def retry_all(self, request): """Queue every retryable staged upload for another pipeline pass.""" from .upload_pipeline import start_pipeline, status_payload now = timezone.now() retryable = self.get_queryset().filter( status__in=[ TempUpload.STATUS_PENDING, TempUpload.STATUS_VISUAL_MATCH, TempUpload.STATUS_ERROR, ] ) retryable.exclude(file="").update( claimed_at=None, attempts=0, pipeline_error="", updated_at=now, ) retryable.exclude(file="").filter(status=TempUpload.STATUS_ERROR).update( status=TempUpload.STATUS_PENDING, e621_checked_at=None, ) start_pipeline(request.user) return Response(status_payload(request.user)) @action(detail=False, methods=["post"], url_path="discard-bulk") def discard_bulk(self, request): """Discard many staged uploads in one request (the board's "all").""" ids = request.data.get("temp_ids") if not isinstance(ids, list) or not ids: return Response( {"detail": "temp_ids must be a non-empty list."}, status=status.HTTP_400_BAD_REQUEST, ) if len(ids) > 1000: return Response( {"detail": "Too many ids in one request (max 1000)."}, status=status.HTTP_400_BAD_REQUEST, ) values = list(dict.fromkeys(str(value) for value in ids)) try: queryset = self.get_queryset().filter(pk__in=values) except (ValidationError, ValueError): return Response( {"detail": "One or more temp_ids are not valid upload ids."}, status=status.HTTP_400_BAD_REQUEST, ) by_id = {str(temp.pk): temp for temp in queryset} discarded: list[str] = [] errors: list[dict[str, str]] = [] for value in values: temp = by_id.get(value) if temp is None: errors.append({"temp_id": value, "error": "not found"}) continue try: self.perform_destroy(temp) discarded.append(value) except Exception as exc: # noqa: BLE001 - report per-file failures logger.exception("Could not discard staged upload %s", value) errors.append({"temp_id": value, "error": str(exc)}) return Response({"discarded": discarded, "errors": errors}) @action(detail=True, methods=["get", "head"], permission_classes=[AllowAny]) def file(self, request, pk=None): """Serve the staged file; accepts a signed URL for media tags.""" user = request.user if request.user.is_authenticated else None if user is None: signature = request.query_params.get("sig") if signature: try: payload = signing.loads( signature, salt=services.UPLOAD_FILE_SALT, max_age=86400, ) except signing.BadSignature: payload = None if payload and str(payload.get("temp")) == str(pk): user = ( get_user_model() .objects.filter(pk=payload.get("user")) .first() ) if user is None or not user.can_upload: return Response( {"detail": "Authentication required."}, status=status.HTTP_401_UNAUTHORIZED, ) temp = TempUpload.objects.filter(pk=pk).first() is_owner = temp is not None and temp.user_id == user.id if temp is None or not (is_owner or user.is_app_staff): raise Http404 if not temp.file: raise Http404 return services.serve_file(request, temp.file.path) @action(detail=True, methods=["post"]) def iqdb(self, request, pk=None): temp = self.get_object() results = request.data.get("results") if not isinstance(results, list): return Response( {"detail": "results must be a list."}, status=status.HTTP_400_BAD_REQUEST, ) temp.iqdb_data = services.sanitize_iqdb_results(results) if temp.status == TempUpload.STATUS_PENDING and temp.iqdb_data: temp.status = TempUpload.STATUS_VISUAL_MATCH temp.save(update_fields=["iqdb_data", "status", "updated_at"]) return Response(self.get_serializer(temp).data) @action(detail=True, methods=["post"]) def resolve(self, request, pk=None): temp = self.get_object() if temp.status == TempUpload.STATUS_COMPLETED: return Response( {"detail": "This upload is already in the library."}, status=status.HTTP_400_BAD_REQUEST, ) mode = str(request.data.get("mode") or "").strip() download_url = None if mode == "link": post = request.data.get("post") post_id = request.data.get("post_id") trimmed = services.trim_e621_post(post) has_numeric_id = post_id is not None and str(post_id).isdigit() if trimmed is None and not has_numeric_id: return Response( {"detail": "A post payload or post_id is required."}, status=status.HTTP_400_BAD_REQUEST, ) if trimmed is not None and trimmed.get("id") is not None: temp.e621_post_id = int(trimmed["id"]) elif has_numeric_id: temp.e621_post_id = int(post_id) temp.e621_data = trimmed temp.resolution = ( TempUpload.RESOLUTION_AUTO_MD5 if request.data.get("auto") else TempUpload.RESOLUTION_LINKED ) rating = str(request.data.get("rating") or "") if rating in {"s", "q", "e"}: temp.custom_rating = rating tags = services.parse_tags(request.data.get("tags")) if tags is not None: temp.custom_tags = tags notes = request.data.get("notes") if notes is not None: temp.custom_notes = str(notes) # Prefer the e621 original; fall back to the staged file when the # post has no URL or its file is byte-identical to the upload. file_data = (trimmed or {}).get("file") or {} candidate_url = str( request.data.get("file_url") or file_data.get("url") or "" ).strip() post_md5 = str(file_data.get("md5") or "").strip().lower() if candidate_url and post_md5 and post_md5 == temp.md5.lower(): candidate_url = "" download_url = candidate_url or None elif mode == "custom": rating = str(request.data.get("rating") or "") temp.custom_rating = rating if rating in {"s", "q", "e"} else "" temp.custom_tags = services.parse_tags(request.data.get("tags")) or [] temp.custom_notes = str(request.data.get("notes") or "") temp.resolution = TempUpload.RESOLUTION_CUSTOM else: return Response( {"detail": "mode must be 'link' or 'custom'."}, status=status.HTTP_400_BAD_REQUEST, ) temp.save() try: complete_temp_upload(temp, download_url=download_url) except Exception as exc: # noqa: BLE001 - report completion failures logger.exception("Could not complete staged upload %s", temp.id) temp.status = TempUpload.STATUS_ERROR temp.save(update_fields=["status", "updated_at"]) return Response( {"detail": f"Could not finish the upload: {exc}"}, status=status.HTTP_400_BAD_REQUEST, ) temp.refresh_from_db() return Response(self.get_serializer(temp).data) @action(detail=False, methods=["post"], url_path="link-bulk") def link_bulk(self, request): """Attach e621 posts to many staged uploads in one request. The upload board's MD5 phase sends one posts.json query per 75 files and hands the matches over here, so a whole batch moves into the library in one update instead of one request per file. """ links = request.data.get("links") if not isinstance(links, list) or not links: return Response( {"detail": "links must be a non-empty list."}, status=status.HTTP_400_BAD_REQUEST, ) ids = [] for link in links: if not isinstance(link, dict) or link.get("temp_id") is None: return Response( {"detail": "Each link needs a temp_id."}, status=status.HTTP_400_BAD_REQUEST, ) ids.append(str(link["temp_id"])) try: queryset = self.get_queryset().filter(pk__in=ids) except (ValidationError, ValueError): return Response( {"detail": "One or more temp_ids are not valid upload ids."}, status=status.HTTP_400_BAD_REQUEST, ) by_id = {str(temp.id): temp for temp in queryset} updated = [] errors = [] for link in links: temp_id = str(link["temp_id"]) temp = by_id.get(temp_id) if temp is None: errors.append({"temp_id": temp_id, "error": "not found"}) continue if temp.status == TempUpload.STATUS_COMPLETED: errors.append({"temp_id": temp_id, "error": "already in the library"}) continue trimmed = services.trim_e621_post(link.get("post")) post_id = link.get("post_id") has_numeric_id = post_id is not None and str(post_id).isdigit() if trimmed is None and not has_numeric_id: errors.append({"temp_id": temp_id, "error": "no post payload"}) continue if trimmed is not None and trimmed.get("id") is not None: temp.e621_post_id = int(trimmed["id"]) elif has_numeric_id: temp.e621_post_id = int(post_id) temp.e621_data = trimmed temp.resolution = ( TempUpload.RESOLUTION_AUTO_MD5 if link.get("auto", True) else TempUpload.RESOLUTION_LINKED ) # An MD5 match means the staged file is byte-identical, so there # is nothing to download from e621; only link a remote URL when # the post's file differs. file_data = (trimmed or {}).get("file") or {} candidate_url = str( link.get("file_url") or file_data.get("url") or "" ).strip() post_md5 = str(file_data.get("md5") or "").strip().lower() if candidate_url and post_md5 and post_md5 == temp.md5.lower(): candidate_url = "" temp.save() try: complete_temp_upload(temp, download_url=candidate_url or None) except Exception as exc: # noqa: BLE001 - report per-file failures logger.exception("Could not complete staged upload %s", temp.id) temp.status = TempUpload.STATUS_ERROR temp.save(update_fields=["status", "updated_at"]) errors.append({"temp_id": temp_id, "error": str(exc)}) continue temp.refresh_from_db() updated.append(self.get_serializer(temp).data) return Response({"updated": updated, "errors": errors}) @action(detail=False, methods=["post"], url_path="resolve-bulk") def resolve_bulk(self, request): """Move many staged uploads into the library with one rating. The upload board's bulk tool: pick a rating, tick pending files, and they are all resolved as custom entries (their staged tags/notes are kept). Own rows only, like the rest of the viewset. """ ids = request.data.get("temp_ids") if not isinstance(ids, list) or not ids: return Response( {"detail": "temp_ids must be a non-empty list."}, status=status.HTTP_400_BAD_REQUEST, ) rating = str(request.data.get("rating") or "").strip() if rating not in {"s", "q", "e"}: return Response( {"detail": "rating must be one of s, q or e."}, status=status.HTTP_400_BAD_REQUEST, ) try: queryset = self.get_queryset().filter( pk__in=[str(value) for value in ids] ) except (ValidationError, ValueError): return Response( {"detail": "One or more ids are not valid upload ids."}, status=status.HTTP_400_BAD_REQUEST, ) resolved: list[str] = [] errors: list[dict[str, str]] = [] for temp in queryset: if temp.status == TempUpload.STATUS_COMPLETED: errors.append( {"temp_id": str(temp.id), "error": "already in the library"} ) continue temp.custom_rating = rating temp.resolution = TempUpload.RESOLUTION_CUSTOM temp.save(update_fields=["custom_rating", "resolution", "updated_at"]) try: complete_temp_upload(temp) except Exception as exc: # noqa: BLE001 - report per-file failures logger.exception("Could not complete staged upload %s", temp.id) temp.status = TempUpload.STATUS_ERROR temp.save(update_fields=["status", "updated_at"]) errors.append({"temp_id": str(temp.id), "error": str(exc)}) continue resolved.append(str(temp.id)) found = {str(temp.id) for temp in queryset} for value in sorted({str(value) for value in ids} - found): errors.append({"temp_id": value, "error": "not found"}) return Response({"resolved": resolved, "errors": errors}) def perform_destroy(self, instance): if instance.file: instance.file.delete(save=False) instance.delete()