initial

app/routers/finish.py

@@ -0,0 +1,19 @@
+import bson
+import fastapi
+import pydantic
+
+from app.storage.mongo import tasks
+
+
+class RequestBody(pydantic.BaseModel):
+    id: str
+
+
+router = fastapi.APIRouter()
+
+
+@router.post('/api/v1/finish', status_code=fastapi.status.HTTP_202_ACCEPTED, responses={'404': {'description': 'Not found'}})
+async def execute(body: RequestBody):
+    if await tasks.finish_task(bson.ObjectId(body.id)):
+        return
+    raise fastapi.HTTPException(404)

app/routers/put.py

@@ -0,0 +1,30 @@
+import datetime
+import fastapi
+import pydantic
+import typing
+
+from app.storage.mongo import tasks
+from app.utils import time
+
+
+class RequestBody(pydantic.BaseModel):
+    payload: dict
+    seconds_to_execute: int
+    delay: int|None = None
+
+
+router = fastapi.APIRouter()
+
+
+@router.post('/api/v1/put', status_code=fastapi.status.HTTP_202_ACCEPTED)
+async def execute(body: RequestBody, queue: typing.Annotated[str, fastapi.Header()]):
+    now = time.now()
+    await tasks.add_task(
+        task=tasks.Task(
+            queue=queue, 
+            payload=body.payload, 
+            put_at=now, 
+            available_from=(now + datetime.timedelta(seconds=body.delay)) if body.delay else now,
+            seconds_to_execute=body.seconds_to_execute,
+        ),
+    )

app/routers/take.py

@@ -0,0 +1,22 @@
+import fastapi
+import pydantic
+import typing
+
+from app.storage.mongo import tasks
+
+
+router = fastapi.APIRouter()
+
+
+class Response(pydantic.BaseModel):
+    id: str
+    attempt: int
+    payload: dict
+
+
+@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)

app/storage/mongo/__init__.py

@@ -0,0 +1,21 @@
+import os
+import motor
+import motor.motor_asyncio
+import pymongo
+
+
+MONGO_HOST = os.getenv('MONGO_HOST', 'localhost')
+MONGO_PASSWORD = os.getenv('MONGO_PASSWORD', 'password')
+
+CONNECTION_STRING = f'mongodb://mongo:{MONGO_PASSWORD}@{MONGO_HOST}:27017/'
+
+
+database: 'motor.MotorDatabase' = motor.motor_asyncio.AsyncIOMotorClient(CONNECTION_STRING).queues
+
+def create_indexes():
+    client = pymongo.MongoClient(CONNECTION_STRING)
+    database = client.get_database('queues')
+    database.get_collection('tasks').create_index([
+        ('queue', 1), 
+        ('available_from', 1),
+    ])

app/storage/mongo/tasks.py

@@ -0,0 +1,44 @@
+import asyncio
+import bson
+import datetime
+import pydantic
+import typing
+
+from app.storage.mongo import database
+from app.utils import time
+from bson import codec_options
+
+
+collection = database.get_collection("tasks", codec_options=codec_options.CodecOptions(tz_aware=True))
+
+
+class Task(pydantic.BaseModel):
+    queue: str
+    payload: dict
+    put_at: pydantic.AwareDatetime
+    available_from: pydantic.AwareDatetime
+    seconds_to_execute: int
+    _id: bson.ObjectId|None = None
+    taken_at: pydantic.AwareDatetime|None = None
+    attempts: int = 0
+
+
+async def add_task(task: Task) -> str:
+    result = await collection.insert_one(task.model_dump())
+    return result.inserted_id
+
+
+async def take_task(queue: str) -> typing.Optional[Task]:
+    now = time.now()
+    async for raw_task in collection.find({'queue': queue, 'available_from': {'$lte': now}}):
+        task = Task.model_validate(raw_task)
+        task._id = raw_task['_id']
+        if not task.taken_at or (task.taken_at + datetime.timedelta(seconds=task.seconds_to_execute)) < now:
+            await collection.update_one({'_id': task._id}, {'$set': {'taken_at': now, 'attempts': task.attempts + 1}})
+            return task
+    return None
+
+
+async def finish_task(id: bson.ObjectId) -> bool:
+    result = await collection.delete_one({'_id': id})
+    return result

app/utils/time.py

@@ -0,0 +1,4 @@
+import datetime
+
+
+now = lambda: datetime.datetime.now(datetime.UTC)