Aislo/db_management/tools_project_backup.py

"""프로젝트 DB 행 백업 · 되살리기 — 프로젝트를 지우기 전에 뜨고, 잘못되면 그대로 되살린다.

무엇을 뜨나
  `projects` 와 그 아래 연쇄(ON DELETE CASCADE)로 지워지는 표 전부 — TABLES 차례.
  id 는 **그대로** 둔다(되살리면 옛 id · 옛 storage 경로가 다시 맞는다).
  파일(storage 폴더)은 이 도구 밖 — robocopy 로 따로 뜨고 되돌린다(tmp/402_계획.md).

쓰는 법(./venv/Scripts/python.exe db_management/tools_project_backup.py …)
    뜨기       dump <파일.json> [--project <UUID> ...]   (없으면 전부)
    되살리기   restore <파일.json> [--project <UUID> ...]
    읽기 확인  check <파일.json> [--project <UUID> ...]

⚠ 되살리기는 **지워진 뒤에만** — 같은 id 가 살아 있으면 INSERT 가 막혀 통째 되돌린다(한 트랜잭션).
"""

import argparse
import asyncio
import json
import sys
from pathlib import Path

sys.path.append(str(Path(__file__).resolve().parent.parent))

import aiomysql

from config.config_db import DB_HOST, DB_NAME, DB_PASSWORD, DB_PORT, DB_USER

# (표, 프로젝트로 거르는 WHERE) — 부모가 앞. 되살리기는 이 차례로 넣는다.
TABLES = (
    ("projects", "id IN ({p})"),
    ("project_members", "project_id IN ({p})"),
    ("input_files", "project_id IN ({p})"),
    ("processed_point_cloud", "project_id IN ({p})"),
    ("surface_models", "project_id IN ({p})"),
    (
        "terrain_layers",
        "surface_model_id IN (SELECT id FROM surface_models WHERE project_id IN ({p}))",
    ),
    ("routes", "project_id IN ({p})"),
    ("route_points", "route_id IN (SELECT id FROM routes WHERE project_id IN ({p}))"),
    ("route_statistics", "route_id IN (SELECT id FROM routes WHERE project_id IN ({p}))"),
    ("longitudinal_sections", "project_id IN ({p})"),
    ("cross_sections", "project_id IN ({p})"),
    ("structures", "project_id IN ({p})"),
    ("project_workflow_stages", "project_id IN ({p})"),
    ("project_versions", "project_id IN ({p})"),
    ("project_automations", "project_id IN ({p})"),
    ("quantity_items", "project_id IN ({p})"),
    ("change_logs", "project_id IN ({p})"),
    ("outputs", "project_id IN ({p})"),
    ("output_files", "output_id IN (SELECT id FROM outputs WHERE project_id IN ({p}))"),
    ("upload_sessions", "project_id IN ({p})"),
    (
        "upload_chunks",
        "session_id IN (SELECT id FROM upload_sessions WHERE project_id IN ({p}))",
    ),
)


async def _connect():
    return await aiomysql.connect(
        host=DB_HOST,
        port=DB_PORT,
        user=DB_USER,
        password=DB_PASSWORD,
        db=DB_NAME,
        charset="utf8mb4",
    )


# 손자 표 → (부모 표, 부모를 가리키는 칸) — 프로젝트 id 를 부모 행에서 찾는다.
PARENTS = {
    "terrain_layers": ("surface_models", "surface_model_id"),
    "route_points": ("routes", "route_id"),
    "route_statistics": ("routes", "route_id"),
    "output_files": ("outputs", "output_id"),
    "upload_chunks": ("upload_sessions", "session_id"),
}


async def dump(path: Path, projects: list[str]) -> None:
    connection = await _connect()
    try:
        async with connection.cursor(aiomysql.DictCursor) as cursor:
            if not projects:
                await cursor.execute("SELECT id FROM projects ORDER BY name")
                projects = [row["id"] for row in await cursor.fetchall()]
            holders = ", ".join(["%s"] * len(projects))
            data: dict = {"projects": projects, "tables": {}}
            for table, where in TABLES:
                where_sql = where.format(p=holders)
                count = where_sql.count("%s") // len(projects)
                await cursor.execute(
                    f"SELECT * FROM `{table}` WHERE {where_sql}", tuple(projects) * count
                )
                data["tables"][table] = list(await cursor.fetchall())
    finally:
        connection.close()
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, ensure_ascii=False, default=str), encoding="utf-8")
    _report(data)


def _report(data: dict) -> None:
    counts = {table: len(rows) for table, rows in data["tables"].items() if rows}
    print(f"프로젝트 {len(data['projects'])} · 표별 행 {counts}")


def _pick(data: dict, projects: list[str]) -> dict:
    """덤프에서 고른 프로젝트 행만."""
    tables = data["tables"]
    if not projects:
        return tables
    wanted = set(projects)
    owner = {
        parent: {row["id"]: row["project_id"] for row in tables.get(parent, [])}
        for parent, _ in PARENTS.values()
    }

    def project_of(table: str, row: dict) -> str | None:
        if table == "projects":
            return row["id"]
        if table in PARENTS:
            parent, column = PARENTS[table]
            return owner[parent].get(row[column])
        return row["project_id"]

    return {
        table: [row for row in tables.get(table, []) if project_of(table, row) in wanted]
        for table, _ in TABLES
    }


async def restore(path: Path, projects: list[str]) -> None:
    data = json.loads(path.read_text(encoding="utf-8"))
    tables = _pick(data, projects)
    connection = await _connect()
    try:
        await connection.begin()
        async with connection.cursor() as cursor:
            for table, _ in TABLES:
                for row in tables.get(table, []):
                    names = ", ".join(f"`{key}`" for key in row)
                    holders = ", ".join(["%s"] * len(row))
                    await cursor.execute(
                        f"INSERT INTO `{table}` ({names}) VALUES ({holders})",
                        tuple(row.values()),
                    )
        await connection.commit()
    except Exception:
        await connection.rollback()
        raise
    finally:
        connection.close()
    print("되살림:", {table: len(rows) for table, rows in tables.items() if rows})


def main() -> None:
    parser = argparse.ArgumentParser(description="프로젝트 DB 행 백업 · 되살리기")
    parser.add_argument("action", choices=("dump", "restore", "check"))
    parser.add_argument("file", type=Path)
    parser.add_argument("--project", action="append", default=[], help="프로젝트 UUID(여러 번)")
    args = parser.parse_args()
    if args.action == "dump":
        asyncio.run(dump(args.file, args.project))
    elif args.action == "restore":
        asyncio.run(restore(args.file, args.project))
    else:
        data = json.loads(args.file.read_text(encoding="utf-8"))
        _report({"projects": data["projects"], "tables": _pick(data, args.project)})


if __name__ == "__main__":
    main()