system-backu/backend/app/core/backup/directory.py

"""目录备份：使用 tarfile 流式打包为 tar.gz。"""
import asyncio
import fnmatch
import os
from datetime import datetime
from pathlib import Path
from typing import AsyncIterator

from app.core.backup.base import BackupMetadata, BaseBackup
from app.utils.logging import get_logger

log = get_logger(__name__)


class DirectoryBackup(BaseBackup):
    def __init__(self, source_config: dict, job_name: str):
        super().__init__(source_config, job_name)
        self.path: str = source_config["path"]
        self.exclude_patterns: list[str] = source_config.get("exclude_patterns", []) or []
        self.base_dir = os.path.basename(os.path.normpath(self.path)) or "root"
        self.root_path = Path(self.path).resolve()

    def metadata(self) -> BackupMetadata:
        ts = datetime.now().strftime("%Y%m%d-%H%M%S")
        safe = "".join(c if c.isalnum() or c in ("_", "-") else "_" for c in self.base_dir)
        return BackupMetadata(
            suggested_filename=f"{self.job_name}-{safe}-{ts}.tar.gz",
            content_type="application/gzip",
            extra={"path": self.path, "base_dir": self.base_dir},
        )

    def _is_excluded(self, rel_path: str) -> bool:
        name = os.path.basename(rel_path)
        for pat in self.exclude_patterns:
            if fnmatch.fnmatch(name, pat) or fnmatch.fnmatch(rel_path, pat):
                return True
        return False

    async def produce(self) -> AsyncIterator[bytes]:
        if not self.root_path.exists():
            raise FileNotFoundError(f"Source path does not exist: {self.path}")
        if not self.root_path.is_dir():
            raise NotADirectoryError(f"Source path is not a directory: {self.path}")

        log.info("Taring directory %s (exclude=%s)", self.path, self.exclude_patterns)

        # 用一个 Queue 在同步线程和异步消费者之间桥接 tar 产出
        queue: asyncio.Queue[bytes | None] = asyncio.Queue(maxsize=32)

        def _tar_worker() -> None:
            import tarfile

            try:
                with tarfile.open(mode="w:gz", bufsize=1024 * 1024) as tar:
                    # 顶层目录名固定为 base_dir，便于解压后还原
                    for dirpath, dirnames, filenames in os.walk(self.root_path):
                        # 排除目录
                        dirnames[:] = [
                            d for d in dirnames
                            if not self._is_excluded(os.path.relpath(os.path.join(dirpath, d), self.root_path))
                        ]
                        for fn in filenames:
                            abs_path = os.path.join(dirpath, fn)
                            rel = os.path.relpath(abs_path, self.root_path)
                            if self._is_excluded(rel):
                                continue
                            try:
                                tar.add(abs_path, arcname=os.path.join(self.base_dir, rel), recursive=False)
                            except (PermissionError, OSError) as e:
                                log.warning("Skip %s: %s", abs_path, e)
                                continue
                            # 周期性 flush 到 queue
                            if tar.fileobj.tell() and (tar.fileobj.tell() % (256 * 1024) == 0):
                                # 这里我们用 tarfile 的 fileobj 直接喂到 queue 更高效
                                pass
            except Exception as e:
                log.exception("tar worker error: %s", e)
            finally:
                # tarfile 关闭后会写完 gzip trailer，我们需要在关闭时把 fileobj 全部送出
                pass

        # 实现简化：使用 tarfile 的 fileobj=BytesIO 会在循环中累积，
        # 对超大目录会爆内存。下面用一个临时文件 + 异步读取的方案。
        import tempfile

        tmp_path = Path(tempfile.gettempdir()) / f"backup_{self.job_name}_{os.getpid()}.tar.gz"

        def _tar_to_file() -> None:
            import tarfile

            try:
                with tarfile.open(str(tmp_path), mode="w:gz") as tar:
                    for dirpath, dirnames, filenames in os.walk(self.root_path):
                        dirnames[:] = [
                            d for d in dirnames
                            if not self._is_excluded(os.path.relpath(os.path.join(dirpath, d), self.root_path))
                        ]
                        for fn in filenames:
                            abs_path = os.path.join(dirpath, fn)
                            rel = os.path.relpath(abs_path, self.root_path)
                            if self._is_excluded(rel):
                                continue
                            try:
                                tar.add(abs_path, arcname=os.path.join(self.base_dir, rel), recursive=False)
                            except (PermissionError, OSError) as e:
                                log.warning("Skip %s: %s", abs_path, e)
            except Exception:
                log.exception("tar file worker failed")

        loop = asyncio.get_event_loop()
        await loop.run_in_executor(None, _tar_to_file)

        if not tmp_path.exists() or tmp_path.stat().st_size == 0:
            raise RuntimeError(f"Failed to create tar.gz for {self.path}")

        try:
            with open(tmp_path, "rb") as f:
                while True:
                    chunk = await loop.run_in_executor(None, f.read, 64 * 1024)
                    if not chunk:
                        break
                    yield chunk
        finally:
            try:
                tmp_path.unlink(missing_ok=True)
            except Exception:
                pass