Add REQUEST_ID_SUBCLASSES

Author: Marishka17

cvat/apps/redis_handler/apps.py

@@ -4,6 +4,7 @@
 
 
 from contextlib import suppress
+from typing import cast
 
 from django.apps import AppConfig
 from django.conf import settings
@@ -23,10 +24,17 @@ def __getitem__(self, key: str | tuple) -> str:
 ACTION_TO_QUEUE = LayeredKeyDict()
 QUEUE_TO_PARSED_JOB_ID_CLS = {}
 
+REQUEST_ID_SUBCLASSES = set()
+
 
 def initialize_mappings():
     from cvat.apps.redis_handler.rq import RequestId
 
+    def init_subclasses(cur_cls: type[RequestId] = RequestId):
+        for subclass in cur_cls.__subclasses__():
+            REQUEST_ID_SUBCLASSES.add(subclass)
+            init_subclasses(subclass)
+
     for queue_name, queue_conf in settings.RQ_QUEUES.items():
         if path_to_parsed_job_id_cls := queue_conf.get("PARSED_JOB_ID_CLASS"):
             parsed_job_id_cls = import_string(path_to_parsed_job_id_cls)
@@ -37,10 +45,28 @@ def initialize_mappings():
                 )
 
             for queue_selector in parsed_job_id_cls.QUEUE_SELECTORS:
+                if not isinstance(queue_selector, (tuple, str)):
+                    raise ImproperlyConfigured("Wrong queue selector, must be either tuple or str")
                 ACTION_TO_QUEUE[queue_selector] = queue_name
 
             QUEUE_TO_PARSED_JOB_ID_CLS[queue_name] = parsed_job_id_cls
 
+    init_subclasses()
+    # check that each subclass that has QUEUE_SELECTORS can be used to determine the queue
+    for subclass in REQUEST_ID_SUBCLASSES:
+        subclass = cast(RequestId, subclass)
+        if subclass.LEGACY_FORMAT_PATTERNS and not subclass.QUEUE_SELECTORS:
+            raise ImproperlyConfigured(
+                f"Subclass {subclass.__name__} has LEGACY_FORMAT_PATTERNS - QUEUE_SELECTORS must be defined"
+            )
+
+        if subclass.QUEUE_SELECTORS:
+            for queue_selector in subclass.QUEUE_SELECTORS:
+                if not ACTION_TO_QUEUE.get(queue_selector):
+                    raise ImproperlyConfigured(
+                        f"Queue selector {queue_selector!r} for the class {subclass.__name__!r} is missed in the queue configuration"
+                    )
+
 
 class RedisHandlerConfig(AppConfig):
     name = "cvat.apps.redis_handler"

cvat/apps/redis_handler/rq.py

@@ -7,11 +7,14 @@
 from uuid import UUID
 
 import attrs
-from django.core.exceptions import ImproperlyConfigured
 from django.utils.html import escape
 from rq.job import Job as RQJob
 
-from cvat.apps.redis_handler.apps import ACTION_TO_QUEUE, QUEUE_TO_PARSED_JOB_ID_CLS
+from cvat.apps.redis_handler.apps import (
+    ACTION_TO_QUEUE,
+    QUEUE_TO_PARSED_JOB_ID_CLS,
+    REQUEST_ID_SUBCLASSES,
+)
 
 
 class IncorrectRequestIdError(ValueError):
@@ -58,7 +61,7 @@ class RequestId:
         default=_default_from_class_attr("ACTION_DEFAULT_VALUE"),
     )
     ACTION_ALLOWED_VALUES: ClassVar[tuple[str]]
-    QUEUE_SELECTORS: ClassVar[tuple]
+    QUEUE_SELECTORS: ClassVar[tuple] = ()
 
     @action.validator
     def validate_action(self, attribute: attrs.Attribute, value: Any):
@@ -128,16 +131,10 @@ def parse(
     ) -> tuple[RequestId, str]:
 
         actual_cls = cls
-        subclasses = set()
         queue: str | None = None
         dict_repr = {}
         fragments = {}
 
-        def init_subclasses(cur_cls: type[RequestId] = RequestId):
-            for subclass in cur_cls.__subclasses__():
-                subclasses.add(subclass)
-                init_subclasses(subclass)
-
         try:
             # try to parse ID as key=value pairs (newly introduced format)
             fragments = dict(urllib.parse.parse_qsl(request_id))
@@ -151,12 +148,7 @@ def init_subclasses(cur_cls: type[RequestId] = RequestId):
 
                 match: re.Match | None = None
 
-                if cls is RequestId:
-                    init_subclasses()
-                else:
-                    subclasses = (cls,)
-
-                for subclass in subclasses:
+                for subclass in REQUEST_ID_SUBCLASSES if cls is RequestId else (cls,):
                     for pattern in subclass.LEGACY_FORMAT_PATTERNS:
                         match = re.match(pattern, request_id)
                         if match:
@@ -169,14 +161,9 @@ def init_subclasses(cur_cls: type[RequestId] = RequestId):
                         f"Unable to parse request ID: {escape(request_id)!r}"
                     )
 
-                queue = ACTION_TO_QUEUE.get(
+                queue = ACTION_TO_QUEUE[
                     actual_cls.QUEUE_SELECTORS[0]
-                )  # each selector match the same queue
-                if not queue:
-                    raise ImproperlyConfigured(
-                        "Job ID class must be set in the related queue config"
-                    )
-
+                ]  # each selector match the same queue
                 fragments = match.groupdict()
                 # "." was replaced with "@" in previous format
                 if "format" in fragments:
@@ -201,9 +188,6 @@ def init_subclasses(cur_cls: type[RequestId] = RequestId):
             result = actual_cls(**dict_repr)
 
             return (result, queue)
-
-        except ImproperlyConfigured:
-            raise
         except Exception as ex:
             raise IncorrectRequestIdError from ex
 

cvat/apps/redis_handler/serializers.py

@@ -10,6 +10,7 @@
 from uuid import UUID
 
 import rq.defaults as rq_defaults
+from django.db.models import TextChoices
 from django.utils import timezone
 from drf_spectacular.utils import extend_schema_field
 from rest_framework import serializers
@@ -23,10 +24,10 @@
 from cvat.apps.engine.utils import parse_exception_message
 from cvat.apps.lambda_manager.rq import LambdaRQMeta
 from cvat.apps.redis_handler.rq import CustomRQJob, RequestId
-from django.db.models import TextChoices
 
 slogger = ServerLogManager(__name__)
 
+
 class RequestStatus(TextChoices):
     QUEUED = "queued"
     STARTED = "started"

cvat/apps/redis_handler/views.py

@@ -23,12 +23,11 @@
     NonModelSimpleFilter,
 )
 from cvat.apps.engine.log import ServerLogManager
-from cvat.apps.redis_handler.serializers import RequestStatus
 from cvat.apps.engine.rq import is_rq_job_owner
 from cvat.apps.engine.types import ExtendedRequest
 from cvat.apps.redis_handler.apps import ACTION_TO_QUEUE
 from cvat.apps.redis_handler.rq import CustomRQJob, RequestId
-from cvat.apps.redis_handler.serializers import RequestSerializer
+from cvat.apps.redis_handler.serializers import RequestSerializer, RequestStatus
 
 slogger = ServerLogManager(__name__)