85 lines
2.1 KiB
Python
85 lines
2.1 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
from typing import Any, Dict, Optional
|
|
|
|
from redis.asyncio import Redis
|
|
|
|
REDIS_URL = os.getenv("REDIS_URL", "redis://redis:6379/0")
|
|
JOB_KEY_PREFIX = "daarion:jobs"
|
|
QUEUE_KEY = "daarion:jobs:queue"
|
|
JOB_TTL_SECONDS = int(os.getenv("DAARION_JOB_TTL_SECONDS", str(72 * 3600)))
|
|
|
|
_redis: Optional[Redis] = None
|
|
|
|
|
|
def _job_key(job_id: str) -> str:
|
|
return f"{JOB_KEY_PREFIX}:{job_id}"
|
|
|
|
|
|
async def redis_client() -> Redis:
|
|
global _redis
|
|
if _redis is None:
|
|
_redis = Redis.from_url(REDIS_URL, decode_responses=True)
|
|
return _redis
|
|
|
|
|
|
async def close_redis() -> None:
|
|
global _redis
|
|
if _redis is not None:
|
|
await _redis.close()
|
|
_redis = None
|
|
|
|
|
|
async def create_job(job_id: str, payload: Dict[str, Any]) -> None:
|
|
r = await redis_client()
|
|
key = _job_key(job_id)
|
|
await r.set(key, json.dumps(payload, ensure_ascii=False), ex=JOB_TTL_SECONDS)
|
|
|
|
|
|
async def get_job(job_id: str) -> Optional[Dict[str, Any]]:
|
|
r = await redis_client()
|
|
raw = await r.get(_job_key(job_id))
|
|
if not raw:
|
|
return None
|
|
try:
|
|
return json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
return None
|
|
|
|
|
|
async def update_job(job_id: str, patch: Dict[str, Any]) -> Optional[Dict[str, Any]]:
|
|
current = await get_job(job_id)
|
|
if not current:
|
|
return None
|
|
current.update(patch)
|
|
await create_job(job_id, current)
|
|
return current
|
|
|
|
|
|
async def enqueue_job(job_id: str) -> None:
|
|
r = await redis_client()
|
|
await r.lpush(QUEUE_KEY, job_id)
|
|
|
|
|
|
async def dequeue_job(block_seconds: int = 5) -> Optional[str]:
|
|
r = await redis_client()
|
|
result = await r.brpop(QUEUE_KEY, timeout=block_seconds)
|
|
if not result:
|
|
return None
|
|
_, job_id = result
|
|
return job_id
|
|
|
|
|
|
async def wait_for_redis(timeout_seconds: int = 30) -> None:
|
|
deadline = asyncio.get_running_loop().time() + timeout_seconds
|
|
while True:
|
|
try:
|
|
r = await redis_client()
|
|
await r.ping()
|
|
return
|
|
except Exception:
|
|
if asyncio.get_running_loop().time() >= deadline:
|
|
raise
|
|
await asyncio.sleep(1)
|