portier/source/concierge/api.py

112 lines
2.6 KiB
Python

from ninja import Router, Schema
from ninja.errors import HttpError
from concierge import models
from typing import List
from django.db import transaction
from django.shortcuts import get_object_or_404
from django.utils.timezone import now
from django.views.decorators.csrf import csrf_exempt
import json
router = Router()
class ClaimReference(Schema):
uuid: str
class AvailableTask(Schema):
uuid: str
type: str
class HeartbeatResponse(Schema):
success: bool = True
claims: List[ClaimReference]
available: List[AvailableTask]
class ClaimResponse(Schema):
uuid: str
type: str
configuration: dict
class ReleaseResponse(Schema):
success: bool = True
uuid: str
type: str
@router.post('/heartbeat/{identity}', response=HeartbeatResponse)
@csrf_exempt
def receive_heartbeat(request, identity: str):
try:
id = models.Identity.objects.get(identity=identity)
except models.Identity.DoesNotExist:
raise HttpError(403, "identity unknown")
# update heartbeat
id.heartbeat = now()
id.save()
# get current claims and available tasks
claims = models.Task.objects.filter(claimed_by=id).all()
available = models.Task.objects.filter(claimed_by=None).all()
return {
'success': True,
'claims': [{'uuid': str(o.uuid)} for o in list(claims)],
'available': [{'uuid': str(o.uuid), 'type': o.type} for o in list(available)],
}
@router.post('/claim/{identity}/{task_uuid}', response=ClaimResponse)
@csrf_exempt
def claim_task(request, identity: str, task_uuid: str):
try:
id = models.Identity.objects.get(identity=identity)
except models.Identity.DoesNotExist:
raise HttpError(403, "identity unknown")
with transaction.atomic():
task = get_object_or_404(models.Task, uuid=task_uuid)
if task.claimed_by:
raise HttpError(423, "task already claimed")
task.claimed_by = id
task.save()
return {
'success': True,
'uuid': task.uuid,
'type': task.type,
'configuration': json.loads(task.configuration)
}
@router.post('/release/{identity}/{task_uuid}')
@csrf_exempt
def release_task(request, identity: str, task_uuid: str):
try:
id = models.Identity.objects.get(identity=identity)
except models.Identity.DoesNotExist:
raise HttpError(403, "identity unknown")
with transaction.atomic():
task = get_object_or_404(models.Task, uuid=task_uuid)
if task.claimed_by != id:
raise HttpError(403, "task not claimed by this identity")
task.claimed_by = None
task.save()
return {
'success': True,
'uuid': task.uuid,
'type': task.type,
}