106 lines
4.0 KiB
Python
106 lines
4.0 KiB
Python
|
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)
|
||
|
|
|