ai-kefu-admin/scripts/archive.py

"""数据归档脚本(Phase 4 Task 3): 租户库老会话 导出→清理→VACUUM。

conversations.created_at 早于 N 天(默认 365)的会话整链归档: 五表行先按
会话分组导出 JSON 到 data/archives/<tenant>/<日期>-<N>d.json, 落盘成功后
才删库(倒 FK: lead_push_logs → evaluations → customer_profiles →
messages → conversations, 单事务), 删了行才 VACUUM。导出/删除集同源于
同一截止串, 无漂移。

纪律:
- 门禁同 import_legacy: 租户须已在 platform.db bootstrap, 否则报错提示
  先跑 bootstrap, 绝不摸租户库;
- 截止时刻 Python 侧算好 ISO 串再入 SQL(不依赖 SQLite datetime() 解析):
  ORM 行存 'YYYY-MM-DD HH:MM:SS.ffffff'(无 tz 后缀), 老库导入行可能带
  '+00:00' 后缀——同族字符串字典序比较对两者皆单调(差异位在 '.'/'+',
  仅整秒边界有微差, 天级截止无实际影响);
- 导出先落盘后删除: 文件写失败 → 数据一字节不动(多一份孤儿导出无害,
  反之是不可逆丢失); 落盘 = 临时文件 fsync 后原子 replace(崩溃级: 断电
  窗口内要么旧文件在、要么新文件完整, "导出先于删除"不击穿); 目录懒建
  (无可归档/失败时不产生任何文件);
- 导出/删除行数对账: 两边计数差一行即非零退出(删除已提交, 不一致必须
  人工响, 不静默放过);
- --dry-run 只列不动: 库零变化、不导出、不 VACUUM;
- VACUUM 被占用锁(SQLITE_BUSY)不判归档失败: 导出/删除已安全完成, 提示
  稍后手动执行;
- archives 目录 = tenants 目录的兄弟 data/archives(同 registry
  uploads_root 的派生纪律: 生产 <repo>/tenants → <repo>/data/archives,
  测试注入 tmp tenants → tmp 内隔离; data/ 在 .gitignore 内天然覆盖);
  导出文件名带时分秒(<日期>-<时分秒>-<N>d.json), 同日同参重跑不覆盖
  唯一副本。

用法(Windows 须 PYTHONUTF8=1, 输出含中文):
  PYTHONUTF8=1 python scripts/archive.py --tenant <slug> \
      [--days 365] [--dry-run] [--platform-db <路径>] [--tenants-dir <目录>]
"""
from __future__ import annotations

import argparse
import asyncio
import json
import os
import sqlite3
import sys
from datetime import datetime, timedelta, timezone
from pathlib import Path

# 仓库根入 sys.path(照 import_legacy.py 锚定纪律: 脚本可从任意 CWD 跑)
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import aiosqlite

REPO_ROOT = Path(__file__).resolve().parents[1]

# 导出分组顺序(父先子后, 与 import_legacy 的 TABLE_ORDER 一致)
TABLE_ORDER = (
    "conversations", "messages", "customer_profiles", "evaluations",
    "lead_push_logs",
)
# 删除序 = 倒 FK(子先父后; SQLite 默认不开 FK 强制, 显式按序删)
DELETE_ORDER = tuple(reversed(TABLE_ORDER))
CHILD_TABLES = TABLE_ORDER[1:]            # 以 conversation_id 关联的子表

# 各表确定性排序(导出顺序稳定; 注意 customer_profiles 无 created_at
# 只有 updated_at, evaluations 是 evaluated_at——五表时间列不齐名)
TABLE_ORDER_BY = {
    "conversations": "created_at, id",
    "messages": "created_at, id",
    "customer_profiles": "updated_at, id",
    "evaluations": "evaluated_at, id",
    "lead_push_logs": "created_at, id",
}


def _tenant_bootstrapped(platform_db: Path, slug: str) -> None:
    """门禁(同 import_legacy): 租户须已在 platform.db 注册;
    只读打开(先看后动, 门禁绝不写 platform.db)。"""
    if not platform_db.is_file():
        raise SystemExit(
            f"platform.db 不存在: {platform_db}\n"
            f"租户须先 bootstrap: 先跑 scripts/bootstrap_platform.py")
    try:
        con = sqlite3.connect(
            platform_db.resolve().as_uri() + "?mode=ro", uri=True)
        try:
            row = con.execute(
                "SELECT 1 FROM tenants WHERE slug = ?", (slug,)).fetchone()
        finally:
            con.close()
    except sqlite3.Error as exc:
        raise SystemExit(
            f"platform.db 读取失败({exc}): 先跑 scripts/bootstrap_platform.py"
        ) from exc
    if row is None:
        raise SystemExit(
            f"租户 {slug} 未在 platform.db 注册: "
            f"先跑 scripts/bootstrap_platform.py 导入该租户档案")


