feat: maintenance worker can pick up prio users

Haidra-Org · Apr 14, 2024 · 15457b9 · 15457b9
1 parent 7d5cd1b
commit 15457b9
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 27 deletions.
diff --git a/horde/classes/stable/waiting_prompt.py b/horde/classes/stable/waiting_prompt.py
@@ -1,6 +1,6 @@
 import copy
-import random
 import os
+import random
 
 from sqlalchemy.sql import expression
 

diff --git a/horde/database/functions.py b/horde/database/functions.py
@@ -1,8 +1,8 @@
 import json
+import os
 import time
 import urllib.parse
 import uuid
-import os
 from datetime import datetime, timedelta
 
 from sqlalchemy import Boolean, and_, func, not_, or_
@@ -804,10 +804,6 @@ def get_sorted_wp_filtered_to_worker(worker, models_list=None, blacklist=None, p
                 ImageWaitingPrompt.nsfw == False,  # noqa E712
                 worker.nsfw == True,  # noqa E712
             ),
-            or_(
-                worker.maintenance == False,  # noqa E712
-                ImageWaitingPrompt.user_id == worker.user_id,
-            ),
             or_(
                 check_bridge_capability("r2", worker.bridge_agent),
                 ImageWaitingPrompt.r2 == False,  # noqa E712
@@ -847,6 +843,11 @@ def get_sorted_wp_filtered_to_worker(worker, models_list=None, blacklist=None, p
     if priority_user_ids:
         final_wp_list = final_wp_list.filter(ImageWaitingPrompt.user_id.in_(priority_user_ids))
         final_wp_list = final_wp_list.filter(
+            # Workers in maintenance can still pick up their owner or their friends
+            or_(
+                worker.maintenance == False,  # noqa E712
+                ImageWaitingPrompt.user_id.in_(priority_user_ids),
+            ),
             or_(
                 WPAllowedWorkers.id.is_(None),
                 and_(
@@ -859,33 +860,40 @@ def get_sorted_wp_filtered_to_worker(worker, models_list=None, blacklist=None, p
                 ),
             ),
         )
-    # If HORDE_REQUIRE_MATCHED_TARGETING is set to 1, we disable using WPAllowedWorkers
-    # Targeted requests will only be picked up in the condition above as it will include the
-    # filter to ensure the worker also has that user as a priority
-    elif os.getenv("HORDE_REQUIRE_MATCHED_TARGETING", 0) == 1:
+    else:
         final_wp_list = final_wp_list.filter(
             or_(
-                WPAllowedWorkers.id.is_(None),
-                and_(
-                    ImageWaitingPrompt.worker_blacklist.is_(True),
-                    WPAllowedWorkers.worker_id != worker.id,
-                ),
+                worker.maintenance == False,  # noqa E712
+                ImageWaitingPrompt.user_id == worker.user_id,
             ),
         )
-    else:
-        final_wp_list = final_wp_list.filter(
-            or_(
-                WPAllowedWorkers.id.is_(None),
-                and_(
-                    ImageWaitingPrompt.worker_blacklist.is_(False),
-                    WPAllowedWorkers.worker_id == worker.id,
+        # If HORDE_REQUIRE_MATCHED_TARGETING is set to 1, we disable using WPAllowedWorkers
+        # Targeted requests will only be picked up in the condition above as it will include the
+        # filter to ensure the worker also has that user as a priority
+        if os.getenv("HORDE_REQUIRE_MATCHED_TARGETING", 0) == 1:
+            final_wp_list = final_wp_list.filter(
+                or_(
+                    WPAllowedWorkers.id.is_(None),
+                    and_(
+                        ImageWaitingPrompt.worker_blacklist.is_(True),
+                        WPAllowedWorkers.worker_id != worker.id,
+                    ),
                 ),
-                and_(
-                    ImageWaitingPrompt.worker_blacklist.is_(True),
-                    WPAllowedWorkers.worker_id != worker.id,
+            )
+        else:
+            final_wp_list = final_wp_list.filter(
+                or_(
+                    WPAllowedWorkers.id.is_(None),
+                    and_(
+                        ImageWaitingPrompt.worker_blacklist.is_(False),
+                        WPAllowedWorkers.worker_id == worker.id,
+                    ),
+                    and_(
+                        ImageWaitingPrompt.worker_blacklist.is_(True),
+                        WPAllowedWorkers.worker_id != worker.id,
+                    ),
                 ),
-            ),
-        )
+            )
 
     # logger.debug(final_wp_list)
     final_wp_list = (