112 lines
2.6 KiB
Python
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,
|
||
|
|
}
|