"""恢复执行器:把 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))