refactor

Author: valentijnscholten

dojo/api_v2/views.py

@@ -35,7 +35,6 @@
 from rest_framework.parsers import MultiPartParser
 from rest_framework.permissions import DjangoModelPermissions, IsAuthenticated
 from rest_framework.response import Response
-from rest_framework.views import APIView
 
 import dojo.finding.helper as finding_helper
 import dojo.jira_link.helper as jira_helper
@@ -3266,20 +3265,21 @@ def get_queryset(self):
         return System_Settings.objects.all().order_by("id")
 
 
-@extend_schema(
-    responses=serializers.CeleryStatusSerializer,
-    summary="Get Celery worker and queue status",
-    description=(
-        "Returns Celery worker liveness, pending queue length, and the active task "
-        "timeout/expiry configuration. Uses the Celery control channel (pidbox) for "
-        "worker status so it works correctly even when the task queue is clogged."
-    ),
-)
-class CeleryStatusView(APIView):
+class CeleryViewSet(viewsets.ViewSet):
     permission_classes = (permissions.IsSuperUser, DjangoModelPermissions)
     queryset = System_Settings.objects.none()
 
-    def get(self, request):
+    @extend_schema(
+        responses=serializers.CeleryStatusSerializer,
+        summary="Get Celery worker and queue status",
+        description=(
+            "Returns Celery worker liveness, pending queue length, and the active task "
+            "timeout/expiry configuration. Uses the Celery control channel (pidbox) for "
+            "worker status so it works correctly even when the task queue is clogged."
+        ),
+    )
+    @action(detail=False, methods=["get"], url_path="status")
+    def status(self, request):
         queue_length = get_celery_queue_length()
         data = {
             "worker_status": get_celery_worker_status(),
@@ -3291,56 +3291,44 @@ def get(self, request):
         }
         return Response(serializers.CeleryStatusSerializer(data).data)
 
-
-@extend_schema(
-    request=None,
-    responses={200: {"type": "object", "properties": {"purged": {"type": "integer"}}}},
-    summary="Purge all pending Celery tasks from the queue",
-    description=(
-        "Removes all pending tasks from the default Celery queue. Tasks already being "
-        "executed by workers are not affected. Note: if deduplication tasks were queued, "
-        "you may need to re-run deduplication manually via `python manage.py dedupe`."
-    ),
-)
-class CeleryQueuePurgeView(APIView):
-    permission_classes = (permissions.IsSuperUser, DjangoModelPermissions)
-    queryset = System_Settings.objects.none()
-
-    def post(self, request):
+    @extend_schema(
+        request=None,
+        responses={200: {"type": "object", "properties": {"purged": {"type": "integer"}}}},
+        summary="Purge all pending Celery tasks from the queue",
+        description=(
+            "Removes all pending tasks from the default Celery queue. Tasks already being "
+            "executed by workers are not affected. Note: if deduplication tasks were queued, "
+            "you may need to re-run deduplication manually via `python manage.py dedupe`."
+        ),
+    )
+    @action(detail=False, methods=["post"], url_path="queue/purge")
+    def queue_purge(self, request):
         purged = purge_celery_queue()
         return Response({"purged": purged})
 
-
-@extend_schema(
-    responses=serializers.CeleryQueueTaskDetailSerializer(many=True),
-    summary="Get per-task breakdown of the Celery queue",
-    description=(
-        "Scans every message in the queue (O(N)) and returns task name, count, and "
-        "oldest/newest queue positions. May be slow for large queues."
-    ),
-)
-class CeleryQueueDetailsView(APIView):
-    permission_classes = (permissions.IsSuperUser, DjangoModelPermissions)
-    queryset = System_Settings.objects.none()
-
-    def get(self, request):
+    @extend_schema(
+        responses=serializers.CeleryQueueTaskDetailSerializer(many=True),
+        summary="Get per-task breakdown of the Celery queue",
+        description=(
+            "Scans every message in the queue (O(N)) and returns task name, count, and "
+            "oldest/newest queue positions. May be slow for large queues."
+        ),
+    )
+    @action(detail=False, methods=["get"], url_path="queue/details")
+    def queue_details(self, request):
         details = get_celery_queue_details()
         if details is None:
             return Response({"error": "Unable to read queue details."}, status=503)
         return Response(serializers.CeleryQueueTaskDetailSerializer(details, many=True).data)
 
-
-@extend_schema(
-    request={"application/json": {"type": "object", "properties": {"task_name": {"type": "string"}}, "required": ["task_name"]}},
-    responses={200: {"type": "object", "properties": {"purged": {"type": "integer"}}}},
-    summary="Purge all queued tasks with a given task name",
-    description="Removes all pending tasks matching the given task name from the default Celery queue.",
-)
-class CeleryQueueTaskPurgeView(APIView):
-    permission_classes = (permissions.IsSuperUser, DjangoModelPermissions)
-    queryset = System_Settings.objects.none()
-
-    def post(self, request):
+    @extend_schema(
+        request={"application/json": {"type": "object", "properties": {"task_name": {"type": "string"}}, "required": ["task_name"]}},
+        responses={200: {"type": "object", "properties": {"purged": {"type": "integer"}}}},
+        summary="Purge all queued tasks with a given task name",
+        description="Removes all pending tasks matching the given task name from the default Celery queue.",
+    )
+    @action(detail=False, methods=["post"], url_path="queue/task/purge")
+    def queue_task_purge(self, request):
         task_name = request.data.get("task_name", "").strip()
         if not task_name:
             return Response({"error": "task_name is required."}, status=400)

