43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
import bson
|
|
import datetime
|
|
import pydantic
|
|
|
|
from storage.mongo import database
|
|
from 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) -> Task|None:
|
|
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
|