Wait for all uploads, then batch MD5 -> visual -> IQDB with bulk links
Batching must not start while files are still being uploaded, and the MD5 phase must move a whole chunk at once instead of one resolve per file: - the upload queue drains completely first (failed uploads included) before any matching starts; - phase 1 asks e621 for every md5 (75 per posts.json request), builds the md5 -> post map from the response, and sends the matches to the new POST /api/uploads/link-bulk/ action, so a whole 75-file chunk moves into Indexed in a single board update; - link-bulk indexes the staged file directly when the post's MD5 matches (identical bytes), so there is no per-file download round trip; - phase 2 runs local visual similarity for whatever stayed pending, phase 3 the IQDB queue. Verified end to end with real e621 files: one md5 query for the batch, one link-bulk call, both matching files flipped to Indexed together, then the visual and IQDB phases. 23 library tests green (link-bulk, visual phase, deferred visual matching).
This commit is contained in:
@@ -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"."""
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user