# ~*~ coding: utf-8 ~*~ import uuid import os from django.core.cache import cache from django.shortcuts import get_object_or_404 from django.utils.translation import ugettext as _ from rest_framework import viewsets, generics from rest_framework.views import Response from .hands import IsSuperUser from .models import Task, AdHoc, AdHocRunHistory, CeleryTask from .serializers import TaskSerializer, AdHocSerializer, \ AdHocRunHistorySerializer from .tasks import run_ansible_task class TaskViewSet(viewsets.ModelViewSet): queryset = Task.objects.all() serializer_class = TaskSerializer permission_classes = (IsSuperUser,) class TaskRun(generics.RetrieveAPIView): queryset = Task.objects.all() serializer_class = TaskViewSet permission_classes = (IsSuperUser,) def retrieve(self, request, *args, **kwargs): task = self.get_object() t = run_ansible_task.delay(str(task.id)) return Response({"task": t.id}) class AdHocViewSet(viewsets.ModelViewSet): queryset = AdHoc.objects.all() serializer_class = AdHocSerializer permission_classes = (IsSuperUser,) def get_queryset(self): task_id = self.request.query_params.get('task') if task_id: task = get_object_or_404(Task, id=task_id) self.queryset = self.queryset.filter(task=task) return self.queryset class AdHocRunHistorySet(viewsets.ModelViewSet): queryset = AdHocRunHistory.objects.all() serializer_class = AdHocRunHistorySerializer permission_classes = (IsSuperUser,) def get_queryset(self): task_id = self.request.query_params.get('task') adhoc_id = self.request.query_params.get('adhoc') if task_id: task = get_object_or_404(Task, id=task_id) adhocs = task.adhoc.all() self.queryset = self.queryset.filter(adhoc__in=adhocs) if adhoc_id: adhoc = get_object_or_404(AdHoc, id=adhoc_id) self.queryset = self.queryset.filter(adhoc=adhoc) return self.queryset class CeleryTaskLogApi(generics.RetrieveAPIView): permission_classes = (IsSuperUser,) buff_size = 1024 * 10 end = False queryset = CeleryTask.objects.all() def get(self, request, *args, **kwargs): mark = request.query_params.get("mark") or str(uuid.uuid4()) task = super().get_object() log_path = task.full_log_path if not log_path or not os.path.isfile(log_path): return Response({"data": _("Waiting ...")}, status=203) with open(log_path, 'r') as f: offset = cache.get(mark, 0) f.seek(offset) data = f.read(self.buff_size).replace('\n', '\r\n') mark = str(uuid.uuid4()) cache.set(mark, f.tell(), 5) if data == '' and task.is_finished(): self.end = True return Response({"data": data, 'end': self.end, 'mark': mark})