queues/app/routers/take.py

import fastapi
import pydantic
import typing

from app.storage.mongo import tasks


router = fastapi.APIRouter()


class Task(pydantic.BaseModel):
    id: str
    attempt: int
    payload: dict


class Response(pydantic.BaseModel):
    id: str
    attempt: int
    payload: dict
    task: Task|None


@router.get('/api/v1/take', responses={404: {'description': 'Not found'}})
async def execute(queue: typing.Annotated[str, fastapi.Header()]) -> Response:
    task = await tasks.take_task(queue)
    if not task:
        raise fastapi.HTTPException(404)
    return Response(id=str(task._id), attempt=task.attempts, payload=task.payload, task=Task(id=str(task._id), attempt=task.attempts, payload=task.payload))