179 lines
6.8 KiB
Python
179 lines
6.8 KiB
Python
"""自动备份与恢复演练（交付报告 P0-5）。
|
||
|
||
设计：
|
||
- 备份内容：SQLite 数据库（backup API 在线快照）+ 知识库原文（kb 目录）打包为 zip；
|
||
- 触发时机：服务启动时一次 + 每 6 小时周期一次（守护线程）；
|
||
- 保留策略：默认保留最近 14 份，超限删最旧；
|
||
- 完整性校验：每次备份后验证 zip 可打开且包含数据库文件；
|
||
- 恢复演练：verify_restore 对最新备份做真实恢复测试——在临时目录解压，
|
||
用 SQLite 打开并查询关键表，验证备份可用（不是只看文件存在）；
|
||
- 运营看板展示最近备份时间/份数/最近一次恢复演练结果。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import shutil
|
||
import sqlite3
|
||
import zipfile
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
from .config import DATA_DIR, KB_DIR
|
||
|
||
logger = logging.getLogger("yawen.backup")
|
||
|
||
BACKUP_DIR = DATA_DIR / "backups"
|
||
KEEP_COUNT = 14
|
||
DB_NAME = "workspace.db"
|
||
|
||
|
||
def _db_path() -> Path:
|
||
return DATA_DIR / DB_NAME
|
||
|
||
|
||
def create_backup(tag: str = "auto") -> Path | None:
|
||
"""创建一次备份，返回备份文件路径；数据库不存在时返回 None。"""
|
||
db_path = _db_path()
|
||
if not db_path.exists():
|
||
return None
|
||
BACKUP_DIR.mkdir(parents=True, exist_ok=True)
|
||
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||
target = BACKUP_DIR / f"backup-{tag}-{stamp}.zip"
|
||
tmp = target.with_suffix(".zip.tmp")
|
||
try:
|
||
with zipfile.ZipFile(tmp, "w", zipfile.ZIP_DEFLATED, compresslevel=1) as zf:
|
||
# SQLite 在线备份 API：一致性快照，不锁库不中断业务
|
||
snap_target = tmp.with_suffix(".db.snapshot")
|
||
src = sqlite3.connect(str(db_path))
|
||
dst = sqlite3.connect(str(snap_target))
|
||
try:
|
||
src.backup(dst)
|
||
finally:
|
||
dst.close()
|
||
src.close()
|
||
zf.write(snap_target, DB_NAME)
|
||
snap_target.unlink(missing_ok=True)
|
||
# 知识库原文（可选，体积可能大；跳过 vector 索引可重建）
|
||
if KB_DIR.exists():
|
||
for file in sorted(KB_DIR.rglob("*")):
|
||
if file.is_file() and "chroma" not in file.parts and "__pycache__" not in file.parts:
|
||
zf.write(file, f"kb/{file.relative_to(KB_DIR)}")
|
||
tmp.replace(target)
|
||
except Exception as exc:
|
||
tmp.unlink(missing_ok=True)
|
||
logger.error("备份失败：%s", exc)
|
||
return None
|
||
_prune_old_backups()
|
||
ok, detail = verify_backup(target)
|
||
if not ok:
|
||
logger.error("备份完整性校验失败 %s：%s", target.name, detail)
|
||
return target
|
||
|
||
|
||
def _prune_old_backups() -> int:
|
||
"""保留最近 KEEP_COUNT 份，返回删除数。"""
|
||
files = sorted(BACKUP_DIR.glob("backup-*.zip"), key=lambda p: p.stat().st_mtime)
|
||
removed = 0
|
||
for old in files[:-KEEP_COUNT]:
|
||
old.unlink(missing_ok=True)
|
||
removed += 1
|
||
return removed
|
||
|
||
|
||
def _backup_dir() -> Path:
|
||
"""动态取备份目录（测试会替换 BACKUP_DIR）。"""
|
||
return BACKUP_DIR
|
||
|
||
|
||
def verify_backup(path: Path) -> tuple[bool, str]:
|
||
"""校验备份 zip：能打开、含数据库、SQLite 能读关键表。"""
|
||
try:
|
||
with zipfile.ZipFile(path) as zf:
|
||
names = zf.namelist()
|
||
if DB_NAME not in names:
|
||
return False, "缺少数据库文件"
|
||
bad = zf.testzip()
|
||
if bad is not None:
|
||
return False, f"损坏条目：{bad}"
|
||
db_bytes = zf.read(DB_NAME)
|
||
# 二进制安全打开：deserialize 直接从字节恢复内存库，避免文本解码产生 NUL
|
||
mem = sqlite3.connect(":memory:")
|
||
mem.deserialize(db_bytes)
|
||
try:
|
||
tables = {row[0] for row in mem.execute(
|
||
"SELECT name FROM sqlite_master WHERE type='table'").fetchall()}
|
||
required = {"campaign", "agent", "appuser"}
|
||
missing = required - tables
|
||
if missing:
|
||
return False, f"缺少关键表：{sorted(missing)}"
|
||
campaign_count = mem.execute("SELECT COUNT(*) FROM campaign").fetchone()[0]
|
||
finally:
|
||
mem.close()
|
||
return True, f"campaign {campaign_count} 行"
|
||
except Exception as exc:
|
||
return False, str(exc)[:200]
|
||
|
||
|
||
def verify_restore_latest() -> dict:
|
||
"""恢复演练：对最新备份执行 verify_backup，并额外验证可解压到临时目录。"""
|
||
backups = sorted(BACKUP_DIR.glob("backup-*.zip"), key=lambda p: p.stat().st_mtime)
|
||
if not backups:
|
||
return {"ok": False, "detail": "尚无备份", "backup": "", "checked_at": datetime.now().isoformat()}
|
||
latest = backups[-1]
|
||
ok, detail = verify_backup(latest)
|
||
# 真实解压演练（写入临时目录后清理）
|
||
if ok:
|
||
try:
|
||
tmp = BACKUP_DIR / "_restore_drill"
|
||
shutil.rmtree(tmp, ignore_errors=True)
|
||
with zipfile.ZipFile(latest) as zf:
|
||
zf.extractall(tmp)
|
||
restored = tmp / DB_NAME
|
||
with sqlite3.connect(str(restored)) as conn:
|
||
conn.execute("SELECT COUNT(*) FROM sqlite_master").fetchone()
|
||
shutil.rmtree(tmp, ignore_errors=True)
|
||
detail += "；解压演练通过"
|
||
except Exception as exc:
|
||
ok = False
|
||
detail = f"解压演练失败：{exc}"[:200]
|
||
result = {"ok": ok, "detail": detail, "backup": latest.name,
|
||
"checked_at": datetime.now().isoformat()}
|
||
_write_drill_report(result)
|
||
return result
|
||
|
||
|
||
def _write_drill_report(result: dict) -> None:
|
||
try:
|
||
import json
|
||
(BACKUP_DIR / "restore-drill.json").write_text(
|
||
json.dumps(result, ensure_ascii=False, indent=1))
|
||
except OSError as exc:
|
||
logger.error("恢复演练报告写入失败：%s", exc)
|
||
|
||
|
||
def backup_status() -> dict:
|
||
"""看板用的备份状态摘要。"""
|
||
files = sorted(BACKUP_DIR.glob("backup-*.zip"), key=lambda p: p.stat().st_mtime)
|
||
latest = files[-1] if files else None
|
||
drill_path = BACKUP_DIR / "restore-drill.json"
|
||
drill = None
|
||
if drill_path.exists():
|
||
try:
|
||
import json
|
||
drill = json.loads(drill_path.read_text())
|
||
except (OSError, ValueError):
|
||
drill = None
|
||
return {
|
||
"count": len(files),
|
||
"keep_count": KEEP_COUNT,
|
||
"latest": latest.name if latest else "",
|
||
"latest_at": datetime.fromtimestamp(latest.stat().st_mtime).isoformat() if latest else "",
|
||
"latest_size_kb": round(latest.stat().st_size / 1024) if latest else 0,
|
||
"last_drill": drill,
|
||
"directory": str(BACKUP_DIR),
|
||
}
|
||
|
||
|
||
__all__ = ["create_backup", "verify_backup", "verify_restore_latest",
|
||
"backup_status", "_prune_old_backups", "BACKUP_DIR", "KEEP_COUNT"]
|