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 restore_postgres(backup_file, target_dbname, runner=None,
                     psql_cmd="psql", allow_production=False):
    """Postgres 恢复：用 psql 把 .sql 备份恢复到目标库，返回目标库名。

    安全规则（恢复演练的核心，绝不放宽）：
    - 目标库名绝不能等于生产库名（CLOUD_DB_NAME 环境变量读到的值），
      相等时直接报错拒绝；只有调用方显式传 allow_production=True
      才放行，而部署脚本（deploy/restore-drill.sh）永远不设这个值。
    - 恢复前要求备份文件存在且非空，否则报错，不跑 psql。
    - 连接参数与 backup_postgres 同款：host/port/user 走 CLOUD_DB_*
      环境变量，密码只以 PGPASSWORD 环境变量传给子进程，绝不
      出现在命令行参数、返回值或日志里。
    - 目标库必须已存在（调用方先 CREATE DATABASE，见恢复演练
      脚本）；psql 非零退出时明确报错（stderr 截断，不含密码）。
    - 本函数只恢复，不删库、不改生产库；删测试库是调用方的事。
    """
    if not backup_file or not os.path.isfile(backup_file):
        raise FileNotFoundError("备份文件不存在：%s" % backup_file)
    if os.path.getsize(backup_file) == 0:
        raise ValueError("备份文件为空，拒绝恢复：%s" % backup_file)
    if not target_dbname or not isinstance(target_dbname, str) \
            or not target_dbname.strip() \
            or any(c.isspace() for c in target_dbname):
        raise ValueError("目标库名非法：%r" % (target_dbname,))
    production = os.environ.get("CLOUD_DB_NAME") or ""
    if production and target_dbname == production \
            and not allow_production:
        raise ValueError(
            "拒绝恢复到生产库 %s：目标库名与 CLOUD_DB_NAME 相同；"
            "恢复演练必须用测试库名" % production)
    host = os.environ.get("CLOUD_DB_HOST") or "127.0.0.1"
    port = os.environ.get("CLOUD_DB_PORT") or "5432"
    user = os.environ.get("CLOUD_DB_USER")
    if not user:
        raise RuntimeError("Postgres 恢复需要 CLOUD_DB_USER")
    if shutil.which(psql_cmd) is None and runner is None:
        raise RuntimeError("psql 不可用：本机找不到 %s" % psql_cmd)
    cmd = [psql_cmd, "-h", host, "-p", str(port), "-U", user,
           "-d", target_dbname, "-v", "ON_ERROR_STOP=1",
           "-f", backup_file]
    env = dict(os.environ)
    password = os.environ.get("CLOUD_DB_PASSWORD")
    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("psql 恢复失败：%s" % err)
    return target_dbname


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()