def _cutoff(days: int) -> str:
    """截止 ISO 串(Python 侧算, 见模块 docstring 的时间形态论证)。"""
    cut = datetime.now(timezone.utc).replace(tzinfo=None) - timedelta(days=days)
    return cut.strftime("%Y-%m-%d %H:%M:%S.%f")


async def _columns(conn, table: str) -> list[str]:
    cur = await conn.execute(f"PRAGMA table_info({table})")
    cols = [row[1] for row in await cur.fetchall()]
    if not cols:
        raise SystemExit(
            f"租户库缺表: {table}(schema 未就位, 先跑应用或 import_legacy)")
    return cols


async def _scalar(conn, sql: str, params=()):
    cur = await conn.execute(sql, params)
    row = await cur.fetchone()
    return row[0]


async def _old_conversations(conn, cutoff: str) -> list:
    """待归档会话 (id, session_id, created_at), 时间序。"""
    cur = await conn.execute(
        "SELECT id, session_id, created_at FROM conversations "
        "WHERE created_at < ? ORDER BY created_at, id", (cutoff,))
    return list(await cur.fetchall())


async def _export_document(conn, cutoff: str) -> dict:
    """按会话分组的五表完整行(conversations 键=会话行; 子表键=会话 id →
    行列表; 无该子表行的会话无键)。SELECT * 全列, 值原样读出(JSON 列即
    库内存储文本, 字节级保真)。"""
    conn.row_factory = aiosqlite.Row
    doc: dict = {t: {} for t in TABLE_ORDER}
    for table in TABLE_ORDER:
        where_id = table == "conversations"
        sql = (f"SELECT * FROM {table} WHERE "
               f"{'id' if where_id else 'conversation_id'} IN "
               f"(SELECT id FROM conversations WHERE created_at < ?) "
               f"ORDER BY {TABLE_ORDER_BY[table]}")
        cur = await conn.execute(sql, (cutoff,))
        for row in await cur.fetchall():
            fields = dict(row)
            if where_id:
                doc[table][row["id"]] = fields
            else:
                doc[table].setdefault(row["conversation_id"],
                                      []).append(fields)
    return doc


def _write_export(path: Path, doc: dict) -> None:
    """先落盘后删库的'落盘'端: 目录懒建 + 临时文件写完 flush+fsync 再
    原子 replace——进程级升级到崩溃级(断电窗口内要么旧文件在、要么新文件
    完整落盘, '导出先于删除'的承诺不被断电击穿)。"""
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".json.tmp")
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(doc, f, ensure_ascii=False, indent=2)
        f.flush()
        os.fsync(f.fileno())
    tmp.replace(path)


