queues/app/routers/take.py
Egor Matveev 9dfc115224
All checks were successful
Deploy Dev / Build (pull_request) Successful in 34s
Deploy Dev / Push (pull_request) Successful in 13s
Deploy Dev / Deploy dev (pull_request) Successful in 11s
fix
2024-12-28 13:19:46 +03:00

36 lines
954 B
Python

import asyncio
import collections
import fastapi
import pydantic
import typing
from app.storage.mongo import tasks
from app.utils.configurator import configurator
DEFAULT_RETRY_AFTER = 0.2
locks = collections.defaultdict(asyncio.Lock)
router = fastapi.APIRouter()
class Task(pydantic.BaseModel):
id: str
attempt: int
payload: dict
class Response(pydantic.BaseModel):
task: Task|None
retry_after: float
@router.get('/api/v1/take')
async def execute(queue: typing.Annotated[str, fastapi.Header()]) -> Response:
async with locks[queue]:
task = await tasks.take_task(queue)
if not task:
return Response(task=None)
retry_after_config = configurator.get_config('retry_after')
retry_after = retry_after_config.get(queue) or retry_after_config.get('default') or DEFAULT_RETRY_AFTER
return Response(task=Task(id=str(task._id), attempt=task.attempts, payload=task.payload), retry_after=retry_after)