763 lines
37 KiB
Python
763 lines
37 KiB
Python
"""Admin controls for resumable monthly PostgreSQL chat archives."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import errno
|
|
import logging
|
|
import shutil
|
|
import sqlite3
|
|
import uuid
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from fastapi import APIRouter, HTTPException, Request
|
|
from pydantic import BaseModel
|
|
from sqlalchemy import extract, func, select
|
|
|
|
from app.gateway.auth.disabled_mode import is_auth_disabled
|
|
from app.gateway.deps import get_current_user_from_request
|
|
from deerflow.persistence.engine import get_session_factory
|
|
from deerflow.persistence.thread_meta.base import is_excluded_system_thread
|
|
from deerflow.persistence.thread_meta.model import ThreadMetaRow
|
|
from deerflow.persistence.types import BEIJING_TZ
|
|
from deerflow.runtime.checkpointer.chat_archive import (
|
|
archive_db_path,
|
|
archive_root,
|
|
ensure_catalog,
|
|
thread_archive_lock,
|
|
try_acquire_archive_job_lock,
|
|
)
|
|
from deerflow.runtime.checkpointer.checkpoint_migration import _sqlite_has_table, copy_thread_state
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter(prefix="/api/admin/chat-archive", tags=["chat-archive"])
|
|
|
|
_ACTIVE = {"queued", "running", "cancelling"}
|
|
_PAGE_SIZE = 40
|
|
_ARCHIVE_DISK_RESERVE_BYTES = 2 * 1024**3
|
|
_PURGE_THREAD_DELAY_SECONDS = 0.5
|
|
_PROCESS_STARTED_AT = datetime.now(BEIJING_TZ).isoformat(timespec="seconds")
|
|
|
|
|
|
class _ArchiveJobCancelled(Exception):
|
|
pass
|
|
|
|
|
|
class ChatArchiveJobResponse(BaseModel):
|
|
job_id: str
|
|
month: str
|
|
action: str
|
|
status: str
|
|
total_threads: int
|
|
processed_threads: int
|
|
failed_threads: int
|
|
current_thread_id: str | None = None
|
|
snapshot_at: str = ""
|
|
processed_through: str | None = None
|
|
message: str = ""
|
|
error: str | None = None
|
|
created_at: str
|
|
updated_at: str
|
|
|
|
|
|
async def _require_admin(request: Request) -> None:
|
|
if is_auth_disabled():
|
|
return
|
|
user = await get_current_user_from_request(request)
|
|
if getattr(user, "system_role", None) != "admin":
|
|
raise HTTPException(status_code=403, detail="Chat archive management requires admin role")
|
|
|
|
|
|
def _now() -> str:
|
|
return datetime.now(BEIJING_TZ).isoformat(timespec="seconds")
|
|
|
|
|
|
def _parse_month(value: str) -> tuple[datetime, datetime]:
|
|
try:
|
|
start = datetime.strptime(value, "%Y-%m")
|
|
except ValueError as exc:
|
|
raise HTTPException(status_code=400, detail="month must use YYYY-MM format") from exc
|
|
if start.strftime("%Y-%m") != value:
|
|
raise HTTPException(status_code=400, detail="month must use YYYY-MM format")
|
|
if start.month == 12:
|
|
end = start.replace(year=start.year + 1, month=1)
|
|
else:
|
|
end = start.replace(month=start.month + 1)
|
|
return start, end
|
|
|
|
|
|
def _primary_checkpointer(request: Request):
|
|
checkpointer = getattr(request.app.state, "checkpointer", None)
|
|
primary = getattr(checkpointer, "primary", checkpointer)
|
|
if primary is None or not hasattr(getattr(primary, "conn", None), "connection"):
|
|
raise HTTPException(status_code=409, detail="当前聊天检查点不是 PostgreSQL，不能执行 PG 归档")
|
|
return checkpointer, primary
|
|
|
|
|
|
def _catalog_connect(root: Path) -> sqlite3.Connection:
|
|
ensure_catalog(root)
|
|
conn = sqlite3.connect(root / "catalog.db", timeout=30)
|
|
conn.row_factory = sqlite3.Row
|
|
conn.execute("PRAGMA busy_timeout=30000")
|
|
return conn
|
|
|
|
|
|
def _recover_jobs(root: Path, *, started_before: str | None = _PROCESS_STARTED_AT, owns_operation_lock: bool = False) -> None:
|
|
operation_lock = None if owns_operation_lock else try_acquire_archive_job_lock(root)
|
|
if not owns_operation_lock and operation_lock is None:
|
|
return # Another gateway/worker still owns the running job.
|
|
cutoff_sql = " AND created_at < ?" if started_before else ""
|
|
cutoff_values = (started_before,) if started_before else ()
|
|
with (contextlib.nullcontext() if owns_operation_lock else operation_lock), _catalog_connect(root) as conn:
|
|
now = _now()
|
|
conn.execute(
|
|
"UPDATE archive_jobs SET status='interrupted', message='执行进程已退出，已保留成功归档；可重新归档或清理已校验部分', "
|
|
"error=COALESCE(error, '归档执行进程已退出或服务重启，任务未完成'), updated_at=? "
|
|
"WHERE status IN ('queued','running','cancelling')" + cutoff_sql,
|
|
(now, *cutoff_values),
|
|
)
|
|
conn.execute(
|
|
"UPDATE archive_months SET status='interrupted', updated_at=? "
|
|
"WHERE status IN ('copying','deleting','restoring') AND current_job_id IN "
|
|
"(SELECT job_id FROM archive_jobs WHERE status='interrupted'" + cutoff_sql + ")",
|
|
(now, *cutoff_values),
|
|
)
|
|
|
|
|
|
def _month_start(month: str) -> tuple[datetime, datetime]:
|
|
start, end = _parse_month(month)
|
|
return start.replace(tzinfo=BEIJING_TZ), end.replace(tzinfo=BEIJING_TZ)
|
|
|
|
|
|
async def _month_thread_count(session_factory, month: str, *, snapshot_at: str | None = None) -> int:
|
|
start, end = _month_start(month)
|
|
conditions = [ThreadMetaRow.created_at >= start, ThreadMetaRow.created_at < end]
|
|
if snapshot_at:
|
|
conditions.append(ThreadMetaRow.created_at <= datetime.fromisoformat(snapshot_at))
|
|
async with session_factory() as session:
|
|
return int(await session.scalar(
|
|
select(func.count()).select_from(ThreadMetaRow).where(*conditions)
|
|
) or 0)
|
|
|
|
|
|
async def _thread_page(session_factory, month: str, offset: int, limit: int, *, snapshot_at: str | None = None):
|
|
start, end = _month_start(month)
|
|
conditions = [ThreadMetaRow.created_at >= start, ThreadMetaRow.created_at < end]
|
|
if snapshot_at:
|
|
conditions.append(ThreadMetaRow.created_at <= datetime.fromisoformat(snapshot_at))
|
|
async with session_factory() as session:
|
|
result = await session.execute(
|
|
select(
|
|
ThreadMetaRow.thread_id,
|
|
ThreadMetaRow.status,
|
|
ThreadMetaRow.metadata_json,
|
|
ThreadMetaRow.created_at,
|
|
)
|
|
.where(*conditions)
|
|
.order_by(ThreadMetaRow.created_at, ThreadMetaRow.thread_id)
|
|
.offset(offset)
|
|
.limit(limit)
|
|
)
|
|
return list(result.all())
|
|
|
|
|
|
async def _thread_status(session_factory, thread_id: str) -> str | None:
|
|
async with session_factory() as session:
|
|
value = await session.scalar(
|
|
select(ThreadMetaRow.status).where(ThreadMetaRow.thread_id == thread_id).limit(1)
|
|
)
|
|
return str(value).lower() if value is not None else None
|
|
|
|
|
|
def _archive_action_thread_count(root: Path, month: str, action: str) -> int:
|
|
"""Count only archived threads that the requested action will process."""
|
|
if not (root / "catalog.db").is_file():
|
|
return 0
|
|
if action == "purge":
|
|
statuses = ("copied", "deleted", "restored")
|
|
placeholders = ",".join("?" for _ in statuses)
|
|
query = f"SELECT COUNT(*) FROM archive_threads WHERE month=? AND status IN ({placeholders})"
|
|
params = (month, *statuses)
|
|
elif action == "rollback":
|
|
query = "SELECT COUNT(*) FROM archive_threads WHERE month=?"
|
|
params = (month,)
|
|
else:
|
|
return 0
|
|
with _catalog_connect(root) as conn:
|
|
row = conn.execute(query, params).fetchone()
|
|
return int(row[0] or 0)
|
|
|
|
|
|
async def _month_inventory(session_factory) -> list[dict[str, Any]]:
|
|
async with session_factory() as session:
|
|
dialect = session.get_bind().dialect.name
|
|
created_at = ThreadMetaRow.created_at
|
|
if dialect in {"postgresql", "opengauss"}:
|
|
local_created = func.timezone("Asia/Shanghai", created_at)
|
|
else:
|
|
local_created = created_at
|
|
year_expr = extract("year", local_created).label("year")
|
|
month_expr = extract("month", local_created).label("month")
|
|
rows = await session.execute(
|
|
select(year_expr, month_expr, func.count())
|
|
.select_from(ThreadMetaRow)
|
|
.group_by(year_expr, month_expr)
|
|
.order_by(year_expr.desc(), month_expr.desc())
|
|
)
|
|
return [
|
|
{"month": f"{int(year):04d}-{int(month):02d}", "thread_count": int(count)}
|
|
for year, month, count in rows.all()
|
|
if year is not None and month is not None
|
|
]
|
|
|
|
|
|
def _catalog_months(root: Path) -> dict[str, dict[str, Any]]:
|
|
if not (root / "catalog.db").exists():
|
|
return {}
|
|
with _catalog_connect(root) as conn:
|
|
months = {
|
|
row["month"]: dict(row)
|
|
for row in conn.execute("SELECT * FROM archive_months")
|
|
}
|
|
for row in conn.execute(
|
|
"SELECT j.* FROM archive_jobs j WHERE j.rowid=(SELECT newest.rowid FROM archive_jobs newest "
|
|
"WHERE newest.month=j.month ORDER BY newest.created_at DESC, newest.rowid DESC LIMIT 1)"
|
|
):
|
|
months.setdefault(row["month"], {})["latest_job"] = dict(row)
|
|
return months
|
|
|
|
|
|
def _save_job(root: Path, job: dict[str, Any]) -> None:
|
|
with _catalog_connect(root) as conn:
|
|
conn.execute(
|
|
"INSERT INTO archive_jobs(job_id,month,action,status,total_threads,processed_threads,failed_threads,"
|
|
"current_thread_id,snapshot_at,processed_through,message,error,created_at,updated_at) "
|
|
"VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
|
|
tuple((job.get(key) or job.get("created_at") or "") if key == "snapshot_at" else job.get(key) for key in (
|
|
"job_id", "month", "action", "status", "total_threads", "processed_threads",
|
|
"failed_threads", "current_thread_id", "snapshot_at", "processed_through", "message", "error", "created_at", "updated_at",
|
|
)),
|
|
)
|
|
|
|
|
|
def _claim_month_job(root: Path, job: dict[str, Any], month_status: str) -> bool:
|
|
"""Atomically allow only one heavy archive operation across processes."""
|
|
conn = _catalog_connect(root)
|
|
try:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
active = conn.execute(
|
|
"SELECT job_id FROM archive_jobs WHERE status IN ('queued','running','cancelling') LIMIT 1"
|
|
).fetchone()
|
|
if active:
|
|
conn.rollback()
|
|
return False
|
|
conn.execute(
|
|
"INSERT INTO archive_jobs(job_id,month,action,status,total_threads,processed_threads,failed_threads,"
|
|
"current_thread_id,snapshot_at,processed_through,message,error,created_at,updated_at) "
|
|
"VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
|
|
tuple((job.get(key) or job.get("created_at") or "") if key == "snapshot_at" else job.get(key) for key in (
|
|
"job_id", "month", "action", "status", "total_threads", "processed_threads",
|
|
"failed_threads", "current_thread_id", "snapshot_at", "processed_through", "message", "error", "created_at", "updated_at",
|
|
)),
|
|
)
|
|
conn.execute(
|
|
"INSERT INTO archive_months(month,status,total_threads,copied_threads,deleted_threads,restored_threads,current_job_id,updated_at) "
|
|
"VALUES(?,?,?,0,0,0,?,?) ON CONFLICT(month) DO UPDATE SET status=excluded.status, "
|
|
"total_threads=excluded.total_threads,current_job_id=excluded.current_job_id,updated_at=excluded.updated_at",
|
|
(job["month"], month_status, job["total_threads"], job["job_id"], _now()),
|
|
)
|
|
conn.commit()
|
|
return True
|
|
except Exception:
|
|
conn.rollback()
|
|
raise
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def _patch_job(root: Path, job_id: str, **values) -> None:
|
|
values["updated_at"] = _now()
|
|
assignments = ",".join(f"{key}=?" for key in values)
|
|
with _catalog_connect(root) as conn:
|
|
conn.execute(f"UPDATE archive_jobs SET {assignments} WHERE job_id=?", (*values.values(), job_id))
|
|
|
|
|
|
async def _patch_job_async(root: Path, job_id: str, **values) -> None:
|
|
await asyncio.to_thread(_patch_job, root, job_id, **values)
|
|
|
|
|
|
def _job_dict(root: Path, job_id: str) -> dict[str, Any] | None:
|
|
if not (root / "catalog.db").exists():
|
|
return None
|
|
with _catalog_connect(root) as conn:
|
|
row = conn.execute("SELECT * FROM archive_jobs WHERE job_id=?", (job_id,)).fetchone()
|
|
return dict(row) if row else None
|
|
|
|
|
|
def _upsert_month(root: Path, month: str, *, status: str, total: int, job_id: str | None = None) -> None:
|
|
with _catalog_connect(root) as conn:
|
|
conn.execute(
|
|
"INSERT INTO archive_months(month,status,total_threads,copied_threads,deleted_threads,restored_threads,current_job_id,updated_at) "
|
|
"VALUES(?,?,?,0,0,0,?,?) ON CONFLICT(month) DO UPDATE SET status=excluded.status, "
|
|
"total_threads=excluded.total_threads,current_job_id=excluded.current_job_id,updated_at=excluded.updated_at",
|
|
(month, status, total, job_id, _now()),
|
|
)
|
|
|
|
|
|
async def _upsert_month_async(root: Path, month: str, *, status: str, total: int, job_id: str | None = None) -> None:
|
|
await asyncio.to_thread(_upsert_month, root, month, status=status, total=total, job_id=job_id)
|
|
|
|
|
|
def _mark_archived_thread(root: Path, thread_id: str, month: str, count: int, status: str = "copied") -> None:
|
|
now = _now()
|
|
with _catalog_connect(root) as conn:
|
|
conn.execute(
|
|
"INSERT INTO archive_threads(thread_id,month,status,checkpoint_count,archived_at,deleted_at,restored_at) "
|
|
"VALUES(?,?,?,?,?,?,?) ON CONFLICT(thread_id) DO UPDATE SET month=excluded.month,status=excluded.status,"
|
|
"checkpoint_count=excluded.checkpoint_count,archived_at=excluded.archived_at,"
|
|
"deleted_at=CASE WHEN excluded.status='deleted' THEN excluded.archived_at ELSE archive_threads.deleted_at END,"
|
|
"restored_at=CASE WHEN excluded.status='restored' THEN excluded.archived_at ELSE archive_threads.restored_at END",
|
|
(thread_id, month, status, count, now, now if status == "deleted" else None, now if status == "restored" else None),
|
|
)
|
|
conn.execute(
|
|
"UPDATE archive_months SET copied_threads=(SELECT COUNT(*) FROM archive_threads WHERE month=?), "
|
|
"deleted_threads=(SELECT COUNT(*) FROM archive_threads WHERE month=? AND status='deleted'), updated_at=? WHERE month=?",
|
|
(month, month, now, month),
|
|
)
|
|
|
|
|
|
async def _checkpoint_ids(saver, thread_id: str) -> list[str]:
|
|
config = {"configurable": {"thread_id": thread_id}}
|
|
ids: list[str] = []
|
|
async for tup in saver.alist(config):
|
|
configurable = (tup.config or {}).get("configurable", {}) or {}
|
|
checkpoint_id = configurable.get("checkpoint_id") or (tup.checkpoint or {}).get("id")
|
|
if checkpoint_id:
|
|
ids.append(str(checkpoint_id))
|
|
return ids
|
|
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def _open_sqlite_archive(path: Path):
|
|
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
context = AsyncSqliteSaver.from_conn_string(str(path))
|
|
saver = await context.__aenter__()
|
|
try:
|
|
await saver.conn.execute("PRAGMA busy_timeout=30000")
|
|
if not await asyncio.to_thread(_sqlite_has_table, str(path), "checkpoints"):
|
|
await saver.conn.execute("PRAGMA journal_mode=WAL")
|
|
await saver.setup()
|
|
yield saver
|
|
finally:
|
|
await context.__aexit__(None, None, None)
|
|
|
|
|
|
class ChatArchiveRunner:
|
|
def __init__(self, root: Path, checkpointer, session_factory):
|
|
self.root = root
|
|
self.checkpointer = checkpointer
|
|
self.primary = getattr(checkpointer, "primary", checkpointer)
|
|
self.session_factory = session_factory
|
|
self._locks: dict[str, asyncio.Lock] = {}
|
|
self._tasks: dict[str, asyncio.Task] = {}
|
|
|
|
async def start(self, month: str, action: str) -> dict[str, Any]:
|
|
_parse_month(month)
|
|
if action not in {"archive", "purge", "rollback"}:
|
|
raise HTTPException(status_code=400, detail="未知归档操作")
|
|
current_month = datetime.now(BEIJING_TZ).strftime("%Y-%m")
|
|
if month > current_month:
|
|
raise HTTPException(status_code=409, detail="不能归档未来月份")
|
|
snapshot_at = _now()
|
|
if action == "archive":
|
|
total = await _month_thread_count(self.session_factory, month, snapshot_at=snapshot_at)
|
|
else:
|
|
total = await asyncio.to_thread(_archive_action_thread_count, self.root, month, action)
|
|
if not total:
|
|
raise HTTPException(status_code=409, detail="该月份尚无已校验归档，不能删除或回滚；请先归档")
|
|
job_id = uuid.uuid4().hex
|
|
job = {
|
|
"job_id": job_id,
|
|
"month": month,
|
|
"action": action,
|
|
"status": "queued",
|
|
"total_threads": total,
|
|
"processed_threads": 0,
|
|
"failed_threads": 0,
|
|
"current_thread_id": None,
|
|
"snapshot_at": snapshot_at,
|
|
"processed_through": None,
|
|
"message": "任务已排队",
|
|
"error": None,
|
|
"created_at": _now(),
|
|
"updated_at": _now(),
|
|
}
|
|
status = {"archive": "copying", "purge": "deleting", "rollback": "restoring"}[action]
|
|
operation_lock = try_acquire_archive_job_lock(self.root)
|
|
if operation_lock is None:
|
|
raise HTTPException(status_code=409, detail="已有月份归档、清理或回滚任务正在运行，请完成后再启动下一项")
|
|
try:
|
|
# The kernel lock is free only if no runner is alive, including
|
|
# one that died after this gateway process finished starting.
|
|
_recover_jobs(self.root, started_before=None, owns_operation_lock=True)
|
|
claimed = _claim_month_job(self.root, job, status)
|
|
if not claimed:
|
|
raise HTTPException(status_code=409, detail="已有月份归档、清理或回滚任务正在运行，请完成后再启动下一项")
|
|
task = asyncio.create_task(self._run(job_id, owns_operation_lock=True), name=f"chat-archive-{action}-{month}")
|
|
except BaseException:
|
|
operation_lock.close()
|
|
raise
|
|
self._tasks[job_id] = task
|
|
def completed(_task):
|
|
operation_lock.close()
|
|
self._tasks.pop(job_id, None)
|
|
task.add_done_callback(completed)
|
|
return job
|
|
|
|
async def request_cancel(self, job_id: str) -> dict[str, Any]:
|
|
job = await asyncio.to_thread(_job_dict, self.root, job_id)
|
|
if not job:
|
|
raise HTTPException(status_code=404, detail="归档任务不存在")
|
|
if job["status"] in _ACTIVE:
|
|
await _patch_job_async(self.root, job_id, status="cancelling", message="将在当前会话处理结束后停止")
|
|
return await asyncio.to_thread(_job_dict, self.root, job_id)
|
|
|
|
async def _check_cancelled(self, job_id: str) -> None:
|
|
job = await asyncio.to_thread(_job_dict, self.root, job_id) or {}
|
|
if job.get("status") == "cancelling":
|
|
raise _ArchiveJobCancelled()
|
|
|
|
async def _record_thread_failure(self, job_id: str, tid: str, exc: Exception) -> None:
|
|
# A corrupt/full archive volume cannot safely accept more threads.
|
|
# Row-specific failures may continue; infrastructure failures must stop.
|
|
sqlite_code = getattr(exc, "sqlite_errorcode", 0) or 0
|
|
if (
|
|
isinstance(exc, MemoryError)
|
|
or isinstance(exc, OSError) and exc.errno in {errno.ENOSPC, errno.EDQUOT}
|
|
or isinstance(exc, sqlite3.Error) and (sqlite_code & 0xFF) in {
|
|
sqlite3.SQLITE_FULL, sqlite3.SQLITE_IOERR, sqlite3.SQLITE_CORRUPT, sqlite3.SQLITE_NOTADB,
|
|
}
|
|
):
|
|
raise exc
|
|
logger.error("Chat archive thread failed; continuing: job=%s thread=%s", job_id, tid, exc_info=True)
|
|
job = await asyncio.to_thread(_job_dict, self.root, job_id) or {}
|
|
await _patch_job_async(
|
|
self.root, job_id,
|
|
failed_threads=int(job.get("failed_threads") or 0) + 1,
|
|
error=f"会话 {tid}: {type(exc).__name__}: {exc}"[:2000],
|
|
message="该会话处理失败，已记录错误，继续处理其他会话",
|
|
)
|
|
|
|
async def _run(self, job_id: str, *, owns_operation_lock: bool = False) -> None:
|
|
operation_lock = None if owns_operation_lock else try_acquire_archive_job_lock(self.root)
|
|
if not owns_operation_lock and operation_lock is None:
|
|
return
|
|
try:
|
|
await self._run_locked(job_id)
|
|
finally:
|
|
if operation_lock is not None:
|
|
operation_lock.close()
|
|
|
|
async def _run_locked(self, job_id: str) -> None:
|
|
job = _job_dict(self.root, job_id)
|
|
if not job:
|
|
return
|
|
month = job["month"]
|
|
action = job["action"]
|
|
lock = self._locks.setdefault(month, asyncio.Lock())
|
|
async with lock:
|
|
try:
|
|
await self._check_cancelled(job_id)
|
|
await _patch_job_async(self.root, job_id, status="running", message="正在处理，页面可安全关闭")
|
|
if action == "archive":
|
|
await self._archive_month(job_id, month)
|
|
elif action == "purge":
|
|
await self._purge_month(job_id, month)
|
|
else:
|
|
await self._rollback_month(job_id, month)
|
|
job = await asyncio.to_thread(_job_dict, self.root, job_id) or job
|
|
failed = int(job.get("failed_threads") or 0)
|
|
status = "completed" if failed == 0 else "partial"
|
|
month_status = {
|
|
"archive": "ready" if failed == 0 else "partial",
|
|
"purge": "deleted" if failed == 0 else "partial",
|
|
"rollback": "ready" if failed == 0 else "partial",
|
|
}[action]
|
|
await _patch_job_async(self.root, job_id, status=status, current_thread_id=None,
|
|
message="处理完成" if failed == 0 else f"完成但有 {failed} 个会话未处理")
|
|
await _upsert_month_async(self.root, month, status=month_status, total=int(job.get("total_threads") or 0), job_id=job_id)
|
|
except (_ArchiveJobCancelled, asyncio.CancelledError) as exc:
|
|
await _patch_job_async(self.root, job_id, status="cancelled" if isinstance(exc, _ArchiveJobCancelled) else "interrupted",
|
|
current_thread_id=None, message="任务已停止，已完成的归档保留，可重新执行",
|
|
error="管理员停止任务" if isinstance(exc, _ArchiveJobCancelled) else "执行任务被取消或服务正在退出")
|
|
await _upsert_month_async(self.root, month, status="interrupted", total=int(job.get("total_threads") or 0), job_id=job_id)
|
|
if isinstance(exc, asyncio.CancelledError):
|
|
raise
|
|
except Exception as exc:
|
|
logger.exception("Chat archive job failed: job=%s month=%s action=%s", job_id, month, action)
|
|
await _patch_job_async(self.root, job_id, status="failed", error=f"{type(exc).__name__}: {exc}"[:2000],
|
|
message="任务停止，已完成的归档保留；可查看错误并重新执行")
|
|
await _upsert_month_async(self.root, month, status="interrupted", total=int(job.get("total_threads") or 0), job_id=job_id)
|
|
|
|
async def _archive_month(self, job_id: str, month: str) -> None:
|
|
path = archive_db_path(self.root, month)
|
|
if (await asyncio.to_thread(shutil.disk_usage, self.root)).free < _ARCHIVE_DISK_RESERVE_BYTES:
|
|
raise RuntimeError("归档目录所在磁盘可用空间不足 2 GiB；请先扩容或清理磁盘，PostgreSQL 数据未删除")
|
|
job = await asyncio.to_thread(_job_dict, self.root, job_id) or {}
|
|
snapshot_at = str(job.get("snapshot_at") or job.get("created_at") or "")
|
|
total = int(job.get("total_threads") or 0)
|
|
await _upsert_month_async(self.root, month, status="copying", total=total, job_id=job_id)
|
|
processed = 0
|
|
offset = 0
|
|
async with _open_sqlite_archive(path) as archive:
|
|
while True:
|
|
rows = await _thread_page(
|
|
self.session_factory, month, offset, _PAGE_SIZE, snapshot_at=snapshot_at
|
|
)
|
|
if not rows:
|
|
break
|
|
offset += len(rows)
|
|
for thread_id, thread_status, metadata, created_at in rows:
|
|
await self._check_cancelled(job_id)
|
|
processed_through = created_at.isoformat() if created_at else None
|
|
if is_excluded_system_thread(metadata if isinstance(metadata, dict) else {}):
|
|
processed += 1
|
|
await _patch_job_async(
|
|
self.root, job_id, processed_threads=processed,
|
|
processed_through=processed_through, message="已跳过系统任务会话"
|
|
)
|
|
continue
|
|
tid = str(thread_id)
|
|
if (await asyncio.to_thread(shutil.disk_usage, self.root)).free < _ARCHIVE_DISK_RESERVE_BYTES:
|
|
raise RuntimeError("归档卷剩余空间低于 2 GiB，任务已暂停；PostgreSQL 数据保留")
|
|
await _patch_job_async(self.root, job_id, current_thread_id=tid, message="正在复制会话检查点")
|
|
if str(thread_status).lower() in {"running", "busy"}:
|
|
current_job = await asyncio.to_thread(_job_dict, self.root, job_id) or {}
|
|
processed += 1
|
|
await _patch_job_async(
|
|
self.root, job_id,
|
|
failed_threads=int(current_job.get("failed_threads") or 0) + 1,
|
|
processed_threads=processed, processed_through=processed_through,
|
|
message="正在跳过运行中的会话"
|
|
)
|
|
continue
|
|
try:
|
|
async with thread_archive_lock(self.root, tid, exclusive=True):
|
|
if await _thread_status(self.session_factory, tid) in {"running", "busy"}:
|
|
raise RuntimeError("会话正在问答中，已跳过；结束后可重新归档")
|
|
source_ids = await _checkpoint_ids(self.primary, tid)
|
|
if not source_ids:
|
|
processed += 1
|
|
await _patch_job_async(
|
|
self.root, job_id, processed_threads=processed,
|
|
processed_through=processed_through, message="已跳过无检查点的空会话"
|
|
)
|
|
continue
|
|
await copy_thread_state(self.primary, archive, tid)
|
|
copied_ids = await _checkpoint_ids(archive, tid)
|
|
if not set(source_ids).issubset(copied_ids):
|
|
raise RuntimeError(f"thread {tid} checkpoint verification failed: source={len(source_ids)} archived={len(copied_ids)}")
|
|
await asyncio.to_thread(_mark_archived_thread, self.root, tid, month, len(copied_ids), "copied")
|
|
except Exception as exc:
|
|
await self._record_thread_failure(job_id, tid, exc)
|
|
processed += 1
|
|
current_job = await asyncio.to_thread(_job_dict, self.root, job_id) or {}
|
|
await _patch_job_async(
|
|
self.root, job_id, processed_threads=processed,
|
|
failed_threads=int(current_job.get("failed_threads") or 0),
|
|
processed_through=processed_through,
|
|
)
|
|
await asyncio.sleep(0.03)
|
|
job = await asyncio.to_thread(_job_dict, self.root, job_id) or {}
|
|
if int(job.get("failed_threads") or 0):
|
|
await _patch_job_async(self.root, job_id, message="部分会话正忙或校验失败；可稍后重新归档")
|
|
|
|
async def _purge_month(self, job_id: str, month: str) -> None:
|
|
path = archive_db_path(self.root, month)
|
|
if not path.is_file():
|
|
raise RuntimeError(f"monthly SQLite archive is missing: {path}")
|
|
with _catalog_connect(self.root) as conn:
|
|
thread_ids = [str(row[0]) for row in conn.execute(
|
|
"SELECT thread_id FROM archive_threads WHERE month=? AND status IN ('copied','deleted','restored') ORDER BY thread_id",
|
|
(month,),
|
|
)]
|
|
processed = 0
|
|
async with _open_sqlite_archive(path) as archive:
|
|
for tid in thread_ids:
|
|
await self._check_cancelled(job_id)
|
|
await _patch_job_async(self.root, job_id, current_thread_id=tid, message="重新校验并清理 PostgreSQL 检查点")
|
|
skipped_running = False
|
|
try:
|
|
async with thread_archive_lock(self.root, tid, exclusive=True):
|
|
thread_status = await _thread_status(self.session_factory, tid)
|
|
if thread_status in {"running", "busy"}:
|
|
skipped_running = True
|
|
else:
|
|
pg_ids = await _checkpoint_ids(self.primary, tid)
|
|
if not pg_ids:
|
|
await asyncio.to_thread(
|
|
_mark_archived_thread,
|
|
self.root,
|
|
tid,
|
|
month,
|
|
len(await _checkpoint_ids(archive, tid)),
|
|
"deleted",
|
|
)
|
|
else:
|
|
await copy_thread_state(self.primary, archive, tid)
|
|
archived_ids = await _checkpoint_ids(archive, tid)
|
|
if not set(pg_ids).issubset(archived_ids):
|
|
raise RuntimeError(f"thread {tid} archive is incomplete; PostgreSQL rows were kept")
|
|
await self.primary.adelete_thread(tid)
|
|
remaining = await _checkpoint_ids(self.primary, tid)
|
|
if remaining:
|
|
raise RuntimeError(f"thread {tid} PostgreSQL cleanup verification failed")
|
|
await asyncio.to_thread(_mark_archived_thread, self.root, tid, month, len(archived_ids), "deleted")
|
|
except Exception as exc:
|
|
await self._record_thread_failure(job_id, tid, exc)
|
|
processed += 1
|
|
job_update: dict[str, Any] = {"processed_threads": processed}
|
|
if skipped_running:
|
|
job = await asyncio.to_thread(_job_dict, self.root, job_id) or {}
|
|
job_update.update(
|
|
failed_threads=int(job.get("failed_threads") or 0) + 1,
|
|
message="会话仍在问答中，已跳过；结束后可重新清理",
|
|
)
|
|
await _patch_job_async(self.root, job_id, **job_update)
|
|
await asyncio.sleep(_PURGE_THREAD_DELAY_SECONDS)
|
|
|
|
async def _rollback_month(self, job_id: str, month: str) -> None:
|
|
path = archive_db_path(self.root, month)
|
|
if not path.is_file():
|
|
raise RuntimeError(f"monthly SQLite archive is missing: {path}")
|
|
with _catalog_connect(self.root) as conn:
|
|
thread_ids = [str(row[0]) for row in conn.execute(
|
|
"SELECT thread_id FROM archive_threads WHERE month=? ORDER BY thread_id", (month,)
|
|
)]
|
|
processed = 0
|
|
async with _open_sqlite_archive(path) as archive:
|
|
for tid in thread_ids:
|
|
await self._check_cancelled(job_id)
|
|
await _patch_job_async(self.root, job_id, current_thread_id=tid, message="回滚检查点到 PostgreSQL")
|
|
try:
|
|
async with thread_archive_lock(self.root, tid, exclusive=True):
|
|
if await _thread_status(self.session_factory, tid) in {"running", "busy"}:
|
|
raise RuntimeError("会话正在问答中，已跳过；结束后可重新回滚")
|
|
archive_ids = await _checkpoint_ids(archive, tid)
|
|
await copy_thread_state(archive, self.primary, tid)
|
|
pg_ids = await _checkpoint_ids(self.primary, tid)
|
|
if not set(archive_ids).issubset(pg_ids):
|
|
raise RuntimeError(f"thread {tid} rollback verification failed")
|
|
await asyncio.to_thread(_mark_archived_thread, self.root, tid, month, len(archive_ids), "restored")
|
|
except Exception as exc:
|
|
await self._record_thread_failure(job_id, tid, exc)
|
|
processed += 1
|
|
await _patch_job_async(self.root, job_id, processed_threads=processed)
|
|
await asyncio.sleep(0.03)
|
|
|
|
|
|
async def _runner(request: Request) -> ChatArchiveRunner:
|
|
runner = getattr(request.app.state, "chat_archive_runner", None)
|
|
if runner is None:
|
|
checkpointer, _ = _primary_checkpointer(request)
|
|
session_factory = get_session_factory()
|
|
if session_factory is None:
|
|
raise HTTPException(status_code=503, detail="聊天元数据数据库未启用")
|
|
root = archive_root()
|
|
ensure_catalog(root)
|
|
await _ensure_archive_recovery_async(request, root)
|
|
runner = ChatArchiveRunner(root, checkpointer, session_factory)
|
|
request.app.state.chat_archive_runner = runner
|
|
return runner
|
|
|
|
|
|
async def _ensure_archive_recovery_async(request: Request, root: Path) -> None:
|
|
"""Recover abandoned work without interrupting another live process.
|
|
|
|
Recheck on inventory requests: another process may die after this process
|
|
starts or after an earlier recovery was skipped because the lock was held.
|
|
"""
|
|
if not (root / "catalog.db").is_file():
|
|
# Merely opening the settings page must not initialize archive storage.
|
|
return
|
|
lock = getattr(request.app.state, "chat_archive_recovery_lock", None)
|
|
if lock is None:
|
|
lock = asyncio.Lock()
|
|
request.app.state.chat_archive_recovery_lock = lock
|
|
async with lock:
|
|
await asyncio.to_thread(_recover_jobs, root, started_before=None)
|
|
|
|
|
|
@router.get("/months")
|
|
async def list_archive_months(request: Request):
|
|
await _require_admin(request)
|
|
_, _primary = _primary_checkpointer(request)
|
|
session_factory = get_session_factory()
|
|
if session_factory is None:
|
|
raise HTTPException(status_code=503, detail="聊天元数据数据库未启用")
|
|
root = archive_root()
|
|
await _ensure_archive_recovery_async(request, root)
|
|
inventory = await _month_inventory(session_factory)
|
|
catalog = _catalog_months(root)
|
|
disk_path = root
|
|
while not disk_path.exists():
|
|
disk_path = disk_path.parent
|
|
disk = await asyncio.to_thread(shutil.disk_usage, disk_path)
|
|
months = []
|
|
for item in inventory:
|
|
month = item["month"]
|
|
state = catalog.get(month, {})
|
|
file_path = archive_db_path(root, month)
|
|
months.append({
|
|
**item,
|
|
"status": state.get("status", "not_archived"),
|
|
"copied_threads": int(state.get("copied_threads") or 0),
|
|
"deleted_threads": int(state.get("deleted_threads") or 0),
|
|
"sqlite_bytes": file_path.stat().st_size if file_path.is_file() else 0,
|
|
"latest_job": state.get("latest_job"),
|
|
})
|
|
return {
|
|
"archive_dir": str(root),
|
|
"disk_total_bytes": disk.total,
|
|
"disk_free_bytes": disk.free,
|
|
"month_basis": "会话创建月份（按业务数据库 threads_meta 时间）",
|
|
"warning": "页面只读取轻量会话元数据，不创建归档文件或扫描 PostgreSQL 检查点表；只有管理员启动归档、清理或回滚任务时才会读写归档数据。",
|
|
"months": months,
|
|
}
|
|
|
|
|
|
async def _start_action(month: str, action: str, request: Request):
|
|
await _require_admin(request)
|
|
_primary_checkpointer(request)
|
|
return await (await _runner(request)).start(month, action)
|
|
|
|
|
|
@router.post("/months/{month}/archive")
|
|
async def start_month_archive(month: str, request: Request):
|
|
return await _start_action(month, "archive", request)
|
|
|
|
|
|
@router.post("/months/{month}/delete-postgres")
|
|
async def delete_archived_month_from_postgres(month: str, request: Request):
|
|
return await _start_action(month, "purge", request)
|
|
|
|
|
|
@router.post("/months/{month}/rollback")
|
|
async def rollback_month_to_postgres(month: str, request: Request):
|
|
return await _start_action(month, "rollback", request)
|
|
|
|
|
|
@router.get("/jobs/{job_id}", response_model=ChatArchiveJobResponse)
|
|
async def get_chat_archive_job(job_id: str, request: Request):
|
|
await _require_admin(request)
|
|
job = _job_dict(archive_root(), job_id)
|
|
if job is None:
|
|
raise HTTPException(status_code=404, detail="归档任务不存在")
|
|
return job
|