36 lines
954 B
Python
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)
|