Kris
2348a04163
本次提交包含多项核心更新： 1. 新增微信4.0浅色主题UI模板资源与说明文档，补充了联系人图标、新的朋友入口等四个UI元素的图像匹配模板 2. 重构UI驱动层，将原2375行的单文件拆分为5个职责清晰的子驱动类，提升代码可维护性 3. 新增批量群发消息与聊天记录导出的完整数据模型、路由与后台协程，支持幂等校验与风控配置 4. 为所有业务接口新增post_verify校验逻辑，补充了verified字段返回操作结果 5. 优化媒体文件解密逻辑，修复了导出功能中的数据库错误处理与分辨率硬编码问题 6. 更新版本配置，将媒体发送功能标记为已落地，调整了能力列表与配置项
1016 lines
37 KiB
Python
1016 lines
37 KiB
Python
"""聊天记录导出路由：同步 + 异步两种模式。
|
||
|
||
同步端点（GET /api/messages/export）：
|
||
limit ≤ 1000 的纯文本/HTML 导出，直接 StreamingResponse 返回。
|
||
支持 HTML / CSV / JSON / TXT 四种格式，HTML 可选 media_inline 内联图片。
|
||
|
||
异步端点（POST /api/messages/export + GET /{task_id}/status +
|
||
GET /{task_id}/download + DELETE /{task_id}）：
|
||
limit > 1000 或 include_media=true 时使用，后台 worker 流式生成，
|
||
完成后下载链接有效 2 小时（expires_at）。include_media=true 且
|
||
media_inline=false 时返回 zip（主文件 + media/ 目录）。
|
||
|
||
约束：
|
||
- 全局同时只允许 1 个异步导出任务运行（_MAX_CONCURRENT_EXPORT_TASKS）
|
||
- 媒体抓取失败写 [媒体缺失] 占位，不中断导出
|
||
- 不引入新第三方依赖
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import base64
|
||
import csv
|
||
import html as _html
|
||
import io
|
||
import json
|
||
import logging
|
||
import os
|
||
import shutil
|
||
import time
|
||
import uuid
|
||
import zipfile
|
||
from datetime import datetime
|
||
from typing import AsyncIterator, Optional
|
||
|
||
from fastapi import APIRouter
|
||
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
|
||
|
||
from woc_bridge.config import _require_db_reader, _state
|
||
from woc_bridge.db.coordinator import _check_db_readable, with_db_retry
|
||
from woc_bridge.models import BridgeError, ExportRequest, ExportTaskStatus
|
||
|
||
logger = logging.getLogger("woc-bridge")
|
||
router = APIRouter(prefix="/api/messages/export", tags=["messages-export"])
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 常量
|
||
# ---------------------------------------------------------------------------
|
||
# 同步导出阈值：limit 超过此值必须走异步
|
||
_SYNC_EXPORT_LIMIT = 1000
|
||
|
||
# 异步任务并发限制：同时只允许 1 个任务运行
|
||
_MAX_CONCURRENT_EXPORT_TASKS = 1
|
||
|
||
# 异步任务分页大小（每批从 DB 拉取 200 条）
|
||
_EXPORT_BATCH_SIZE = 200
|
||
|
||
# 任务存储（内存，进程级）。task_id → 任务状态 dict
|
||
# dict 结构与 ExportTaskStatus 对齐，外加内部字段：
|
||
# _task: asyncio.Task | None 后台 worker 句柄
|
||
# _file_path: str | None 生成的文件绝对路径
|
||
# _task_dir: str | None 任务目录（zip 模式含 media/）
|
||
# _cancelled: bool 是否被主动取消
|
||
_export_tasks: dict[str, dict] = {}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 辅助函数
|
||
# ---------------------------------------------------------------------------
|
||
def _now_ts() -> int:
|
||
"""当前 Unix 时间戳（秒）。"""
|
||
return int(time.time())
|
||
|
||
|
||
def _format_time(ts: int) -> str:
|
||
"""时间戳格式化为 YYYY-MM-DD HH:MM:SS（本地时区）。"""
|
||
if not ts:
|
||
return "未知时间"
|
||
try:
|
||
return datetime.fromtimestamp(ts).strftime("%Y-%m-%d %H:%M:%S")
|
||
except (OSError, ValueError, OverflowError):
|
||
return "未知时间"
|
||
|
||
|
||
def _format_date(ts: int) -> str:
|
||
"""时间戳格式化为 YYYY-MM-DD（用于时间分隔条）。"""
|
||
if not ts:
|
||
return ""
|
||
try:
|
||
return datetime.fromtimestamp(ts).strftime("%Y-%m-%d")
|
||
except (OSError, ValueError, OverflowError):
|
||
return ""
|
||
|
||
|
||
def _ext_for_format(fmt: str) -> str:
|
||
"""格式 → 文件扩展名。"""
|
||
return {"html": "html", "csv": "csv", "json": "json", "txt": "txt"}.get(
|
||
fmt, "txt"
|
||
)
|
||
|
||
|
||
def _mime_for_format(fmt: str) -> str:
|
||
"""格式 → MIME 类型。"""
|
||
return {
|
||
"html": "text/html; charset=utf-8",
|
||
"csv": "text/csv; charset=utf-8",
|
||
"json": "application/json; charset=utf-8",
|
||
"txt": "text/plain; charset=utf-8",
|
||
"zip": "application/zip",
|
||
}.get(fmt, "text/plain; charset=utf-8")
|
||
|
||
|
||
def _safe_filename(talker: str) -> str:
|
||
"""把 wxid/chatroom id 转为文件名安全字符串。"""
|
||
safe = "".join(c if c.isalnum() or c in ("-", "_", ".") else "_" for c in talker)
|
||
return safe or "export"
|
||
|
||
|
||
def _count_active_tasks() -> int:
|
||
"""统计当前 pending/running 状态的任务数（用于并发限制）。"""
|
||
return sum(
|
||
1 for t in _export_tasks.values()
|
||
if t.get("status") in ("pending", "running")
|
||
)
|
||
|
||
|
||
def _lazy_cleanup_expired() -> None:
|
||
"""懒清理过期任务（GET /status 时触发）。
|
||
|
||
P1 修复：过期判定扩展到 completed/failed/cancelled 三种终态。
|
||
running/pending 不清理（仍在执行中）。
|
||
清理动作：删除任务目录与文件，从 _export_tasks 移除。
|
||
"""
|
||
now = _now_ts()
|
||
expired_ids: list[str] = []
|
||
for tid, task in list(_export_tasks.items()):
|
||
status = task.get("status")
|
||
# P1 修复：failed/cancelled 也需清理，否则内存泄漏
|
||
if status not in ("completed", "failed", "cancelled"):
|
||
continue
|
||
exp = task.get("expires_at")
|
||
if exp is not None and exp < now:
|
||
expired_ids.append(tid)
|
||
for tid in expired_ids:
|
||
task = _export_tasks.pop(tid, None)
|
||
if task is None:
|
||
continue
|
||
# 删除任务目录（含主文件与 media/）
|
||
task_dir = task.get("_task_dir")
|
||
if task_dir and os.path.isdir(task_dir):
|
||
try:
|
||
shutil.rmtree(task_dir, ignore_errors=True)
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
def _safe_remove(path: str) -> None:
|
||
"""安全删除文件，吞掉 OSError。"""
|
||
try:
|
||
if path and os.path.isfile(path):
|
||
os.remove(path)
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 4 个格式生成器（同步版本，用于 StreamingResponse）
|
||
# ---------------------------------------------------------------------------
|
||
def _html_head(talker: str, nickname: str) -> bytes:
|
||
"""生成 HTML 头部 + CSS（仿微信气泡样式）。"""
|
||
display_name = nickname or talker
|
||
return (
|
||
"<!DOCTYPE html>\n"
|
||
'<html lang="zh-CN">\n<head>\n'
|
||
'<meta charset="UTF-8">\n'
|
||
f'<title>聊天记录 - {_html.escape(display_name)}</title>\n'
|
||
"<style>\n"
|
||
" body { background:#EDEDED; margin:0; padding:20px; "
|
||
"font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',sans-serif; "
|
||
"font-size:14px; color:#333; }\n"
|
||
" .header { text-align:center; padding:10px 0 20px; "
|
||
"font-size:18px; color:#111; font-weight:600; }\n"
|
||
" .msg { margin:6px 0; display:flex; }\n"
|
||
" .msg.self { justify-content:flex-end; }\n"
|
||
" .msg.other { justify-content:flex-start; }\n"
|
||
" .bubble { max-width:60%; padding:8px 12px; border-radius:6px; "
|
||
"word-break:break-word; box-shadow:0 1px 1px rgba(0,0,0,0.05); }\n"
|
||
" .self .bubble { background:#95EC69; }\n"
|
||
" .other .bubble { background:#FFFFFF; }\n"
|
||
" .system { text-align:center; margin:10px 0; }\n"
|
||
" .system span { background:#DADADA; color:#FFF; padding:2px 8px; "
|
||
"border-radius:3px; font-size:12px; }\n"
|
||
" .time-sep { text-align:center; margin:14px 0; }\n"
|
||
" .time-sep span { background:#CECECE; color:#FFF; padding:2px 8px; "
|
||
"border-radius:3px; font-size:12px; }\n"
|
||
" .bubble img { max-width:100%; height:auto; display:block; "
|
||
"border-radius:4px; }\n"
|
||
" .media-missing { color:#999; font-style:italic; }\n"
|
||
" .sender { font-size:12px; color:#888; margin-bottom:2px; }\n"
|
||
"</style>\n</head>\n<body>\n"
|
||
f'<div class="header">{_html.escape(display_name)}</div>\n'
|
||
).encode("utf-8")
|
||
|
||
|
||
def _html_tail() -> bytes:
|
||
"""HTML 尾部。"""
|
||
return b"</body>\n</html>\n"
|
||
|
||
|
||
def _render_media_html(msg: dict, media_inline: bool, db_reader=None) -> str:
|
||
"""渲染单条媒体消息的 HTML 内部内容。
|
||
|
||
media_inline=true 时图片转 base64 内联（经 db_reader 解密 .dat）；
|
||
否则用 media/<filename> 相对路径（由 _copy_media_files 写入解密后的文件）。
|
||
媒体不可读时返回 [媒体缺失] 占位。
|
||
|
||
P0 修复：media_inline=true 时通过 db_reader.read_media_file_by_path 读取
|
||
并解密 .dat 文件，不再直接读取原始密文。
|
||
"""
|
||
media_path = msg.get("media_path")
|
||
render_type = msg.get("render_type", "")
|
||
if not media_path or not os.path.isfile(media_path) or not os.access(media_path, os.R_OK):
|
||
label = render_type or "媒体"
|
||
return f'<span class="media-missing">[{_html.escape(label)}缺失]</span>'
|
||
|
||
filename = os.path.basename(media_path)
|
||
|
||
if render_type == "image":
|
||
if media_inline:
|
||
# P0 修复：通过 db_reader 读取并解密 .dat，不再直接读取原始密文
|
||
media_data = None
|
||
if db_reader is not None:
|
||
try:
|
||
result = db_reader.read_media_file_by_path(media_path)
|
||
if result is not None:
|
||
media_data, mime_from_db = result
|
||
except Exception:
|
||
pass
|
||
if media_data is None:
|
||
return '<span class="media-missing">[图片缺失]</span>'
|
||
b64 = base64.b64encode(media_data).decode("ascii")
|
||
ext = os.path.splitext(filename)[1].lower().lstrip(".")
|
||
mime = {
|
||
"jpg": "image/jpeg", "jpeg": "image/jpeg",
|
||
"png": "image/png", "gif": "image/gif",
|
||
"bmp": "image/bmp", "webp": "image/webp",
|
||
}.get(ext, "image/jpeg")
|
||
return f'<img src="data:{mime};base64,{b64}" alt="{_html.escape(filename)}">'
|
||
return f'<img src="media/{_html.escape(filename)}" alt="{_html.escape(filename)}">'
|
||
|
||
# 语音/视频/文件：给出下载链接
|
||
label = {
|
||
"voice": "语音", "video": "视频", "file": "文件",
|
||
}.get(render_type, "媒体")
|
||
return (
|
||
f'<a href="media/{_html.escape(filename)}" '
|
||
f'download="{_html.escape(filename)}">[{label}] {_html.escape(filename)}</a>'
|
||
)
|
||
|
||
|
||
def _render_html_messages(
|
||
messages: list[dict],
|
||
talker: str,
|
||
media_inline: bool,
|
||
last_ts: Optional[int],
|
||
db_reader=None,
|
||
) -> tuple[bytes, Optional[int]]:
|
||
"""渲染一批消息为 HTML 片段，返回 (bytes, new_last_ts)。
|
||
|
||
时间分隔条：相邻消息间隔 > 5 分钟插入。
|
||
"""
|
||
buf = io.StringIO()
|
||
for msg in messages:
|
||
ts = int(msg.get("create_time", 0) or 0)
|
||
if last_ts is not None and ts - last_ts > 300:
|
||
if _format_date(ts) != _format_date(last_ts):
|
||
sep_text = _format_date(ts)
|
||
else:
|
||
sep_text = _format_time(ts)
|
||
buf.write(
|
||
f'<div class="time-sep"><span>{_html.escape(sep_text)}</span></div>\n'
|
||
)
|
||
last_ts = ts
|
||
|
||
render_type = msg.get("render_type", "text")
|
||
is_self = bool(msg.get("is_sender"))
|
||
sender = msg.get("sender", "") or ""
|
||
content = msg.get("content", "") or ""
|
||
|
||
if render_type == "system":
|
||
sys_text = _html.escape(content) or "[系统消息]"
|
||
buf.write(f'<div class="system"><span>{sys_text}</span></div>\n')
|
||
continue
|
||
|
||
side = "self" if is_self else "other"
|
||
if render_type in ("image", "voice", "video", "file"):
|
||
bubble_inner = _render_media_html(msg, media_inline, db_reader)
|
||
else:
|
||
text = _html.escape(content)
|
||
bubble_inner = text.replace("\n", "<br>") if text else "[空消息]"
|
||
|
||
sender_html = ""
|
||
if msg.get("session_type") == "group" and not is_self and sender:
|
||
sender_html = f'<div class="sender">{_html.escape(sender)}</div>'
|
||
|
||
buf.write(
|
||
f'<div class="msg {side}">'
|
||
f'<div>{sender_html}<div class="bubble">{bubble_inner}</div></div>'
|
||
f"</div>\n"
|
||
)
|
||
return buf.getvalue().encode("utf-8"), last_ts
|
||
|
||
|
||
def _generate_html(
|
||
messages: list[dict],
|
||
talker: str,
|
||
nickname: str,
|
||
media_inline: bool = False,
|
||
db_reader=None,
|
||
) -> AsyncIterator[bytes]:
|
||
"""HTML 生成器（同步端点用）：内联 CSS 气泡样式。
|
||
|
||
本人消息右对齐 #95EC69，对方左对齐 #FFFFFF，系统居中灰，
|
||
相邻消息间隔 > 5 分钟插入时间分隔条，
|
||
媒体消息渲染 <img> 或 [媒体缺失] 占位。
|
||
"""
|
||
yield _html_head(talker, nickname)
|
||
last_ts: Optional[int] = None
|
||
chunk, last_ts = _render_html_messages(
|
||
messages, talker, media_inline, last_ts, db_reader
|
||
)
|
||
yield chunk
|
||
yield _html_tail()
|
||
|
||
|
||
def _generate_csv(messages: list[dict]) -> AsyncIterator[bytes]:
|
||
"""CSV 生成器：msg_id, create_time, is_self, type, content 列。
|
||
|
||
content 中的换行符由 csv 模块自动处理（双引号包裹）。
|
||
首行写 UTF-8 BOM 让 Excel 正确识别编码。
|
||
"""
|
||
buf = io.StringIO()
|
||
writer = csv.writer(buf)
|
||
writer.writerow(["msg_id", "create_time", "is_self", "type", "content"])
|
||
yield buf.getvalue().encode("utf-8-sig")
|
||
|
||
buf.seek(0)
|
||
buf.truncate(0)
|
||
for msg in messages:
|
||
writer.writerow([
|
||
msg.get("msg_id", ""),
|
||
msg.get("create_time", 0),
|
||
"true" if msg.get("is_sender") else "false",
|
||
msg.get("type", 0),
|
||
msg.get("content", "") or "",
|
||
])
|
||
if buf.tell() > 8192:
|
||
yield buf.getvalue().encode("utf-8")
|
||
buf.seek(0)
|
||
buf.truncate(0)
|
||
if buf.tell() > 0:
|
||
yield buf.getvalue().encode("utf-8")
|
||
|
||
|
||
def _make_json_item(msg: dict) -> dict:
|
||
"""构造单条消息的 JSON dict。"""
|
||
return {
|
||
"msg_id": msg.get("msg_id", ""),
|
||
"create_time": msg.get("create_time", 0),
|
||
"is_sender": bool(msg.get("is_sender")),
|
||
"type": msg.get("type", 0),
|
||
"render_type": msg.get("render_type", "text"),
|
||
"sender": msg.get("sender", ""),
|
||
"content": msg.get("content", "") or "",
|
||
"session_type": msg.get("session_type", ""),
|
||
"media_path": msg.get("media_path"),
|
||
}
|
||
|
||
|
||
def _generate_json(messages: list[dict], talker: str) -> AsyncIterator[bytes]:
|
||
"""JSON 生成器：结构化字段，含 media_path。
|
||
|
||
输出格式：{"talker": "...", "count": N, "messages": [...]}
|
||
流式输出：head → 每条 item（用逗号分隔）→ tail。
|
||
"""
|
||
yield (
|
||
f'{{"talker": {json.dumps(talker, ensure_ascii=False)}, '
|
||
f'"count": {len(messages)}, "messages": [\n'
|
||
).encode("utf-8")
|
||
for i, msg in enumerate(messages):
|
||
item = _make_json_item(msg)
|
||
sep = "," if i < len(messages) - 1 else ""
|
||
yield f" {json.dumps(item, ensure_ascii=False)}{sep}\n".encode("utf-8")
|
||
yield b"]}\n"
|
||
|
||
|
||
def _generate_txt(messages: list[dict], nickname: str) -> AsyncIterator[bytes]:
|
||
"""TXT 生成器：[时间] 昵称: 内容 纯文本。
|
||
|
||
本人消息用 "我" 标记，对方用 nickname 或 talker 显示。
|
||
群消息用 sender 字段（如果存在）。
|
||
"""
|
||
display_name = nickname or "对方"
|
||
for msg in messages:
|
||
ts = int(msg.get("create_time", 0) or 0)
|
||
time_str = _format_time(ts)
|
||
is_self = bool(msg.get("is_sender"))
|
||
sender = msg.get("sender", "") or ""
|
||
content = msg.get("content", "") or ""
|
||
render_type = msg.get("render_type", "text")
|
||
|
||
if render_type in ("image", "voice", "video", "file"):
|
||
label = {
|
||
"image": "图片", "voice": "语音",
|
||
"video": "视频", "file": "文件",
|
||
}.get(render_type, "媒体")
|
||
if not msg.get("media_path"):
|
||
content_text = f"[{label}缺失]"
|
||
else:
|
||
content_text = f"[{label}]"
|
||
else:
|
||
content_text = content
|
||
|
||
if is_self:
|
||
who = "我"
|
||
elif sender and msg.get("session_type") == "group":
|
||
who = sender
|
||
else:
|
||
who = display_name
|
||
|
||
yield f"[{time_str}] {who}: {content_text}\n".encode("utf-8")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 同步导出端点：GET /api/messages/export
|
||
# ---------------------------------------------------------------------------
|
||
@router.get("")
|
||
@with_db_retry
|
||
async def sync_export(
|
||
talker: str,
|
||
format: str = "html",
|
||
limit: int = 1000,
|
||
include_media: bool = False,
|
||
media_inline: bool = False,
|
||
start_time: Optional[int] = None,
|
||
end_time: Optional[int] = None,
|
||
):
|
||
"""同步导出（limit ≤ 1000）。
|
||
|
||
直接 StreamingResponse 返回生成内容，Content-Disposition 触发浏览器下载。
|
||
HTML 格式可设 media_inline=true 将图片以 base64 内联（仅适合少量图片）。
|
||
|
||
Args:
|
||
talker: 会话 wxid 或群 id
|
||
format: html / csv / json / txt
|
||
limit: 1~1000
|
||
include_media: 同步端点忽略此参数（媒体通过 media_inline 内联或路径引用）
|
||
media_inline: 仅 HTML，图片转 base64 内联
|
||
start_time: 起始时间戳（含）
|
||
end_time: 结束时间戳（含）
|
||
|
||
Raises:
|
||
BridgeError(INVALID_PARAMS): talker 为空、format 非法、limit 越界
|
||
BridgeError(RATE_LIMITED): 已有异步导出任务运行中
|
||
BridgeError(DB_NOT_FOUND / DB_ENCRYPTED): DB 不可读
|
||
"""
|
||
if not talker:
|
||
raise BridgeError(code="INVALID_PARAMS", message="talker 不能为空")
|
||
if format not in ("html", "csv", "json", "txt"):
|
||
raise BridgeError(
|
||
code="INVALID_PARAMS",
|
||
message=f"format 必须为 html/csv/json/txt，收到 {format}",
|
||
)
|
||
if limit < 1 or limit > _SYNC_EXPORT_LIMIT:
|
||
raise BridgeError(
|
||
code="INVALID_PARAMS",
|
||
message=f"同步导出 limit 必须在 1~{_SYNC_EXPORT_LIMIT} 之间，收到 {limit}，"
|
||
"如需导出更多请用 POST /api/messages/export 异步端点",
|
||
)
|
||
if start_time is not None and end_time is not None and start_time > end_time:
|
||
raise BridgeError(
|
||
code="INVALID_PARAMS",
|
||
message="start_time 不能大于 end_time",
|
||
)
|
||
|
||
# 并发限制：异步任务运行中时拒绝同步导出，避免 DB I/O 抢占
|
||
if _count_active_tasks() >= _MAX_CONCURRENT_EXPORT_TASKS:
|
||
raise BridgeError(
|
||
code="RATE_LIMITED",
|
||
message="已有异步导出任务运行中，请稍后重试或先取消该任务",
|
||
details={"retry_after": 30},
|
||
)
|
||
|
||
db_reader = _require_db_reader()
|
||
await _check_db_readable()
|
||
|
||
# 一次性查 DB（limit ≤ 1000，内存可控）
|
||
messages = await asyncio.to_thread(
|
||
db_reader.get_messages_for_export,
|
||
talker,
|
||
start_time,
|
||
end_time,
|
||
limit,
|
||
0,
|
||
)
|
||
if messages is None:
|
||
raise BridgeError(
|
||
code="DB_NOT_FOUND",
|
||
message="DB 不可读，无法导出消息",
|
||
)
|
||
|
||
# 取昵称（HTML/TXT 用）
|
||
nickname: Optional[str] = ""
|
||
if format in ("html", "txt"):
|
||
try:
|
||
nickname = await asyncio.to_thread(db_reader.get_contact_nickname, talker)
|
||
except Exception:
|
||
nickname = None
|
||
|
||
filename = f"export_{_safe_filename(talker)}.{_ext_for_format(format)}"
|
||
media_type = _mime_for_format(format)
|
||
headers = {
|
||
"Content-Disposition": f'attachment; filename="{filename}"',
|
||
}
|
||
|
||
if format == "html":
|
||
gen = _generate_html(messages, talker, nickname or "", media_inline, db_reader)
|
||
elif format == "csv":
|
||
gen = _generate_csv(messages)
|
||
elif format == "json":
|
||
gen = _generate_json(messages, talker)
|
||
else: # txt
|
||
gen = _generate_txt(messages, nickname or "")
|
||
|
||
logger.info(
|
||
"export sync: talker=%s format=%s limit=%d → %d 条",
|
||
talker, format, limit, len(messages),
|
||
)
|
||
return StreamingResponse(gen, media_type=media_type, headers=headers)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 异步导出端点：POST /api/messages/export
|
||
# ---------------------------------------------------------------------------
|
||
@router.post("")
|
||
async def submit_export_task(req: ExportRequest):
|
||
"""提交异步导出任务。
|
||
|
||
适用场景：limit > 1000 或 include_media=true。
|
||
立即返回 task_id + status=pending，后台 worker 流式生成文件。
|
||
|
||
全局同时只允许 1 个任务运行（_MAX_CONCURRENT_EXPORT_TASKS）。
|
||
提交时即把状态置为 running（避免并发提交绕过限制）。
|
||
|
||
Returns:
|
||
ExportTaskStatus：task_id + status=pending + progress=0
|
||
|
||
Raises:
|
||
BridgeError(INVALID_PARAMS): start_time > end_time
|
||
BridgeError(RATE_LIMITED): 已有任务运行中
|
||
"""
|
||
if req.start_time is not None and req.end_time is not None and req.start_time > req.end_time:
|
||
raise BridgeError(
|
||
code="INVALID_PARAMS",
|
||
message="start_time 不能大于 end_time",
|
||
)
|
||
|
||
# 懒清理过期任务（顺带回收槽位）
|
||
_lazy_cleanup_expired()
|
||
|
||
# 并发限制：check-then-set 在 asyncio 单线程内原子（无 await 间隙）
|
||
if _count_active_tasks() >= _MAX_CONCURRENT_EXPORT_TASKS:
|
||
raise BridgeError(
|
||
code="RATE_LIMITED",
|
||
message="已有导出任务运行中，请稍后重试或先取消该任务（DELETE /api/messages/export/{task_id}）",
|
||
details={"retry_after": 30},
|
||
)
|
||
|
||
task_id = uuid.uuid4().hex
|
||
_export_tasks[task_id] = {
|
||
"task_id": task_id,
|
||
# 直接置 running，避免下一并发提交绕过限制
|
||
"status": "running",
|
||
"progress": 0.0,
|
||
"total": 0,
|
||
"processed": 0,
|
||
"download_url": None,
|
||
"expires_at": None,
|
||
"error_message": None,
|
||
"_task": None,
|
||
"_file_path": None,
|
||
"_task_dir": None,
|
||
"_cancelled": False,
|
||
}
|
||
|
||
asyncio_task = asyncio.create_task(_run_export_task(task_id, req))
|
||
_export_tasks[task_id]["_task"] = asyncio_task
|
||
|
||
logger.info(
|
||
"export async submit: task_id=%s talker=%s format=%s limit=%d include_media=%s",
|
||
task_id, req.talker, req.format, req.limit, req.include_media,
|
||
)
|
||
return ExportTaskStatus(
|
||
task_id=task_id,
|
||
status="running",
|
||
progress=0.0,
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 异步导出：GET /api/messages/export/{task_id}/status
|
||
# ---------------------------------------------------------------------------
|
||
@router.get("/{task_id}/status")
|
||
async def get_export_status(task_id: str):
|
||
"""查询导出任务状态。
|
||
|
||
顺带触发懒清理过期任务。
|
||
"""
|
||
_lazy_cleanup_expired()
|
||
task = _export_tasks.get(task_id)
|
||
if task is None:
|
||
raise BridgeError(
|
||
code="INVALID_PARAMS",
|
||
message=f"任务不存在或已过期: {task_id}",
|
||
)
|
||
return ExportTaskStatus(
|
||
task_id=task["task_id"],
|
||
status=task["status"],
|
||
progress=task["progress"],
|
||
total=task["total"],
|
||
processed=task["processed"],
|
||
download_url=task.get("download_url"),
|
||
expires_at=task.get("expires_at"),
|
||
error_message=task.get("error_message"),
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 异步导出：GET /api/messages/export/{task_id}/download
|
||
# ---------------------------------------------------------------------------
|
||
@router.get("/{task_id}/download")
|
||
async def download_export(task_id: str):
|
||
"""下载导出文件。
|
||
|
||
- include_media=true 且 media_inline=false 时返回 zip（主文件 + media/）
|
||
- 其他情况返回单文件（HTML/CSV/JSON/TXT）
|
||
|
||
Raises:
|
||
BridgeError(INVALID_PARAMS): 任务不存在
|
||
BridgeError(STATE_DIRTY): 任务未完成或已过期
|
||
"""
|
||
_lazy_cleanup_expired()
|
||
task = _export_tasks.get(task_id)
|
||
if task is None:
|
||
raise BridgeError(
|
||
code="INVALID_PARAMS",
|
||
message=f"任务不存在或已过期: {task_id}",
|
||
)
|
||
if task["status"] != "completed":
|
||
raise BridgeError(
|
||
code="STATE_DIRTY",
|
||
message=f"任务未完成，当前状态: {task['status']}",
|
||
)
|
||
file_path = task.get("_file_path")
|
||
if not file_path or not os.path.isfile(file_path):
|
||
raise BridgeError(
|
||
code="BRIDGE_INTERNAL_ERROR",
|
||
message="导出文件丢失，请重新提交任务",
|
||
)
|
||
filename = os.path.basename(file_path)
|
||
# 按扩展名决定 MIME
|
||
if filename.endswith(".zip"):
|
||
media_type = _mime_for_format("zip")
|
||
elif filename.endswith(".html"):
|
||
media_type = _mime_for_format("html")
|
||
elif filename.endswith(".csv"):
|
||
media_type = _mime_for_format("csv")
|
||
elif filename.endswith(".json"):
|
||
media_type = _mime_for_format("json")
|
||
else:
|
||
media_type = _mime_for_format("txt")
|
||
return FileResponse(file_path, media_type=media_type, filename=filename)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 异步导出：DELETE /api/messages/export/{task_id}
|
||
# ---------------------------------------------------------------------------
|
||
@router.delete("/{task_id}")
|
||
async def cancel_export(task_id: str):
|
||
"""主动取消任务 + 清理文件。
|
||
|
||
- running 中：取消 asyncio.Task，置 status=cancelled，删除文件
|
||
- pending/failed：直接清理
|
||
- completed：删除文件并从内存移除（提前回收）
|
||
"""
|
||
task = _export_tasks.get(task_id)
|
||
if task is None:
|
||
raise BridgeError(
|
||
code="INVALID_PARAMS",
|
||
message=f"任务不存在: {task_id}",
|
||
)
|
||
|
||
asyncio_task = task.get("_task")
|
||
if asyncio_task is not None and not asyncio_task.done():
|
||
task["_cancelled"] = True
|
||
asyncio_task.cancel()
|
||
try:
|
||
await asyncio.wait_for(asyncio_task, timeout=5.0)
|
||
except (asyncio.TimeoutError, asyncio.CancelledError, Exception):
|
||
pass
|
||
|
||
# 清理任务目录（含主文件与 media/）
|
||
task_dir = task.get("_task_dir")
|
||
if task_dir and os.path.isdir(task_dir):
|
||
try:
|
||
shutil.rmtree(task_dir, ignore_errors=True)
|
||
except OSError:
|
||
pass
|
||
|
||
# 内存移除
|
||
_export_tasks.pop(task_id, None)
|
||
|
||
logger.info("export cancel: task_id=%s", task_id)
|
||
return JSONResponse({"success": True, "task_id": task_id, "status": "cancelled"})
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 后台 worker
|
||
# ---------------------------------------------------------------------------
|
||
async def _run_export_task(task_id: str, req: ExportRequest) -> None:
|
||
"""后台 worker：分页查 DB（每批 _EXPORT_BATCH_SIZE 条）流式写入。
|
||
|
||
流程：
|
||
1. 创建 <export_dir>/<task_id>/ 目录
|
||
2. 分页查 DB，每批写入文件
|
||
3. include_media=true 且 media_inline=false 时拷贝媒体文件到 media/
|
||
4. 完成后：
|
||
- media_inline=false 且 include_media=true：打包 zip
|
||
- 否则：保留单文件
|
||
5. 设置 download_url + expires_at（默认 2 小时后），status=completed
|
||
6. 异常：status=failed + error_message
|
||
7. 取消（_cancelled=true）：status=cancelled，删文件
|
||
"""
|
||
task = _export_tasks.get(task_id)
|
||
if task is None:
|
||
return
|
||
|
||
# 导出目录（来自 BridgeConfig.export_dir）
|
||
export_dir = (
|
||
_state.config.export_dir if _state.config else "/config/woc-export"
|
||
)
|
||
task_dir = os.path.join(export_dir, task_id)
|
||
media_dir = os.path.join(task_dir, "media")
|
||
try:
|
||
os.makedirs(task_dir, exist_ok=True)
|
||
if req.include_media and not req.media_inline:
|
||
os.makedirs(media_dir, exist_ok=True)
|
||
except OSError as e:
|
||
task["status"] = "failed"
|
||
task["error_message"] = f"创建任务目录失败: {e}"
|
||
# P0 修复：failed 任务也设置 expires_at，否则 _lazy_cleanup_expired 永不清理
|
||
task["expires_at"] = _now_ts() + 300
|
||
return
|
||
task["_task_dir"] = task_dir
|
||
|
||
# 主文件路径
|
||
ext = _ext_for_format(req.format)
|
||
main_filename = f"export_{_safe_filename(req.talker)}.{ext}"
|
||
main_path = os.path.join(task_dir, main_filename)
|
||
|
||
try:
|
||
# P1 修复：_require_db_reader 移入 try 块，DB 不可用时任务标 failed
|
||
# 而非永久卡在 running
|
||
db_reader = _require_db_reader()
|
||
except BridgeError as e:
|
||
task["status"] = "failed"
|
||
task["error_message"] = f"DB 不可读: {e.message}"
|
||
# P0 修复：设置 expires_at + 清理 task_dir
|
||
task["expires_at"] = _now_ts() + 300
|
||
shutil.rmtree(task_dir, ignore_errors=True)
|
||
return
|
||
|
||
nickname: Optional[str] = ""
|
||
if req.format in ("html", "txt"):
|
||
try:
|
||
nickname = await asyncio.to_thread(db_reader.get_contact_nickname, req.talker)
|
||
except Exception:
|
||
nickname = None
|
||
|
||
try:
|
||
# 流式写入主文件
|
||
with open(main_path, "wb") as f:
|
||
await _stream_write_main_file(
|
||
f, task, req, db_reader, nickname or "", media_dir
|
||
)
|
||
|
||
if task.get("_cancelled"):
|
||
task["status"] = "cancelled"
|
||
# P0 修复：设置 expires_at + 清理 task_dir
|
||
task["expires_at"] = _now_ts() + 300
|
||
shutil.rmtree(task_dir, ignore_errors=True)
|
||
return
|
||
|
||
# include_media=true 且 media_inline=false：打包 zip
|
||
final_path = main_path
|
||
if (
|
||
req.include_media
|
||
and not req.media_inline
|
||
and os.path.isdir(media_dir)
|
||
and os.listdir(media_dir)
|
||
):
|
||
zip_path = os.path.join(task_dir, f"export_{_safe_filename(req.talker)}.zip")
|
||
with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
|
||
zf.write(main_path, main_filename)
|
||
for root, _, files in os.walk(media_dir):
|
||
for name in files:
|
||
full = os.path.join(root, name)
|
||
arc = os.path.relpath(full, task_dir)
|
||
zf.write(full, arc)
|
||
# zip 已生成，删除主文件原拷贝（zip 内已归档）
|
||
_safe_remove(main_path)
|
||
shutil.rmtree(media_dir, ignore_errors=True)
|
||
final_path = zip_path
|
||
|
||
# 完成
|
||
task["status"] = "completed"
|
||
task["progress"] = 1.0
|
||
task["_file_path"] = final_path
|
||
ttl = (
|
||
_state.config.export_task_ttl_sec if _state.config else 7200
|
||
)
|
||
task["expires_at"] = _now_ts() + ttl
|
||
task["download_url"] = f"/api/messages/export/{task_id}/download"
|
||
logger.info(
|
||
"export async completed: task_id=%s file=%s total=%d",
|
||
task_id, os.path.basename(final_path), task["total"],
|
||
)
|
||
except asyncio.CancelledError:
|
||
task["status"] = "cancelled"
|
||
# P0 修复：设置 expires_at + 清理 task_dir
|
||
task["expires_at"] = _now_ts() + 300
|
||
shutil.rmtree(task_dir, ignore_errors=True)
|
||
raise
|
||
except Exception as e:
|
||
task["status"] = "failed"
|
||
task["error_message"] = str(e)[:500]
|
||
# P0 修复：设置 expires_at + 清理 task_dir
|
||
task["expires_at"] = _now_ts() + 300
|
||
shutil.rmtree(task_dir, ignore_errors=True)
|
||
logger.exception("export async failed: task_id=%s", task_id)
|
||
|
||
|
||
async def _stream_write_main_file(
|
||
f,
|
||
task: dict,
|
||
req: ExportRequest,
|
||
db_reader,
|
||
nickname: str,
|
||
media_dir: str,
|
||
) -> None:
|
||
"""分页查 DB + 流式写入主文件 + 可选拷贝媒体。
|
||
|
||
每批 _EXPORT_BATCH_SIZE 条，调用 db_reader.get_messages_for_export。
|
||
JSON 用 "pending item" 模式处理跨批逗号分隔；HTML 用 head/tail 包裹。
|
||
"""
|
||
fmt = req.format
|
||
limit = req.limit
|
||
batch_size = _EXPORT_BATCH_SIZE
|
||
|
||
fetched = 0
|
||
first_batch = True
|
||
last_ts: Optional[int] = None # HTML 跨批时间分隔条状态
|
||
json_pending: Optional[str] = None # JSON 跨批逗号分隔的 pending item
|
||
|
||
# 格式头部
|
||
if fmt == "json":
|
||
# P1 修复：count 移到尾部（流式无法预知实际条数，写 limit 会误导）
|
||
f.write(
|
||
f'{{"talker": {json.dumps(req.talker, ensure_ascii=False)}, '
|
||
f'"messages": [\n'.encode("utf-8")
|
||
)
|
||
elif fmt == "html":
|
||
f.write(_html_head(req.talker, nickname))
|
||
elif fmt == "csv":
|
||
buf = io.StringIO()
|
||
w = csv.writer(buf)
|
||
w.writerow(["msg_id", "create_time", "is_self", "type", "content"])
|
||
f.write(buf.getvalue().encode("utf-8-sig"))
|
||
|
||
try:
|
||
while fetched < limit:
|
||
if task.get("_cancelled"):
|
||
return
|
||
|
||
batch_limit = min(batch_size, limit - fetched)
|
||
batch = await asyncio.to_thread(
|
||
db_reader.get_messages_for_export,
|
||
req.talker,
|
||
req.start_time,
|
||
req.end_time,
|
||
batch_limit,
|
||
fetched,
|
||
)
|
||
if batch is None:
|
||
raise BridgeError(code="DB_NOT_FOUND", message="DB 不可读，无法导出")
|
||
if not batch:
|
||
break
|
||
|
||
# 媒体拷贝（include_media=true 且 media_inline=false）
|
||
if req.include_media and not req.media_inline:
|
||
_copy_media_files(batch, media_dir, db_reader)
|
||
|
||
# 写入文件（按格式）
|
||
if fmt == "json":
|
||
for msg in batch:
|
||
item_json = f" {json.dumps(_make_json_item(msg), ensure_ascii=False)}"
|
||
if json_pending is not None:
|
||
f.write((json_pending + ",\n").encode("utf-8"))
|
||
json_pending = item_json
|
||
elif fmt == "html":
|
||
chunk, last_ts = _render_html_messages(
|
||
batch, req.talker, req.media_inline, last_ts, db_reader
|
||
)
|
||
f.write(chunk)
|
||
elif fmt == "csv":
|
||
buf = io.StringIO()
|
||
w = csv.writer(buf)
|
||
for msg in batch:
|
||
w.writerow([
|
||
msg.get("msg_id", ""),
|
||
msg.get("create_time", 0),
|
||
"true" if msg.get("is_sender") else "false",
|
||
msg.get("type", 0),
|
||
msg.get("content", "") or "",
|
||
])
|
||
f.write(buf.getvalue().encode("utf-8"))
|
||
else: # txt
|
||
for chunk in _generate_txt(batch, nickname):
|
||
f.write(chunk)
|
||
|
||
fetched += len(batch)
|
||
task["processed"] = fetched
|
||
task["total"] = max(task["total"], fetched)
|
||
task["progress"] = min(1.0, fetched / limit) if limit > 0 else 1.0
|
||
first_batch = False
|
||
|
||
# 不足一批说明已无更多消息
|
||
if len(batch) < batch_limit:
|
||
break
|
||
finally:
|
||
# 格式尾部
|
||
if fmt == "json":
|
||
if json_pending is not None:
|
||
f.write((json_pending + "\n").encode("utf-8"))
|
||
# P1 修复：count 写实际条数（fetched），而非 limit
|
||
f.write(f'], "count": {fetched}}}\n'.encode("utf-8"))
|
||
elif fmt == "html":
|
||
f.write(_html_tail())
|
||
|
||
|
||
def _copy_media_files(messages: list[dict], media_dir: str, db_reader=None) -> None:
|
||
"""把消息中的媒体文件解密后写入 media_dir。
|
||
|
||
P0 修复：通过 db_reader.read_media_file_by_path 读取并解密 .dat 文件，
|
||
不再用 shutil.copyfile 直接拷贝原始密文。db_reader=None 时回退到直接拷贝
|
||
（兼容旧调用方，但 .dat 媒体将不可读）。
|
||
|
||
媒体缺失（path=None 或文件不可读）时跳过，由 HTML 渲染时写 [媒体缺失] 占位。
|
||
写入失败（权限/磁盘）仅记日志，不中断导出。
|
||
"""
|
||
for msg in messages:
|
||
src = msg.get("media_path")
|
||
if not src or not os.path.isfile(src) or not os.access(src, os.R_OK):
|
||
continue
|
||
dst = os.path.join(media_dir, os.path.basename(src))
|
||
# 同名冲突：追加 msg_id 前缀
|
||
if os.path.exists(dst):
|
||
base, ext = os.path.splitext(os.path.basename(src))
|
||
dst = os.path.join(media_dir, f"{base}_{msg.get('msg_id', '')}{ext}")
|
||
try:
|
||
# P0 修复：优先通过 db_reader 读取并解密 .dat
|
||
if db_reader is not None:
|
||
result = db_reader.read_media_file_by_path(src)
|
||
if result is not None:
|
||
data, _mime = result
|
||
with open(dst, "wb") as f:
|
||
f.write(data)
|
||
continue
|
||
# 解密失败（返回 None）：跳过该文件
|
||
logger.warning("解密媒体失败，跳过 src=%s", src)
|
||
continue
|
||
# 回退：直接拷贝（.dat 媒体将不可读）
|
||
shutil.copyfile(src, dst)
|
||
except OSError as e:
|
||
logger.warning("拷贝媒体失败 src=%s: %s", src, e)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# lifespan 钩子：进程退出时取消所有运行中任务
|
||
# ---------------------------------------------------------------------------
|
||
async def cancel_all_export_tasks_on_shutdown() -> None:
|
||
"""lifespan finally 段调用：取消所有运行中导出任务。
|
||
|
||
避免进程退出时残留 asyncio.Task 触发未捕获异常。
|
||
"""
|
||
for tid, task in list(_export_tasks.items()):
|
||
asyncio_task = task.get("_task")
|
||
if asyncio_task is not None and not asyncio_task.done():
|
||
task["_cancelled"] = True
|
||
asyncio_task.cancel()
|
||
try:
|
||
await asyncio.wait_for(asyncio_task, timeout=2.0)
|
||
except (asyncio.TimeoutError, asyncio.CancelledError, Exception):
|
||
pass
|
||
_export_tasks.clear()
|