Upload updates

backend/apps/library/management/commands/process_uploads.py

@@ -0,0 +1,68 @@
+"""Process staged uploads: e621 MD5, visual similarity, IQDB.
+
+Runs synchronously, unlike the daemon thread the API starts on demand. Useful
+for tests, for a manual drain after an outage and for the scheduler if a
+deployment wants a periodic safety net.
+
+    python manage.py process_uploads             # every user with queued work, once
+    python manage.py process_uploads --user 3    # one user
+    python manage.py process_uploads --loop 60   # keep draining every 60s
+"""
+
+import time
+
+from django.core.management.base import BaseCommand
+
+from apps.library.models import TempUpload
+from apps.library.upload_pipeline import (
+    MAX_ATTEMPTS,
+    OUTSTANDING_Q,
+    WORK_STATUSES,
+    reap_stale_claims,
+    run_pipeline,
+)
+
+
+class Command(BaseCommand):
+    help = "Run the staged-upload pipeline (MD5 -> visual similarity -> IQDB)."
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            "--user",
+            type=int,
+            default=None,
+            help="Only process this user id.",
+        )
+        parser.add_argument(
+            "--loop",
+            type=int,
+            default=0,
+            metavar="SECONDS",
+            help="Keep draining every SECONDS seconds instead of exiting.",
+        )
+
+    def handle(self, *args, **options):
+        interval = options["loop"] or 0
+        while True:
+            self.drain(user_id=options["user"])
+            if interval <= 0:
+                return
+            time.sleep(interval)
+
+    def drain(self, user_id=None):
+        reap_stale_claims()
+        queryset = (
+            TempUpload.objects.filter(status__in=WORK_STATUSES)
+            .filter(OUTSTANDING_Q)
+            .filter(attempts__lt=MAX_ATTEMPTS)
+        )
+        if user_id is not None:
+            queryset = queryset.filter(user_id=user_id)
+        user_ids = list(queryset.values_list("user_id", flat=True).distinct())
+        if not user_ids:
+            self.stdout.write("No staged uploads need processing.")
+            return
+        for value in user_ids:
+            self.stdout.write(f"Processing staged uploads for user {value}...")
+            run_pipeline(value)
+        self.stdout.write(self.style.SUCCESS(f"Processed {len(user_ids)} queue(s)."))