69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""WebSocket 日志推送端点
|
||
|
||
提供 WS /ws/logs/{execution_id} 端点，实时推送 ETL 任务执行日志。
|
||
客户端连接后，先发送已有的历史日志行，再实时推送新日志，
|
||
直到执行结束（收到 None 哨兵）或客户端断开。
|
||
|
||
设计要点：
|
||
- 利用 TaskExecutor 已有的 subscribe/unsubscribe 机制
|
||
- 连接时先回放内存缓冲区中的历史日志，避免丢失已产生的行
|
||
- 通过 asyncio.Queue 接收实时日志，None 表示执行结束
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
|
||
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
||
|
||
from ..services.task_executor import task_executor
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
ws_router = APIRouter()
|
||
|
||
|
||
@ws_router.websocket("/ws/logs/{execution_id}")
|
||
async def ws_logs(websocket: WebSocket, execution_id: str) -> None:
|
||
"""实时推送指定 execution_id 的任务执行日志。
|
||
|
||
流程：
|
||
1. 接受 WebSocket 连接
|
||
2. 回放内存缓冲区中已有的日志行
|
||
3. 订阅 TaskExecutor，持续推送新日志
|
||
4. 收到 None（执行结束）或客户端断开时关闭
|
||
"""
|
||
await websocket.accept()
|
||
logger.info("WebSocket 连接已建立: execution_id=%s", execution_id)
|
||
|
||
# 订阅日志流
|
||
queue = task_executor.subscribe(execution_id)
|
||
|
||
try:
|
||
# 回放已有的历史日志行
|
||
for line in task_executor.get_logs(execution_id):
|
||
await websocket.send_text(line)
|
||
|
||
# 如果任务已经不在运行且没有订阅者队列中的数据，
|
||
# 仍然保持连接等待——可能是任务刚结束但 queue 里还有未消费的消息
|
||
while True:
|
||
msg = await queue.get()
|
||
if msg is None:
|
||
# 执行结束哨兵
|
||
break
|
||
await websocket.send_text(msg)
|
||
|
||
except WebSocketDisconnect:
|
||
logger.info("WebSocket 客户端断开: execution_id=%s", execution_id)
|
||
except Exception:
|
||
logger.exception("WebSocket 异常: execution_id=%s", execution_id)
|
||
finally:
|
||
task_executor.unsubscribe(execution_id, queue)
|
||
# 安全关闭连接（客户端可能已断开，忽略错误）
|
||
try:
|
||
await websocket.close()
|
||
except Exception:
|
||
pass
|
||
logger.info("WebSocket 连接已清理: execution_id=%s", execution_id)
|