system-backu/backend/app/core/restore.py

"""恢复执行器：把 backup artifact 还原到目标（MySQL 库 / 目录路径）。"""
import asyncio
import gzip
import io
import os
import shutil
import tarfile
import tempfile
from datetime import datetime
from pathlib import Path
from typing import Optional

from app.core.storage.factory import create_storage
from app.database import SessionLocal
from app.models.job import Job
from app.models.restore import RestoreHistory
from app.models.run import Run, RunLog
from app.utils.logging import get_logger

log = get_logger(__name__)


async def execute_restore(restore_id: int) -> None:
    """执行一次恢复。"""
    db = SessionLocal()
    try:
        rh: RestoreHistory | None = db.get(RestoreHistory, restore_id)
        if not rh:
            return

        rh.status = "running"
        rh.started_at = datetime.utcnow()
        db.commit()

        run: Run | None = db.get(Run, rh.run_id)
        if not run or not run.artifact_path:
            raise RuntimeError("源 run 或 artifact 不存在")
        job: Job | None = db.get(Job, run.job_id)
        if not job:
            raise RuntimeError("源 job 不存在")

        storage = create_storage(job.storage)
        log_text = io.StringIO()

        def logl(msg: str) -> None:
            line = f"[{datetime.utcnow():%H:%M:%S}] {msg}"
            log.info(line)
            log_text.write(line + "\n")

        try:
            logl(f"Restoring run_id={run.id} -> {rh.target} (type={job.type})")

            # 先下载到临时文件
            tmpdir = Path(tempfile.gettempdir()) / f"restore_{rh.id}_{os.getpid()}"
            tmpdir.mkdir(parents=True, exist_ok=True)
            local_tmp = tmpdir / Path(run.artifact_path).name

            logl(f"Downloading {run.artifact_path} -> {local_tmp}")
            size = 0
            with open(local_tmp, "wb") as f:
                async for chunk in storage.download_stream(run.artifact_path):
                    f.write(chunk)
                    size += len(chunk)
            logl(f"Downloaded {size} bytes")

            if job.type == "mysql":
                await _restore_mysql(local_tmp, rh.target, job, logl)
            elif job.type == "directory":
                _restore_directory(local_tmp, rh.target, logl)
            else:
                raise RuntimeError(f"Unsupported job type: {job.type}")

            rh.status = "success"
            rh.finished_at = datetime.utcnow()
            logl("Restore success")
        except Exception as e:
            log.exception("Restore failed")
            logl(f"Restore failed: {type(e).__name__}: {e}")
            rh.status = "failed"
            rh.error_message = f"{type(e).__name__}: {e}"[:2000]
            rh.finished_at = datetime.utcnow()
        finally:
            # 写日志
            content = log_text.getvalue()
            db.add(RunLog(run_id=run.id, content=f"[restore #{rh.id}]\n{content}"))
            db.commit()
            try:
                shutil.rmtree(tmpdir, ignore_errors=True)
            except Exception:
                pass
    finally:
        db.close()


async def _restore_mysql(archive_path: Path, target_db: str, job: Job, logl) -> None:
    """把 sql.gz 通过 gunzip | mysql 还原到 target_db。"""
    from app.utils.crypto import decrypt_dict

    src = decrypt_dict(job.source_config_json)
    host = src["host"]
    port = int(src.get("port", 3306))
    user = src["user"]
    password = src.get("password", "")

    logl(f"Decompressing and piping to mysql db={target_db} host={host}:{port}")
    # 解压
    decompressed = archive_path.with_suffix("")  # .sql
    if archive_path.name.endswith(".gz"):
        with gzip.open(archive_path, "rb") as gz, open(decompressed, "wb") as out:
            shutil.copyfileobj(gz, out)
        sql_file = decompressed
    else:
        sql_file = archive_path

    # 调用 mysql client
    args = ["mysql", "-h", host, "-P", str(port), "-u", user, target_db]
    if password:
        # 在 -u 后插入 -pPWD
        idx = args.index("-u")
        args.insert(idx + 2, f"-p{password}")

    proc = await asyncio.create_subprocess_exec(
        *args,
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    assert proc.stdin is not None
    with open(sql_file, "rb") as f:
        while True:
            chunk = f.read(64 * 1024)
            if not chunk:
                break
            proc.stdin.write(chunk)
            await proc.stdin.drain()
    proc.stdin.close()
    rc = await proc.wait()
    if rc != 0:
        stderr = (await proc.stderr.read()).decode(errors="replace") if proc.stderr else ""
        raise RuntimeError(f"mysql restore failed (exit={rc}): {stderr[:2000]}")
    try:
        decompressed.unlink(missing_ok=True)
    except Exception:
        pass


def _restore_directory(archive_path: Path, target_dir: str, logl) -> None:
    """解压 tar.gz 到 target_dir（必须为空或不存在）。"""
    target = Path(target_dir).resolve()
    if target.exists() and any(target.iterdir()):
        raise RuntimeError(f"目标目录非空: {target}")
    target.mkdir(parents=True, exist_ok=True)

    logl(f"Extracting tar.gz to {target}")
    if archive_path.name.endswith(".gz"):
        mode = "r:gz"
    else:
        mode = "r:"
    with tarfile.open(str(archive_path), mode) as tar:
        # 安全检查：拒绝路径逃逸
        for member in tar.getmembers():
            member_path = (target / member.name).resolve()
            if not str(member_path).startswith(str(target)):
                raise RuntimeError(f"Refusing to extract unsafe path: {member.name}")
        tar.extractall(path=str(target))