dojo/system_settings/urls.py

@@ -9,8 +9,8 @@
         name="system_settings",
     ),
     re_path(
-        r"^system_status$",
-        views.SystemStatusView.as_view(),
-        name="system_status",
+        r"^celery_status$",
+        views.CeleryStatusView.as_view(),
+        name="celery_status",
     ),
 ]

dojo/system_settings/views.py

@@ -120,12 +120,12 @@ def post(
         return render(request, self.get_template(), context)
 
 
-class SystemStatusView(View):
+class CeleryStatusView(View):
     def get(
         self,
         request: HttpRequest,
     ) -> HttpResponse:
         if not request.user.is_superuser:
             raise PermissionDenied
-        add_breadcrumb(title="System status", top_level=False, request=request)
-        return render(request, "dojo/system_status.html")
+        add_breadcrumb(title="Celery status", top_level=False, request=request)
+        return render(request, "dojo/celery_status.html")

dojo/templates/base.html

@@ -614,8 +614,8 @@
                                                         </a>
                                                     </li>
                                                     <li>
-                                                        <a href="{% url 'system_status' %}">
-                                                            {% trans "System Status" %}
+                                                        <a href="{% url 'celery_status' %}">
+                                                            {% trans "Celery Status" %}
                                                         </a>
                                                     </li>
                                                 {% endif %}

dojo/templates/dojo/system_status.html dojo/templates/dojo/celery_status.htmldojo/templates/dojo/system_status.html renamed to dojo/templates/dojo/celery_status.html

@@ -4,8 +4,9 @@
 {% block content %}
     {{ block.super }}
     <div class="row">
-        <h3>System Status</h3>
+        <h3>Celery Status</h3>
         <br>
+        {% if show_celery_status %}
         <div id="celery-status-panel" class="col-md-5">
             <div class="panel panel-default">
                 <div class="panel-heading">
@@ -103,17 +104,19 @@ <h4>Celery Settings</h4>
                 </div>
             </div>
         </div>
+        {% endif %}
     </div>
 {% endblock content %}
 
 {% block postscript %}
     {{ block.super }}
+    {% if show_celery_status %}
     <script>
         (function() {
-            var purgeUrl = "{% url 'celery_queue_purge_api' %}";
-            var statusUrl = "{% url 'celery_status_api' %}";
-            var detailsUrl = "{% url 'celery_queue_details_api' %}";
-            var taskPurgeUrl = "{% url 'celery_queue_task_purge_api' %}";
+            var purgeUrl = "{% url 'celery-queue-purge' %}";
+            var statusUrl = "{% url 'celery-status' %}";
+            var detailsUrl = "{% url 'celery-queue-details' %}";
+            var taskPurgeUrl = "{% url 'celery-queue-task-purge' %}";
 
             function setStatusBadge(loadingId, badgeId, isOk, okText, failText) {
                 $('#' + loadingId).hide();
@@ -302,4 +305,5 @@ <h4>Celery Settings</h4>
             });
         })();
     </script>
+    {% endif %}
 {% endblock %}

dojo/urls.py

@@ -15,10 +15,7 @@
     AnnouncementViewSet,
     AppAnalysisViewSet,
     BurpRawRequestResponseViewSet,
-    CeleryQueueDetailsView,
-    CeleryQueuePurgeView,
-    CeleryQueueTaskPurgeView,
-    CeleryStatusView,
+    CeleryViewSet,
     ConfigurationPermissionViewSet,
     CredentialsMappingViewSet,
     CredentialsViewSet,
@@ -198,6 +195,7 @@
 else:
     v2_api.register(r"endpoints", EndPointViewSet, basename="endpoint")
     v2_api.register(r"endpoint_status", EndpointStatusViewSet, basename="endpoint_status")
+v2_api.register(r"celery", CeleryViewSet, basename="celery")
 # V3
 add_asset_urls(v2_api)
 add_organization_urls(v2_api)
@@ -246,10 +244,6 @@
     #  Django Rest Framework API v2
     re_path(r"^{}api/v2/".format(get_system_setting("url_prefix")), include(v2_api.urls)),
     re_path(r"^{}api/v2/user_profile/".format(get_system_setting("url_prefix")), UserProfileView.as_view(), name="user_profile"),
-    re_path(r"^{}api/v2/celery/status/$".format(get_system_setting("url_prefix")), CeleryStatusView.as_view(), name="celery_status_api"),
-    re_path(r"^{}api/v2/celery/queue/purge/$".format(get_system_setting("url_prefix")), CeleryQueuePurgeView.as_view(), name="celery_queue_purge_api"),
-    re_path(r"^{}api/v2/celery/queue/details/$".format(get_system_setting("url_prefix")), CeleryQueueDetailsView.as_view(), name="celery_queue_details_api"),
-    re_path(r"^{}api/v2/celery/queue/task/purge/$".format(get_system_setting("url_prefix")), CeleryQueueTaskPurgeView.as_view(), name="celery_queue_task_purge_api"),
 ]
 
 if hasattr(settings, "API_TOKENS_ENABLED") and hasattr(settings, "API_TOKEN_AUTH_ENDPOINT_ENABLED"):