codex-app-server/cloud/backup.py

"""云库备份：SQLite 在线备份 + 轮转。

用法：python3 -m cloud.backup --db <库路径> --out <目录> [--keep 3]

- 用 sqlite3 的 backup API 做在线备份（源库正在被 server 读写时
  也能安全拷，不用先停服务、不用直接 cp 正在写的文件）。
- 轮转：备份文件名含时间戳，写完后按时间倒序只保留最近 keep
  份（默认 3），多余的删掉；keep < 1 时报错，不删光。
- 备份失败（源库不存在、目标目录建不出来、backup API 报错）
  明确抛错、不静默，CLI 退出码非零。
- 备份文件本身就是 SQLite 库，可直接用 sqlite3 打开验证含
  tasks 表；备份里不含任何密钥（库里本来就只存 API key 哈希）。
"""

import argparse
import glob
import os
import shutil
import sqlite3
import subprocess
import time

from cloud.db import connect as _db_connect


def backup_db(db_path, out_dir, keep=3, timestamp=None):
    """执行一次备份并轮转，返回备份文件路径。

    timestamp 可注入（单测用固定时间戳造多份）；缺省用当前时间。
    """
    if keep < 1:
        raise ValueError("keep 必须 >= 1，否则会删光所有备份")
    if not os.path.exists(db_path):
        raise FileNotFoundError("源库不存在：%s" % db_path)
    os.makedirs(out_dir, exist_ok=True)
    stamp = timestamp or time.strftime("%Y%m%d%H%M%S", time.localtime())
    dest = os.path.join(out_dir, "cloud-backup-%s.db" % stamp)
    src = _db_connect(str(db_path))
    try:
        dst = sqlite3.connect(dest)
        try:
            src.backup(dst)
        finally:
            dst.close()
    finally:
        src.close()
    # 轮转：按文件名时间戳倒序（文件名即时间戳，排序即时间序），
    # 只留最近 keep 份。
    backups = sorted(glob.glob(os.path.join(out_dir, "cloud-backup-*.db")))
    for old in backups[:-keep] if len(backups) > keep else []:
        os.remove(old)
    return dest


def backup_postgres(out_dir, keep=3, timestamp=None, runner=None,
                    pg_dump_cmd="pg_dump"):
    """Postgres 备份：调 pg_dump 写出 .sql 文件并轮转，返回文件路径。

    连接参数只走 CLOUD_DB_* 环境变量（host/port/dbname/user/
    password），密码只以 PGPASSWORD 环境变量传给子进程，绝不
    出现在命令行参数、返回值或日志里。runner 可注入（单测用假
    runner）：调用约定 runner(cmd, env) -> 具 returncode/stdout/
    stderr 的对象；未注入时用 subprocess.run。pg_dump 不存在时
    明确报错；pg_dump 非零退出时明确报错（stderr 截断，不含密码）。
    轮转规则与 SQLite 分支一致：只留最近 keep 份，keep < 1 报错。
    SQLite 分支原样不动，本函数不碰任何 .db 文件。
    """
    if keep < 1:
        raise ValueError("keep 必须 >= 1，否则会删光所有备份")
    host = os.environ.get("CLOUD_DB_HOST") or "127.0.0.1"
    port = os.environ.get("CLOUD_DB_PORT") or "5432"
    dbname = os.environ.get("CLOUD_DB_NAME")
    user = os.environ.get("CLOUD_DB_USER")
    password = os.environ.get("CLOUD_DB_PASSWORD")
    if not dbname or not user:
        raise RuntimeError(
            "Postgres 备份需要 CLOUD_DB_NAME 与 CLOUD_DB_USER")
    if shutil.which(pg_dump_cmd) is None and runner is None:
        raise RuntimeError("pg_dump 不可用：本机找不到 %s" % pg_dump_cmd)
    os.makedirs(out_dir, exist_ok=True)
    stamp = timestamp or time.strftime("%Y%m%d%H%M%S", time.localtime())
    dest = os.path.join(out_dir, "cloud-backup-%s.sql" % stamp)
    cmd = [pg_dump_cmd, "-h", host, "-p", str(port), "-U", user,
           "-d", dbname, "-f", dest]
    env = dict(os.environ)
    if password:
        env["PGPASSWORD"] = password
    if runner is not None:
        result = runner(cmd, env)
    else:
        result = subprocess.run(cmd, capture_output=True, text=True,
                                env=env)
    if getattr(result, "returncode", 1) != 0:
        err = (getattr(result, "stderr", "") or "")[:500]
        raise RuntimeError("pg_dump 失败：%s" % err)
    if not os.path.exists(dest):
        raise RuntimeError("pg_dump 未生成备份文件：%s" % dest)
    backups = sorted(glob.glob(os.path.join(out_dir, "cloud-backup-*.sql")))
    for old in backups[:-keep] if len(backups) > keep else []:
        os.remove(old)
    return dest


def main(argv=None):
    parser = argparse.ArgumentParser(description="cloud control-plane backup")
    parser.add_argument("--db", default=None,
                        help="SQLite 库路径；postgres 模式可省略")
    parser.add_argument("--out", required=True)
    parser.add_argument("--keep", type=int, default=3)
    parser.add_argument("--backend", default="sqlite",
                        choices=["sqlite", "postgres"],
                        help="sqlite 时 --db 必填；postgres 时忽略 --db，"
                             "连接参数走 CLOUD_DB_* 环境变量")
    args = parser.parse_args(argv)
    if args.backend == "postgres":
        dest = backup_postgres(args.out, keep=args.keep)
    else:
        if not args.db:
            parser.error("sqlite 模式需要 --db")
        dest = backup_db(args.db, args.out, keep=args.keep)
    print(dest)


if __name__ == "__main__":
    main()