创建中心任务的管理系统

tasks/views.py

@@ -0,0 +1,105 @@
+from django.shortcuts import get_object_or_404
+from django.utils import timezone
+from rest_framework import viewsets, status
+from rest_framework.decorators import action
+from rest_framework.permissions import IsAuthenticated
+from rest_framework.response import Response
+from rest_framework.parsers import MultiPartParser, FormParser
+from .models import Client, Task, TaskResult
+from .serializers import (
+    ClientSerializer,
+    TaskSerializer,
+    TaskResultSerializer,
+    TaskClaimSerializer,
+    TaskStartSerializer,
+    TaskCompleteSerializer
+)
+
+class ClientViewSet(viewsets.ModelViewSet):
+    queryset = Client.objects.all()
+    serializer_class = ClientSerializer
+    permission_classes = [IsAuthenticated]
+
+class TaskViewSet(viewsets.ModelViewSet):
+    queryset = Task.objects.all()
+    serializer_class = TaskSerializer
+    permission_classes = [IsAuthenticated]
+
+    @action(detail=False, methods=['post'], serializer_class=TaskClaimSerializer)
+    def claim(self, request):
+        """Client claims an available task"""
+        serializer = self.get_serializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+        client_name = serializer.validated_data['client_name']
+
+        # Try to find a pending task, either assigned to this client or unassigned
+        with self.queryset.model.objects._base_manager._db.transaction.atomic():
+            # First try to find tasks assigned to this client
+            task = Task.objects.filter(
+                status='pending',
+                client_name=client_name
+            ).select_for_update().first()
+
+            # If no task assigned to this client, try to find unassigned tasks
+            if not task:
+                task = Task.objects.filter(
+                    status='pending',
+                    client_name__isnull=True
+                ).select_for_update().first()
+
+            if not task:
+                return Response({'detail': 'No available tasks'}, status=status.HTTP_404_NOT_FOUND)
+
+            # Assign the task to the client
+            task.status = 'assigned'
+            task.assigned_to = client_name
+            task.save()
+
+        return Response(TaskSerializer(task).data, status=status.HTTP_200_OK)
+
+    @action(detail=True, methods=['post'], serializer_class=TaskStartSerializer)
+    def start(self, request, pk=None):
+        """Client starts a task"""
+        task = self.get_object()
+        if task.status != 'assigned':
+            return Response({'detail': 'Task is not assigned'}, status=status.HTTP_400_BAD_REQUEST)
+
+        task.status = 'running'
+        task.started_at = timezone.now()
+        task.save()
+
+        return Response(TaskSerializer(task).data, status=status.HTTP_200_OK)
+
+    @action(detail=True, methods=['post'], serializer_class=TaskCompleteSerializer)
+    def complete(self, request, pk=None):
+        """Client completes a task"""
+        task = self.get_object()
+        serializer = self.get_serializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+
+        if task.status != 'running':
+            return Response({'detail': 'Task is not running'}, status=status.HTTP_400_BAD_REQUEST)
+
+        task.status = serializer.validated_data['status']
+        task.completed_at = timezone.now()
+        task.save()
+
+        return Response(TaskSerializer(task).data, status=status.HTTP_200_OK)
+
+class TaskResultViewSet(viewsets.ModelViewSet):
+    queryset = TaskResult.objects.all()
+    serializer_class = TaskResultSerializer
+    permission_classes = [IsAuthenticated]
+    parser_classes = [MultiPartParser, FormParser]
+
+    @action(detail=True, methods=['get'])
+    def download(self, request, pk=None):
+        """Download task result file"""
+        task_result = self.get_object()
+        if not task_result.result_file:
+            return Response({'detail': 'No file available'}, status=status.HTTP_404_NOT_FOUND)
+
+        # Use Django's built-in FileResponse for downloading
+        from django.http import FileResponse
+        return FileResponse(task_result.result_file.open('rb'), as_attachment=True)
+