NDOJ/judge/views/tasks.py

import json
from functools import partial
from uuid import UUID

from celery.result import AsyncResult
from django.core.exceptions import PermissionDenied
from django.http import (
    Http404,
    HttpResponseBadRequest,
    HttpResponseRedirect,
    JsonResponse,
)
from django.shortcuts import render
from django.urls import reverse
from django.utils.http import is_safe_url

from judge.tasks import failure, progress, success
from judge.utils.celery import redirect_to_task_status
from judge.utils.views import short_circuit_middleware


def get_task_status(task_id):
    result = AsyncResult(task_id)
    info = result.result
    if result.state == "PROGRESS":
        return {
            "code": "PROGRESS",
            "done": info["done"],
            "total": info["total"],
            "stage": info["stage"],
        }
    elif result.state == "SUCCESS":
        return {"code": "SUCCESS"}
    elif result.state == "FAILURE":
        return {"code": "FAILURE", "error": str(info)}
    else:
        return {"code": "WORKING"}


def task_status(request, task_id):
    try:
        UUID(task_id)
    except ValueError:
        raise Http404()

    redirect = request.GET.get("redirect")
    if not is_safe_url(redirect, allowed_hosts={request.get_host()}):
        redirect = None

    status = get_task_status(task_id)
    if status["code"] == "SUCCESS" and redirect:
        return HttpResponseRedirect(redirect)

    return render(
        request,
        "task_status.html",
        {
            "task_id": task_id,
            "task_status": json.dumps(status),
            "message": request.GET.get("message", ""),
            "redirect": redirect or "",
        },
    )


@short_circuit_middleware
def task_status_ajax(request):
    if "id" not in request.GET:
        return HttpResponseBadRequest(
            'Need to pass GET parameter "id"', content_type="text/plain"
        )
    return JsonResponse(get_task_status(request.GET["id"]))


def demo_task(request, task, message):
    if not request.user.is_superuser:
        raise PermissionDenied()
    result = task.delay()
    return redirect_to_task_status(result, message=message, redirect=reverse("home"))


demo_success = partial(
    demo_task, task=success, message="Running example task that succeeds..."
)
demo_failure = partial(
    demo_task, task=failure, message="Running example task that fails..."
)
demo_progress = partial(
    demo_task, task=progress, message="Running example task that waits 10 seconds..."
)