zhifu-platform/scripts/mysql_recovery.py

import argparse
import hashlib
import json
import os
import secrets
import subprocess
import time
from datetime import UTC, datetime
from pathlib import Path
from uuid import uuid4

ROOT = Path(__file__).resolve().parents[1]
BACKUPS = ROOT / ".work" / "backups"
DATABASES = {"mysql": "zhifu", "mysql-oms": "oms"}
DUMP = (
    'MYSQL_PWD="$MYSQL_ROOT_PASSWORD" exec mysqldump -uroot '
    "--single-transaction --quick --hex-blob --order-by-primary "
    "--skip-comments --skip-dump-date --skip-add-drop-table "
    "--set-gtid-purged=OFF --no-tablespaces "
)


def docker(*args, **kwargs):
    return subprocess.run(["docker", *args], check=True, timeout=180, **kwargs)


def digest(path):
    with path.open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def backup():
    directory = BACKUPS / (datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") + "-" + uuid4().hex[:8])
    directory.mkdir(parents=True)
    manifest = {"format": 1, "started_at": datetime.now(UTC).isoformat(), "databases": {}}
    for service, database in DATABASES.items():
        source = docker("compose", "ps", "-q", service, cwd=ROOT, capture_output=True).stdout
        source_id = source.decode().strip()
        if not source_id:
            raise RuntimeError(f"{service}: MySQL 容器未启动")
        metadata = json.loads(docker("inspect", source_id, capture_output=True).stdout)[0]
        path = directory / f"{database}.sql"
        started = datetime.now(UTC).isoformat()
        with path.open("xb") as output:
            docker("exec", source_id, "sh", "-c", DUMP + database, stdout=output)
        manifest["databases"][database] = {
            "file": path.name,
            "sha256": digest(path),
            "bytes": path.stat().st_size,
            "image_id": metadata["Image"],
            "started_at": started,
            "finished_at": datetime.now(UTC).isoformat(),
        }
    manifest["finished_at"] = datetime.now(UTC).isoformat()
    (directory / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    print(directory)
    return directory


def restore(directory):
    directory = directory.resolve(strict=True)
    if not directory.is_relative_to(BACKUPS.resolve()):
        raise ValueError("恢复文件必须位于项目 .work/backups")
    manifest = json.loads((directory / "manifest.json").read_text(encoding="utf-8"))
    if manifest["format"] != 1 or set(manifest["databases"]) != set(DATABASES.values()):
        raise ValueError("备份清单格式不符合要求")
    for database, entry in manifest["databases"].items():
        if entry["file"] != f"{database}.sql":
            raise ValueError("备份文件名称不符合要求")
        path = (directory / entry["file"]).resolve(strict=True)
        if not path.is_relative_to(directory) or digest(path) != entry["sha256"]:
            raise ValueError("备份文件 SHA256 校验失败")
    started = time.monotonic()
    report = {"started_at": datetime.now(UTC).isoformat(), "databases": {}}
    for database, entry in manifest["databases"].items():
        name = "zhifu-restore-" + database + "-" + uuid4().hex[:12]
        environment = {**os.environ, "MYSQL_ROOT_PASSWORD": secrets.token_urlsafe(32)}
        docker(
            "run",
            "-d",
            "--name",
            name,
            "--network",
            "none",
            "--tmpfs",
            "/var/lib/mysql:rw,size=1g",
            "--memory",
            "2g",
            "--label",
            "zhifu.purpose=recovery-verification",
            "--label",
            "zhifu.recovery-run=" + os.environ.get("ZHIFU_RECOVERY_RUN_ID", "standalone"),
            "-e",
            "MYSQL_ROOT_PASSWORD",
            "-e",
            f"MYSQL_DATABASE={database}",
            entry["image_id"],
            env=environment,
            capture_output=True,
        )
        try:
            for _ in range(90):
                ready = subprocess.run(
                    [
                        "docker",
                        "exec",
                        name,
                        "sh",
                        "-c",
                        'MYSQL_PWD="$MYSQL_ROOT_PASSWORD" mysql -h127.0.0.1 -uroot -e "SELECT 1"',
                    ],
                    capture_output=True,
                    timeout=5,
                )
                if ready.returncode == 0:
                    break
                time.sleep(1)
            else:
                raise TimeoutError("独立 MySQL 恢复容器启动超时")
            with (directory / entry["file"]).open("rb") as source:
                docker(
                    "exec",
                    "-i",
                    name,
                    "sh",
                    "-c",
                    'MYSQL_PWD="$MYSQL_ROOT_PASSWORD" exec mysql -uroot ' + database,
                    stdin=source,
                )
            restored = directory / f"{database}.restored.sql"
            with restored.open("wb") as output:
                docker("exec", name, "sh", "-c", DUMP + database, stdout=output)
            if digest(restored) != entry["sha256"]:
                raise RuntimeError(f"{database}: 恢复后的完整 SQL 内容校验失败")
            counts = (
                "'orders', (SELECT COUNT(*) FROM order_projections), "
                "'tickets', (SELECT COUNT(*) FROM tickets), "
                "'tenants', (SELECT COUNT(*) FROM tenants), "
                if database == "zhifu"
                else "'orders', (SELECT COUNT(*) FROM orders), "
            )
            query = (
                "SELECT JSON_OBJECT('revision', (SELECT version_num FROM alembic_version), "
                + counts
                + "'tables', (SELECT COUNT(*) FROM information_schema.tables "
                "WHERE table_schema=DATABASE()))"
            )
            result = docker(
                "exec",
                name,
                "sh",
                "-c",
                'MYSQL_PWD="$MYSQL_ROOT_PASSWORD" exec mysql -uroot -N -B ' + database + ' -e "$1"',
                "query",
                query,
                capture_output=True,
            )
            report["databases"][database] = {
                **json.loads(result.stdout),
                "checksum_equal": True,
                "bytes": entry["bytes"],
            }
        finally:
            docker("rm", "-f", name, capture_output=True)
    report["seconds"] = round(time.monotonic() - started, 2)
    (directory / "restore-report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(json.dumps(report, indent=2))
    return report


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="MySQL 备份与独立容器恢复验证")
    parser.add_argument("action", choices=["backup", "verify", "exercise"])
    parser.add_argument("--directory", type=Path)
    args = parser.parse_args()
    if args.action == "verify":
        if args.directory is None:
            parser.error("verify 需要 --directory")
        restore(args.directory)
    else:
        directory = backup()
        if args.action == "exercise":
            restore(directory)