7848 Add RQ API (#17938)

* 7848 Add Background Tasks (RQ) to API * 7848 Tasks * 7848 cleanup * 7848 add worker support * 7848 switch to APIView * 7848 Task detail view * 7848 Task enqueue, requeue, stop * 7848 Task enqueue, requeue, stop * 7848 Task enqueue, requeue, stop * 7848 tests * 7848 tests * 7848 OpenAPI doc generation * 7848 OpenAPI doc generation * 7848 review changes * 7848 viewset * 7848 viewset * 7848 fix tests * 7848 more viewsets * 7848 fix docstring * 7848 review comments * 7848 review comments - get all tasks * 7848 queue detail view * 7848 cleanup * 7848 cleanup * 7848 cleanup * 7848 cleanup * Rename viewsets for consistency w/serializers * Misc cleanup * 7848 review changes * 7848 review changes * 7848 add test * 7848 queue detail view * 7848 fix tests * 7848 fix the spectacular test failure * 7848 fix the spectacular test failure * Misc cleanup --------- Co-authored-by: Jeremy Stretch <jstretch@netboxlabs.com>
2024-11-26 07:01:06 -08:00
parent 9b4b56febc
commit a24576f126
9 changed files with 612 additions and 99 deletions
--- a/netbox/core/api/views.py
+++ b/netbox/core/api/views.py
@@ -1,5 +1,8 @@
+from django.http import Http404, HttpResponse
 from django.shortcuts import get_object_or_404
 from django.utils.translation import gettext_lazy as _
+from drf_spectacular.types import OpenApiTypes
+from drf_spectacular.utils import extend_schema
 from rest_framework.decorators import action
 from rest_framework.exceptions import PermissionDenied
 from rest_framework.response import Response
@@ -10,8 +13,17 @@ from core import filtersets
 from core.choices import DataSourceStatusChoices
 from core.jobs import SyncDataSourceJob
 from core.models import *
+from core.utils import delete_rq_job, enqueue_rq_job, get_rq_jobs, requeue_rq_job, stop_rq_job
+from django_rq.queues import get_redis_connection
+from django_rq.utils import get_statistics
+from django_rq.settings import QUEUES_LIST
 from netbox.api.metadata import ContentTypeMetadata
+from netbox.api.pagination import LimitOffsetListPagination
 from netbox.api.viewsets import NetBoxModelViewSet, NetBoxReadOnlyModelViewSet
+from rest_framework import viewsets
+from rest_framework.permissions import IsAdminUser
+from rq.job import Job as RQ_Job
+from rq.worker import Worker
 from . import serializers


@@ -71,3 +83,152 @@ class ObjectChangeViewSet(ReadOnlyModelViewSet):
    queryset = ObjectChange.objects.valid_models()
    serializer_class = serializers.ObjectChangeSerializer
    filterset_class = filtersets.ObjectChangeFilterSet
+
+
+class BaseRQViewSet(viewsets.ViewSet):
+    """
+    Base class for RQ view sets. Provides a list() method. Subclasses must implement get_data().
+    """
+    permission_classes = [IsAdminUser]
+    serializer_class = None
+
+    def get_data(self):
+        raise NotImplementedError()
+
+    @extend_schema(responses={200: OpenApiTypes.OBJECT})
+    def list(self, request):
+        data = self.get_data()
+        paginator = LimitOffsetListPagination()
+        data = paginator.paginate_list(data, request)
+
+        serializer = self.serializer_class(data, many=True, context={'request': request})
+        return paginator.get_paginated_response(serializer.data)
+
+    def get_serializer(self, *args, **kwargs):
+        """
+        Return the serializer instance that should be used for validating and
+        deserializing input, and for serializing output.
+        """
+        serializer_class = self.get_serializer_class()
+        kwargs['context'] = self.get_serializer_context()
+        return serializer_class(*args, **kwargs)
+
+
+class BackgroundQueueViewSet(BaseRQViewSet):
+    """
+    Retrieve a list of RQ Queues.
+    Note: Queue names are not URL safe so not returning a detail view.
+    """
+    serializer_class = serializers.BackgroundQueueSerializer
+    lookup_field = 'name'
+    lookup_value_regex = r'[\w.@+-]+'
+
+    def get_view_name(self):
+        return "Background Queues"
+
+    def get_data(self):
+        return get_statistics(run_maintenance_tasks=True)["queues"]
+
+    @extend_schema(responses={200: OpenApiTypes.OBJECT})
+    def retrieve(self, request, name):
+        data = self.get_data()
+        if not data:
+            raise Http404
+
+        for queue in data:
+            if queue['name'] == name:
+                serializer = self.serializer_class(queue, context={'request': request})
+                return Response(serializer.data)
+
+        raise Http404
+
+
+class BackgroundWorkerViewSet(BaseRQViewSet):
+    """
+    Retrieve a list of RQ Workers.
+    """
+    serializer_class = serializers.BackgroundWorkerSerializer
+    lookup_field = 'name'
+
+    def get_view_name(self):
+        return "Background Workers"
+
+    def get_data(self):
+        config = QUEUES_LIST[0]
+        return Worker.all(get_redis_connection(config['connection_config']))
+
+    def retrieve(self, request, name):
+        # all the RQ queues should use the same connection
+        config = QUEUES_LIST[0]
+        workers = Worker.all(get_redis_connection(config['connection_config']))
+        worker = next((item for item in workers if item.name == name), None)
+        if not worker:
+            raise Http404
+
+        serializer = serializers.BackgroundWorkerSerializer(worker, context={'request': request})
+        return Response(serializer.data)
+
+
+class BackgroundTaskViewSet(BaseRQViewSet):
+    """
+    Retrieve a list of RQ Tasks.
+    """
+    serializer_class = serializers.BackgroundTaskSerializer
+
+    def get_view_name(self):
+        return "Background Tasks"
+
+    def get_data(self):
+        return get_rq_jobs()
+
+    def get_task_from_id(self, task_id):
+        config = QUEUES_LIST[0]
+        task = RQ_Job.fetch(task_id, connection=get_redis_connection(config['connection_config']))
+        if not task:
+            raise Http404
+
+        return task
+
+    @extend_schema(responses={200: OpenApiTypes.OBJECT})
+    def retrieve(self, request, pk):
+        """
+        Retrieve the details of the specified RQ Task.
+        """
+        task = self.get_task_from_id(pk)
+        serializer = self.serializer_class(task, context={'request': request})
+        return Response(serializer.data)
+
+    @action(methods=["POST"], detail=True)
+    def delete(self, request, pk):
+        """
+        Delete the specified RQ Task.
+        """
+        delete_rq_job(pk)
+        return HttpResponse(status=200)
+
+    @action(methods=["POST"], detail=True)
+    def requeue(self, request, pk):
+        """
+        Requeues the specified RQ Task.
+        """
+        requeue_rq_job(pk)
+        return HttpResponse(status=200)
+
+    @action(methods=["POST"], detail=True)
+    def enqueue(self, request, pk):
+        """
+        Enqueues the specified RQ Task.
+        """
+        enqueue_rq_job(pk)
+        return HttpResponse(status=200)
+
+    @action(methods=["POST"], detail=True)
+    def stop(self, request, pk):
+        """
+        Stops the specified RQ Task.
+        """
+        stopped_jobs = stop_rq_job(pk)
+        if len(stopped_jobs) == 1:
+            return HttpResponse(status=200)
+        else:
+            return HttpResponse(status=204)