119 lines
3.5 KiB
Python
119 lines
3.5 KiB
Python
from typing import Any, Dict, List
|
|
import json
|
|
|
|
from loguru import logger
|
|
from app.core.config import settings
|
|
|
|
try:
|
|
import redis
|
|
except Exception: # pragma: no cover - optional dependency
|
|
redis = None
|
|
|
|
|
|
class InMemoryTaskStore:
|
|
def __init__(self) -> None:
|
|
self._tasks: Dict[str, Dict[str, Any]] = {}
|
|
|
|
def create(self, task_id: str, user_id: str) -> Dict[str, Any]:
|
|
task = {
|
|
"status": "pending",
|
|
"task_id": task_id,
|
|
"progress": 0,
|
|
"user_id": user_id,
|
|
}
|
|
self._tasks[task_id] = task
|
|
return task
|
|
|
|
def get(self, task_id: str) -> Dict[str, Any]:
|
|
return self._tasks.get(task_id, {"status": "not_found"})
|
|
|
|
def list(self) -> List[Dict[str, Any]]:
|
|
return list(self._tasks.values())
|
|
|
|
def update(self, task_id: str, updates: Dict[str, Any]) -> Dict[str, Any]:
|
|
task = self._tasks.get(task_id)
|
|
if not task:
|
|
task = {"status": "pending", "task_id": task_id}
|
|
self._tasks[task_id] = task
|
|
task.update(updates)
|
|
return task
|
|
|
|
|
|
class RedisTaskStore:
|
|
def __init__(self, client: "redis.Redis") -> None:
|
|
self._client = client
|
|
self._index_key = "vigent:tasks:index"
|
|
|
|
def _key(self, task_id: str) -> str:
|
|
return f"vigent:tasks:{task_id}"
|
|
|
|
def create(self, task_id: str, user_id: str) -> Dict[str, Any]:
|
|
task = {
|
|
"status": "pending",
|
|
"task_id": task_id,
|
|
"progress": 0,
|
|
"user_id": user_id,
|
|
}
|
|
self._client.set(self._key(task_id), json.dumps(task, ensure_ascii=False))
|
|
self._client.sadd(self._index_key, task_id)
|
|
return task
|
|
|
|
def get(self, task_id: str) -> Dict[str, Any]:
|
|
raw = self._client.get(self._key(task_id))
|
|
if not raw:
|
|
return {"status": "not_found"}
|
|
return json.loads(raw)
|
|
|
|
def list(self) -> List[Dict[str, Any]]:
|
|
task_ids = list(self._client.smembers(self._index_key) or [])
|
|
if not task_ids:
|
|
return []
|
|
keys = [self._key(task_id) for task_id in task_ids]
|
|
raw_items = self._client.mget(keys)
|
|
tasks = []
|
|
for raw in raw_items:
|
|
if raw:
|
|
try:
|
|
tasks.append(json.loads(raw))
|
|
except Exception:
|
|
continue
|
|
return tasks
|
|
|
|
def update(self, task_id: str, updates: Dict[str, Any]) -> Dict[str, Any]:
|
|
task = self.get(task_id)
|
|
if task.get("status") == "not_found":
|
|
task = {"status": "pending", "task_id": task_id}
|
|
task.update(updates)
|
|
self._client.set(self._key(task_id), json.dumps(task, ensure_ascii=False))
|
|
self._client.sadd(self._index_key, task_id)
|
|
return task
|
|
|
|
|
|
def _build_task_store():
|
|
if redis is None:
|
|
logger.warning("Redis not available, using in-memory task store")
|
|
return InMemoryTaskStore()
|
|
try:
|
|
client = redis.Redis.from_url(settings.REDIS_URL, decode_responses=True)
|
|
client.ping()
|
|
logger.info("Using Redis task store")
|
|
return RedisTaskStore(client)
|
|
except Exception as e:
|
|
logger.warning(f"Redis connection failed, using in-memory task store: {e}")
|
|
return InMemoryTaskStore()
|
|
|
|
|
|
task_store = _build_task_store()
|
|
|
|
|
|
def create_task(task_id: str, user_id: str) -> Dict[str, Any]:
|
|
return task_store.create(task_id, user_id)
|
|
|
|
|
|
def get_task(task_id: str) -> Dict[str, Any]:
|
|
return task_store.get(task_id)
|
|
|
|
|
|
def list_tasks() -> List[Dict[str, Any]]:
|
|
return task_store.list()
|