diff --git a/backend/apps/library/tests/test_uploads.py b/backend/apps/library/tests/test_uploads.py index e8a6f4d..ea44c7a 100644 --- a/backend/apps/library/tests/test_uploads.py +++ b/backend/apps/library/tests/test_uploads.py @@ -132,13 +132,17 @@ class StagedUploadWorkflowTests(TestCase): ) return client - def make_temp(self, user, label, status=TempUpload.STATUS_PENDING): + def make_temp( + self, user, label, status=TempUpload.STATUS_PENDING, payload=None + ): return TempUpload.objects.create( user=user, - file=SimpleUploadedFile(f"{label}.png", TINY_PNG, content_type="image/png"), + file=SimpleUploadedFile( + f"{label}.png", payload or TINY_PNG, content_type="image/png" + ), original_filename=f"{label}.png", md5=hashlib.md5(label.encode()).hexdigest(), - size=len(TINY_PNG), + size=len(payload or TINY_PNG), status=status, ) @@ -261,6 +265,57 @@ class StagedUploadWorkflowTests(TestCase): response = client.post(f"/api/uploads/{temp.id}/visual-match/") self.assertEqual(response.status_code, 400) + def test_link_bulk_moves_a_batch_with_one_call(self): + client = self.api_client(self.uploader) + first = self.make_temp(self.uploader, "bulk-link-1") + second = self.make_temp( + self.uploader, "bulk-link-2", payload=similar_png((10, 200, 30)) + ) + response = jpost( + client, + "/api/uploads/link-bulk/", + { + "links": [ + { + "temp_id": str(first.id), + "post": { + "id": 900001, + "rating": "s", + "file": { + "md5": first.md5, + # Identical MD5: the staged file is indexed + # without downloading the URL. + "url": "https://static1.e621.net/data/fake-1.png", + }, + }, + }, + { + "temp_id": str(second.id), + "post": { + "id": 900002, + "rating": "q", + "file": { + "md5": second.md5, + "url": "https://static1.e621.net/data/fake-2.png", + }, + }, + }, + ] + }, + ) + self.assertEqual(response.status_code, 200) + body = response.json() + self.assertEqual(body["errors"], []) + self.assertEqual(len(body["updated"]), 2) + for temp, post_id in ((first, 900001), (second, 900002)): + temp.refresh_from_db() + self.assertEqual(temp.status, TempUpload.STATUS_COMPLETED) + self.assertEqual(temp.library_item_id is not None, True) + self.assertEqual(temp.e621_post_id, post_id) + self.assertEqual(temp.resolution, TempUpload.RESOLUTION_AUTO_MD5) + self.assertEqual(temp.library_item.e621_post_id, post_id) + self.assertEqual(MediaItem.objects.count(), 2) + class IqdbRecordingTests(TestCase): """The modal needs to tell "checked, no match" from "never checked".""" diff --git a/backend/apps/library/uploads.py b/backend/apps/library/uploads.py index 6bae436..09ef3f1 100644 --- a/backend/apps/library/uploads.py +++ b/backend/apps/library/uploads.py @@ -348,6 +348,88 @@ class TempUploadViewSet( 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. diff --git a/frontend/src/features/upload/UploadPage.tsx b/frontend/src/features/upload/UploadPage.tsx index 9ad34b6..194cbe1 100644 --- a/frontend/src/features/upload/UploadPage.tsx +++ b/frontend/src/features/upload/UploadPage.tsx @@ -28,6 +28,7 @@ import { fetchPostsByMd5, iqdbSearch, RATING_LABELS, + type E621Post, } from "@/lib/e621"; import { formatBytes } from "@/lib/format"; import type { E621IqdbCandidate, Rating, TempUpload } from "@/lib/types"; @@ -857,12 +858,6 @@ export default function UploadPage() { }); } - async function refreshUploads(): Promise { - const fresh = await api("/api/uploads/"); - queryClient.setQueryData(["uploads"], fresh); - return fresh; - } - // One serial drain for all IQDB checks: clicking "Check similarity" many // times used to start overlapping runs that re-downloaded the same staging // blobs and piled requests onto the e621 queue until everything stalled. @@ -962,61 +957,88 @@ export default function UploadPage() { } } + /** One link-bulk request for the batch, with a smaller-chunk fallback. */ + async function bulkLink( + links: { temp_id: string; post: E621Post }[], + ): Promise { + try { + const result = await api<{ updated: TempUpload[] }>( + "/api/uploads/link-bulk/", + { method: "POST", json: { links } }, + ); + return result.updated; + } catch (error) { + if (links.length <= 15) throw error; + // Very large files can make one batch request slow; retry in pieces. + const updated: TempUpload[] = []; + for (let index = 0; index < links.length; index += 15) { + const result = await api<{ updated: TempUpload[] }>( + "/api/uploads/link-bulk/", + { method: "POST", json: { links: links.slice(index, index + 15) } }, + ); + updated.push(...result.updated); + } + return updated; + } + } + /** - * Phased processing for one finished upload batch: every file goes through - * the e621 MD5 lookup (batched), then local visual similarity, then IQDB. - * Phases never interleave per file, and each step lands on the board as it - * finishes. + * Phases for one finished upload batch, strictly in order: every file's + * e621 MD5 (one query + one bulk link call per 75 files, so a chunk moves + * into the library in a single board update), then every file's local + * visual similarity, then IQDB. */ - async function runUploadPipeline(created: TempUpload[]) { + async function processUploads(created: TempUpload[]) { const creds = effectiveCredentials(credentials); - const targets = created.filter((temp) => temp.status === "pending"); - if (targets.length === 0) { + const pending = created.filter((temp) => temp.status === "pending"); + if (pending.length === 0) { invalidateUploads(); return; } - const targetIds = new Set(targets.map((temp) => temp.temp_id)); + const ids = new Set(pending.map((temp) => temp.temp_id)); - // Phase 1: e621 MD5 lookup, 75 md5: metatags per query. - setPipeline({ phase: "md5", done: 0, total: targets.length }); + // Phase 1: ask e621 for every md5, 75 per request; each chunk is then + // linked with a single bulk call so it flips to Indexed together. + setPipeline({ phase: "md5", done: 0, total: pending.length }); try { - for (let index = 0; index < targets.length; index += MD5_BATCH_SIZE) { - const md5Batch = targets.slice(index, index + MD5_BATCH_SIZE); + for (let index = 0; index < pending.length; index += MD5_BATCH_SIZE) { + const chunk = pending.slice(index, index + MD5_BATCH_SIZE); const posts = await fetchPostsByMd5( creds, - md5Batch.map((temp) => temp.md5), + chunk.map((temp) => temp.md5), ); - for (const post of posts) { - const match = md5Batch.find((temp) => temp.md5 === post.file.md5); - if (!match) continue; - try { - const updated = await api( - `/api/uploads/${match.temp_id}/resolve/`, - { method: "POST", json: { mode: "link", post, auto: true } }, + const byMd5 = new Map(posts.map((post) => [post.file.md5, post])); + const links = chunk.flatMap((temp) => { + const post = temp.md5 ? byMd5.get(temp.md5) : undefined; + return post ? [{ temp_id: temp.temp_id, post }] : []; + }); + if (links.length > 0) { + const updated = await bulkLink(links); + if (updated.length > 0) { + // One request, one board update: the chunk moves together. + const byId = new Map(updated.map((temp) => [temp.temp_id, temp])); + queryClient.setQueryData(["uploads"], (current) => + current + ? current.map((item) => byId.get(item.temp_id) ?? item) + : current, ); - upsertUpload(updated); - } catch { - // Leave it pending. } } setPipeline((current) => current - ? { - ...current, - done: Math.min(index + MD5_BATCH_SIZE, targets.length), - } + ? { ...current, done: Math.min(index + MD5_BATCH_SIZE, pending.length) } : current, ); } } catch { - // e621 unavailable; those files stay pending. + // e621 unavailable or a chunk failed; those files stay pending. } setPipeline(null); - // Phase 2: local visual similarity, one file at a time. + // Phase 2: local visual similarity for whatever is still pending. const cached = queryClient.getQueryData(["uploads"]) ?? []; const visualTargets = cached.filter( - (temp) => targetIds.has(temp.temp_id) && temp.status === "pending", + (temp) => ids.has(temp.temp_id) && temp.status === "pending", ); setPipeline({ phase: "visual", done: 0, total: visualTargets.length }); for (const temp of visualTargets) { @@ -1035,10 +1057,12 @@ export default function UploadPage() { } setPipeline(null); - // Phase 3: IQDB for everything still unresolved (serial queue with its + // Phase 3: IQDB for the batch's unresolved files (serial queue with its // own progress, cooldown and fatal-error stop). - const fresh = await refreshUploads(); - checkSimilarity(unresolvedIqdbIds(fresh)); + const remaining = ( + queryClient.getQueryData(["uploads"]) ?? [] + ).filter((temp) => ids.has(temp.temp_id)); + checkSimilarity(unresolvedIqdbIds(remaining)); } /** Mark an entry done, fade it out and drop it once it has been seen. */ @@ -1096,9 +1120,10 @@ export default function UploadPage() { runningRef.current = false; setBusy(false); invalidateUploads(); + // Batching only starts once every upload settled (success or failure). if (created.length > 0) { setProcessing(true); - void runUploadPipeline(created).finally(() => setProcessing(false)); + void processUploads(created).finally(() => setProcessing(false)); } } @@ -1199,12 +1224,6 @@ export default function UploadPage() { retry checks - ) : busy ? ( - ) : pipeline ? ( + ) : busy ? ( + ) : processing && iqdbProgress ? (