piexian
d25d2eae18
- 修复裸 except 为 except Exception - 修复单行多语句 (E701) - 移除未使用的导入 (F401) - 添加 noqa 注释忽略 main.py 的 E402 (设计决策) - 代码格式化
153 lines
4.0 KiB
Python
153 lines
4.0 KiB
Python
"""
|
|
Batch task manager for admin batch operations (SSE progress).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
import uuid
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
|
|
class BatchTask:
|
|
def __init__(self, total: int):
|
|
self.id = uuid.uuid4().hex
|
|
self.total = int(total)
|
|
self.processed = 0
|
|
self.ok = 0
|
|
self.fail = 0
|
|
self.status = "running"
|
|
self.warning: Optional[str] = None
|
|
self.result: Optional[Dict[str, Any]] = None
|
|
self.error: Optional[str] = None
|
|
self.created_at = time.time()
|
|
self._queues: List[asyncio.Queue] = []
|
|
self._final_event: Optional[Dict[str, Any]] = None
|
|
self.cancelled = False
|
|
|
|
def snapshot(self) -> Dict[str, Any]:
|
|
return {
|
|
"task_id": self.id,
|
|
"status": self.status,
|
|
"total": self.total,
|
|
"processed": self.processed,
|
|
"ok": self.ok,
|
|
"fail": self.fail,
|
|
"warning": self.warning,
|
|
}
|
|
|
|
def attach(self) -> asyncio.Queue:
|
|
q: asyncio.Queue = asyncio.Queue(maxsize=200)
|
|
self._queues.append(q)
|
|
return q
|
|
|
|
def detach(self, q: asyncio.Queue) -> None:
|
|
if q in self._queues:
|
|
self._queues.remove(q)
|
|
|
|
def _publish(self, event: Dict[str, Any]) -> None:
|
|
for q in list(self._queues):
|
|
try:
|
|
q.put_nowait(event)
|
|
except Exception:
|
|
# Drop if queue is full or closed
|
|
pass
|
|
|
|
def record(
|
|
self, ok: bool, *, item: Any = None, detail: Any = None, error: str = ""
|
|
) -> None:
|
|
self.processed += 1
|
|
if ok:
|
|
self.ok += 1
|
|
else:
|
|
self.fail += 1
|
|
event: Dict[str, Any] = {
|
|
"type": "progress",
|
|
"task_id": self.id,
|
|
"total": self.total,
|
|
"processed": self.processed,
|
|
"ok": self.ok,
|
|
"fail": self.fail,
|
|
}
|
|
if item is not None:
|
|
event["item"] = item
|
|
if detail is not None:
|
|
event["detail"] = detail
|
|
if error:
|
|
event["error"] = error
|
|
self._publish(event)
|
|
|
|
def finish(self, result: Dict[str, Any], *, warning: Optional[str] = None) -> None:
|
|
self.status = "done"
|
|
self.result = result
|
|
self.warning = warning
|
|
event = {
|
|
"type": "done",
|
|
"task_id": self.id,
|
|
"total": self.total,
|
|
"processed": self.processed,
|
|
"ok": self.ok,
|
|
"fail": self.fail,
|
|
"warning": self.warning,
|
|
"result": result,
|
|
}
|
|
self._final_event = event
|
|
self._publish(event)
|
|
|
|
def fail_task(self, error: str) -> None:
|
|
self.status = "error"
|
|
self.error = error
|
|
event = {
|
|
"type": "error",
|
|
"task_id": self.id,
|
|
"total": self.total,
|
|
"processed": self.processed,
|
|
"ok": self.ok,
|
|
"fail": self.fail,
|
|
"error": error,
|
|
}
|
|
self._final_event = event
|
|
self._publish(event)
|
|
|
|
def cancel(self) -> None:
|
|
self.cancelled = True
|
|
|
|
def finish_cancelled(self) -> None:
|
|
self.status = "cancelled"
|
|
event = {
|
|
"type": "cancelled",
|
|
"task_id": self.id,
|
|
"total": self.total,
|
|
"processed": self.processed,
|
|
"ok": self.ok,
|
|
"fail": self.fail,
|
|
}
|
|
self._final_event = event
|
|
self._publish(event)
|
|
|
|
def final_event(self) -> Optional[Dict[str, Any]]:
|
|
return self._final_event
|
|
|
|
|
|
_TASKS: Dict[str, BatchTask] = {}
|
|
|
|
|
|
def create_task(total: int) -> BatchTask:
|
|
task = BatchTask(total)
|
|
_TASKS[task.id] = task
|
|
return task
|
|
|
|
|
|
def get_task(task_id: str) -> Optional[BatchTask]:
|
|
return _TASKS.get(task_id)
|
|
|
|
|
|
def delete_task(task_id: str) -> None:
|
|
_TASKS.pop(task_id, None)
|
|
|
|
|
|
async def expire_task(task_id: str, delay: int = 300) -> None:
|
|
await asyncio.sleep(delay)
|
|
delete_task(task_id)
|