async def run(args) -> None:
    days: int = args.days
    if days < 0:
        raise SystemExit(f"--days 不得为负: {days}")
    platform_db = Path(args.platform_db)
    _tenant_bootstrapped(platform_db, args.tenant)      # 门禁先于一切

    tenants_dir = Path(args.tenants_dir)
    tenant_db = tenants_dir / f"{args.tenant}.db"
    if not tenant_db.is_file():
        raise SystemExit(f"租户库不存在: {tenant_db}")
    archives_dir = tenants_dir.parent / "data" / "archives"
    cutoff = _cutoff(days)
    size_before = tenant_db.stat().st_size

    async with aiosqlite.connect(str(tenant_db), timeout=30) as conn:
        for table in TABLE_ORDER:                       # 缺表显性报错
            await _columns(conn, table)
        old = await _old_conversations(conn, cutoff)
        old_ids = [row[0] for row in old]

        print(f"租户库: {tenant_db}")
        print(f"归档策略: created_at < {cutoff}({days} 天前)")

        if args.dry_run:
            print("模式: dry-run(只列不动, 库零变化)")
            print(f"待归档会话: {len(old_ids)} 条")
            for cid, session_id, created_at in old:
                print(f"  {cid} session={session_id} "
                      f"created_at={created_at}")
            if old_ids:
                print("将删除的各表行数(未动):")
                for table in ("conversations", *CHILD_TABLES):
                    n = len(old_ids) if table == "conversations" else \
                        await _scalar(conn, (
                            f"SELECT COUNT(*) FROM {table} "
                            "WHERE conversation_id IN (SELECT id FROM "
                            "conversations WHERE created_at < ?)"), (cutoff,))
                    print(f"  {table}: {n}")
            print("dry-run 结论: 零改动(未导出/未删除/未 VACUUM)")
            return

        export_path = archives_dir / args.tenant / (
            f"{datetime.now(timezone.utc).strftime('%Y%m%d-%H%M%S')}"
            f"-{days}d.json")
        deleted: dict[str, int] = {}
        if not old_ids:
            print("归档会话: 0 条")
            print("结论: 无可归档会话, 零改动(未导出/未删除/未 VACUUM)")
            print(f"VACUUM: 未执行(无删除行)")
            print(f"库文件: {size_before} → {tenant_db.stat().st_size} 字节")
            return

        # ---- 导出(先落盘, 失败即中止, 数据不动) ---------------------------
        doc = {"tenant": args.tenant,
               "generated_at": datetime.now(timezone.utc).isoformat(),
               "days": days, "cutoff": cutoff,
               "conversation_count": len(old_ids)}
        doc.update(await _export_document(conn, cutoff))
        try:
            _write_export(export_path, doc)
        except OSError as exc:
            raise SystemExit(
                f"导出写盘失败({exc}): 中止, 库零变化(未删除)") from exc

        # ---- 清库(倒 FK, 单事务; 子查询同源截止串) ------------------------
        for table in DELETE_ORDER:
            col = "id" if table == "conversations" else "conversation_id"
            cur = await conn.execute(
                f"DELETE FROM {table} WHERE {col} IN "
                f"(SELECT id FROM conversations WHERE created_at < ?)",
                (cutoff,))
            deleted[table] = cur.rowcount
        await conn.commit()

        # ---- 对账: 导出行数 vs 删除行数(删除已提交, 差一行必须人工响) ------
        mismatched = []
        for table in DELETE_ORDER:
            exported = (len(doc[table]) if table == "conversations"
                        else sum(len(rows) for rows in doc[table].values()))
            if exported != deleted[table]:
                mismatched.append(f"{table}(导出={exported} "
                                  f"删除={deleted[table]})")
        if mismatched:
            raise SystemExit("导出/删除对账不一致: "
                             + ", ".join(mismatched)
                             + "; 导出文件已落盘, 请人工核查")
        print(f"导出: {export_path}")
        print(f"归档会话: {len(old_ids)} 条; 各表删除行数(倒 FK 序):")
        for table in DELETE_ORDER:
            print(f"  {table}: {deleted[table]}")
        print("导出/删除对账: 一致")

        # ---- VACUUM(删了行才跑; 已 commit, 事务外; 占用锁不判归档失败) ----
        try:
            pages_before = await _scalar(conn, "PRAGMA page_count")
            await conn.execute("VACUUM")
            pages_after = await _scalar(conn, "PRAGMA page_count")
            print(f"VACUUM: 已执行(page_count {pages_before} "
                  f"→ {pages_after})")
            print("归档结论: 导出已落盘, 删除已提交, VACUUM 完成")
        except sqlite3.OperationalError as exc:
            print(f"VACUUM: 被占用锁跳过({exc})——归档已完成, "
                  f"请稍后手动执行 VACUUM")
            print("归档结论: 导出已落盘, 删除已提交")

        # ---- 报告收尾: 库文件前后大小 --------------------------------------
        size_after = tenant_db.stat().st_size
        print(f"库文件: {size_before} → {size_after} 字节")


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--tenant", required=True,
                        help="租户 slug(须已 bootstrap 注册)")
    parser.add_argument("--days", type=int, default=365,
                        help="归档阈值: created_at 早于 N 天前的会话"
                             "(默认 365)")
    parser.add_argument("--dry-run", action="store_true",
                        help="只列不动: 不导出/不删除/不 VACUUM, 库零变化")
    parser.add_argument("--platform-db",
                        default=str(REPO_ROOT / "data" / "platform.db"),
                        help="platform.db 路径(默认 <repo>/data/, 门禁查"
                             "租户是否已 bootstrap)")
    parser.add_argument("--tenants-dir",
                        default=str(REPO_ROOT / "tenants"),
                        help="租户库目录(默认 <repo>/tenants; 导出落其"
                             "兄弟目录 data/archives/<tenant>/)")
    args = parser.parse_args()
    asyncio.run(run(args))


if __name__ == "__main__":
    main()