Codex
321 lines
9.0 KiB
Python
321 lines
9.0 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
任务执行日志服务
|
|
通过 Redis Pub/Sub 实现任务执行过程中的实时日志推送
|
|
"""
|
|
import json
|
|
import asyncio
|
|
from datetime import datetime
|
|
from typing import Optional, AsyncGenerator
|
|
from dataclasses import dataclass, asdict
|
|
from enum import Enum
|
|
|
|
from utils.redis import RedisClient
|
|
|
|
|
|
class LogLevel(str, Enum):
|
|
"""日志级别"""
|
|
DEBUG = "debug"
|
|
INFO = "info"
|
|
WARNING = "warning"
|
|
ERROR = "error"
|
|
SUCCESS = "success"
|
|
|
|
|
|
@dataclass
|
|
class TaskLogEntry:
|
|
"""任务日志条目"""
|
|
timestamp: str
|
|
level: str
|
|
message: str
|
|
step: Optional[str] = None
|
|
progress: Optional[int] = None # 0-100
|
|
data: Optional[dict] = None
|
|
|
|
def to_dict(self) -> dict:
|
|
result = asdict(self)
|
|
# 移除 None 值
|
|
return {k: v for k, v in result.items() if v is not None}
|
|
|
|
|
|
class TaskLogService:
|
|
"""
|
|
任务日志服务
|
|
|
|
用于在任务执行过程中发布实时日志，前端通过 SSE 订阅
|
|
使用 Redis List 存储历史日志 + Pub/Sub 推送新消息
|
|
"""
|
|
|
|
# Redis 频道前缀
|
|
CHANNEL_PREFIX = "scheduler:task_log:"
|
|
# Redis List 前缀（存储历史日志）
|
|
LIST_PREFIX = "scheduler:task_log_history:"
|
|
# 日志过期时间（秒）- 1小时
|
|
LOG_EXPIRE_SECONDS = 3600
|
|
|
|
@classmethod
|
|
def _get_channel(cls, log_id: str) -> str:
|
|
"""获取 Redis Pub/Sub 频道名称"""
|
|
return f"{cls.CHANNEL_PREFIX}{log_id}"
|
|
|
|
@classmethod
|
|
def _get_list_key(cls, log_id: str) -> str:
|
|
"""获取 Redis List 键名"""
|
|
return f"{cls.LIST_PREFIX}{log_id}"
|
|
|
|
@classmethod
|
|
async def publish(
|
|
cls,
|
|
log_id: str,
|
|
message: str,
|
|
level: LogLevel = LogLevel.INFO,
|
|
step: Optional[str] = None,
|
|
progress: Optional[int] = None,
|
|
data: Optional[dict] = None
|
|
):
|
|
"""
|
|
发布日志消息
|
|
|
|
同时存储到 Redis List（历史）和发布到 Pub/Sub（实时）
|
|
|
|
Args:
|
|
log_id: 日志记录 ID
|
|
message: 日志消息
|
|
level: 日志级别
|
|
step: 当前步骤名称
|
|
progress: 进度百分比 (0-100)
|
|
data: 附加数据
|
|
"""
|
|
try:
|
|
redis = await RedisClient.get_client()
|
|
except Exception:
|
|
return
|
|
|
|
entry = TaskLogEntry(
|
|
timestamp=datetime.now().isoformat(),
|
|
level=level.value,
|
|
message=message,
|
|
step=step,
|
|
progress=progress,
|
|
data=data
|
|
)
|
|
|
|
entry_json = json.dumps(entry.to_dict(), ensure_ascii=False)
|
|
list_key = cls._get_list_key(log_id)
|
|
channel = cls._get_channel(log_id)
|
|
|
|
# 存储到 List（历史日志）
|
|
await redis.rpush(list_key, entry_json)
|
|
# 设置过期时间
|
|
await redis.expire(list_key, cls.LOG_EXPIRE_SECONDS)
|
|
# 发布到 Pub/Sub（实时推送）
|
|
await redis.publish(channel, entry_json)
|
|
|
|
@classmethod
|
|
async def publish_start(cls, log_id: str, job_name: str):
|
|
"""发布任务开始事件"""
|
|
await cls.publish(
|
|
log_id=log_id,
|
|
message=f"任务 {job_name} 开始执行",
|
|
level=LogLevel.INFO,
|
|
step="start",
|
|
progress=0
|
|
)
|
|
|
|
@classmethod
|
|
async def publish_complete(cls, log_id: str, job_name: str, result: Optional[str] = None):
|
|
"""发布任务完成事件"""
|
|
await cls.publish(
|
|
log_id=log_id,
|
|
message=f"任务 {job_name} 执行完成",
|
|
level=LogLevel.SUCCESS,
|
|
step="complete",
|
|
progress=100,
|
|
data={"result": result} if result else None
|
|
)
|
|
|
|
@classmethod
|
|
async def publish_error(cls, log_id: str, job_name: str, error: str):
|
|
"""发布任务错误事件"""
|
|
await cls.publish(
|
|
log_id=log_id,
|
|
message=f"任务 {job_name} 执行失败: {error}",
|
|
level=LogLevel.ERROR,
|
|
step="error",
|
|
data={"error": error}
|
|
)
|
|
|
|
@classmethod
|
|
async def get_history(cls, log_id: str) -> list:
|
|
"""
|
|
获取历史日志
|
|
|
|
Args:
|
|
log_id: 日志记录 ID
|
|
|
|
Returns:
|
|
日志条目列表
|
|
"""
|
|
try:
|
|
redis = await RedisClient.get_client()
|
|
except Exception:
|
|
return []
|
|
|
|
list_key = cls._get_list_key(log_id)
|
|
entries = await redis.lrange(list_key, 0, -1)
|
|
|
|
result = []
|
|
for entry_json in entries:
|
|
try:
|
|
result.append(json.loads(entry_json))
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
return result
|
|
|
|
@classmethod
|
|
async def subscribe(cls, log_id: str) -> AsyncGenerator[dict, None]:
|
|
"""
|
|
订阅日志消息
|
|
|
|
先返回历史日志，然后订阅新消息
|
|
|
|
Args:
|
|
log_id: 日志记录 ID
|
|
|
|
Yields:
|
|
日志条目字典
|
|
"""
|
|
try:
|
|
redis = await RedisClient.get_client()
|
|
except Exception:
|
|
return
|
|
|
|
list_key = cls._get_list_key(log_id)
|
|
channel = cls._get_channel(log_id)
|
|
|
|
# 先获取历史日志
|
|
history = await redis.lrange(list_key, 0, -1)
|
|
|
|
# 返回历史日志
|
|
for entry_json in history:
|
|
try:
|
|
data = json.loads(entry_json)
|
|
yield data
|
|
|
|
# 如果历史日志中已经有完成或错误事件，直接结束
|
|
if data.get("step") in ("complete", "error"):
|
|
return
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
# 订阅新消息
|
|
pubsub = redis.pubsub()
|
|
await pubsub.subscribe(channel)
|
|
|
|
try:
|
|
async for message in pubsub.listen():
|
|
if message["type"] == "message":
|
|
data = json.loads(message["data"])
|
|
yield data
|
|
|
|
# 如果是完成或错误事件，结束订阅
|
|
if data.get("step") in ("complete", "error"):
|
|
break
|
|
finally:
|
|
await pubsub.unsubscribe(channel)
|
|
await pubsub.close()
|
|
|
|
|
|
class TaskLogger:
|
|
"""
|
|
任务日志记录器
|
|
|
|
在任务函数中使用，用于记录执行过程中的日志
|
|
"""
|
|
|
|
def __init__(self, log_id: str, job_name: str = ""):
|
|
self.log_id = log_id
|
|
self.job_name = job_name
|
|
self._current_step = None
|
|
self._progress = 0
|
|
|
|
async def debug(self, message: str, **kwargs):
|
|
"""记录调试日志"""
|
|
await TaskLogService.publish(
|
|
log_id=self.log_id,
|
|
message=message,
|
|
level=LogLevel.DEBUG,
|
|
step=self._current_step,
|
|
progress=self._progress,
|
|
**kwargs
|
|
)
|
|
|
|
async def info(self, message: str, **kwargs):
|
|
"""记录信息日志"""
|
|
await TaskLogService.publish(
|
|
log_id=self.log_id,
|
|
message=message,
|
|
level=LogLevel.INFO,
|
|
step=self._current_step,
|
|
progress=self._progress,
|
|
**kwargs
|
|
)
|
|
|
|
async def warning(self, message: str, **kwargs):
|
|
"""记录警告日志"""
|
|
await TaskLogService.publish(
|
|
log_id=self.log_id,
|
|
message=message,
|
|
level=LogLevel.WARNING,
|
|
step=self._current_step,
|
|
progress=self._progress,
|
|
**kwargs
|
|
)
|
|
|
|
async def error(self, message: str, **kwargs):
|
|
"""记录错误日志"""
|
|
await TaskLogService.publish(
|
|
log_id=self.log_id,
|
|
message=message,
|
|
level=LogLevel.ERROR,
|
|
step=self._current_step,
|
|
progress=self._progress,
|
|
**kwargs
|
|
)
|
|
|
|
async def success(self, message: str, **kwargs):
|
|
"""记录成功日志"""
|
|
await TaskLogService.publish(
|
|
log_id=self.log_id,
|
|
message=message,
|
|
level=LogLevel.SUCCESS,
|
|
step=self._current_step,
|
|
progress=self._progress,
|
|
**kwargs
|
|
)
|
|
|
|
def set_step(self, step: str):
|
|
"""设置当前步骤"""
|
|
self._current_step = step
|
|
|
|
def set_progress(self, progress: int):
|
|
"""设置进度 (0-100)"""
|
|
self._progress = max(0, min(100, progress))
|
|
|
|
async def step(self, step: str, message: str, progress: Optional[int] = None):
|
|
"""
|
|
记录步骤日志
|
|
|
|
Args:
|
|
step: 步骤名称
|
|
message: 日志消息
|
|
progress: 进度百分比
|
|
"""
|
|
self._current_step = step
|
|
if progress is not None:
|
|
self._progress = progress
|
|
|
|
await self.info(message, step=step, progress=self._progress)
|