Aislo/common_util/common_util_save_guard.py

"""[저장] · [확정] 전부 아니면 없음(PLAN 472-6 · 브레인 A) — 떠 두고 실패면 되돌림.

차례마다 안쪽 창구가 저마다 커넥션 · 파일을 쓰므로 DB 트랜잭션 하나로 못 꿴다. 대신 시작 전에
  - DB 행 — `projects` 한 줄 · `project_workflow_stages` · `routes` 와 자식 넷(`route_points` ·
    `route_statistics` · `longitudinal_sections` · `cross_sections`)
  - storage — 저장 차례가 쓰는 폴더(`STORAGE_SCOPE` · 관 차례가 있으면 배수 폴더도)
를 떠 두고, 하나라도 실패하면 둘 다 되돌린다. 되돌림 쓰기는 바뀐 것만(시놀로지 동기 — 안 바뀐
파일 · 표는 안 씀). 같은 프로젝트 저장이 겹치면 뒤 요청이 기다림(`project_lock`).
"""

from __future__ import annotations

import asyncio
import shutil
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import aiomysql

#: 저장 차례가 쓰는 storage 자리(프로젝트 뿌리 기준) — 추천 노선 미리 셈(`route_recommend`)은 뒤에서
#: 따로 도는 곳이라 안 넣음.
STORAGE_SCOPE = (
    "B05_Profile/route",
    "B05_Profile/corridor",
    "B06_Section",
    "B07_DesignDetail",
    "_derived",
    "templates",
    "workflow.json",
    "project_manifest.json",
)
#: 관 · 관 옵션 차례 · [확정](배수유역 다시 셈)이 쓰는 자리.
PIPE_SCOPE = ("B04_PreProcess/drainage",)
#: 노선 아래 표(노선 행이 지워지면 함께 지워짐 · CASCADE).
ROUTE_CHILDREN = ("route_points", "route_statistics", "longitudinal_sections", "cross_sections")

_LOCKS: dict[str, asyncio.Lock] = {}


def project_lock(project_id: str) -> asyncio.Lock:
    """같은 프로젝트 저장 한 번에 하나 — 뒤 요청은 앞 저장이 끝날 때까지 기다림.
    ponytail: 프로세스 메모리 잠금 · 서버 여러 대면 DB 잠금으로."""
    return _LOCKS.setdefault(str(project_id), asyncio.Lock())


# ── storage ────────────────────────────────────────────────────────────────


def _stamp(path: Path) -> tuple[int, int]:
    stat = path.stat()
    return stat.st_size, stat.st_mtime_ns


def _files(root: Path, scope: tuple[str, ...]) -> dict[str, Path]:
    out: dict[str, Path] = {}
    for relative in scope:
        target = root / relative
        if target.is_file():
            out[relative] = target
        elif target.is_dir():
            for path in target.rglob("*"):
                if path.is_file():
                    out[path.relative_to(root).as_posix()] = path
    return out


@dataclass
class StorageSnapshot:
    """떠 둔 파일 — 사본 폴더(임시) · 파일마다 (크기, 바뀐 때)."""

    root: Path
    scope: tuple[str, ...]
    copy_dir: Path
    stamps: dict[str, tuple[int, int]] = field(default_factory=dict)

    def restore(self) -> int:
        """바뀐 · 새로 생긴 · 지워진 파일만 되돌림 → 되쓴 수(지움 포함)."""
        written = 0
        now = _files(self.root, self.scope)
        for relative, path in now.items():
            if relative not in self.stamps:
                path.unlink()
                written += 1
            elif _stamp(path) != self.stamps[relative]:
                shutil.copy2(self.copy_dir / relative, path)
                written += 1
        for relative in self.stamps.keys() - now.keys():
            target = self.root / relative
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(self.copy_dir / relative, target)
            written += 1
        return written

    def changed(self) -> list[str]:
        """떠 둔 뒤 바뀐 자리(시험 · 로그)."""
        now = _files(self.root, self.scope)
        moved = [r for r, p in now.items() if self.stamps.get(r) != _stamp(p)]
        return sorted(moved + list(self.stamps.keys() - now.keys()))

    def discard(self) -> None:
        shutil.rmtree(self.copy_dir, ignore_errors=True)


def snapshot_storage(root: Path, scope: tuple[str, ...]) -> StorageSnapshot:
    copy_dir = Path(tempfile.mkdtemp(prefix="aislo_save_"))
    snap = StorageSnapshot(Path(root), scope, copy_dir)
    for relative, path in _files(snap.root, scope).items():
        target = copy_dir / relative
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(path, target)
        snap.stamps[relative] = _stamp(path)
    return snap


# ── DB ─────────────────────────────────────────────────────────────────────


async def _select(cursor: Any, sql: str, args: tuple | list) -> list[dict[str, Any]]:
    await cursor.execute(sql, args)
    return [dict(row) for row in await cursor.fetchall()]


async def snapshot_rows(connection: aiomysql.Connection, project_id: str) -> dict[str, list]:
    """프로젝트 한 줄 · 단계 · 노선과 자식 넷의 행(값 그대로)."""
    async with connection.cursor(aiomysql.DictCursor) as cursor:
        rows = {
            "projects": await _select(cursor, "SELECT * FROM projects WHERE id=%s", (project_id,)),
            "project_workflow_stages": await _select(
                cursor,
                "SELECT * FROM project_workflow_stages WHERE project_id=%s ORDER BY id",
                (project_id,),
            ),
            "routes": await _select(
                cursor, "SELECT * FROM routes WHERE project_id=%s ORDER BY id", (project_id,)
            ),
        }
        ids = [row["id"] for row in rows["routes"]] or [-1]
        marks = ",".join(["%s"] * len(ids))
        for table in ROUTE_CHILDREN:
            rows[table] = await _select(
                cursor, f"SELECT * FROM {table} WHERE route_id IN ({marks}) ORDER BY id", ids
            )
    return rows


async def _insert(cursor: Any, table: str, rows: list[dict[str, Any]]) -> None:
    if not rows:
        return
    columns = list(rows[0])
    names = ",".join(f"`{name}`" for name in columns)
    marks = ",".join(["%s"] * len(columns))
    await cursor.executemany(
        f"INSERT INTO {table} ({names}) VALUES ({marks})",
        [[row[name] for name in columns] for row in rows],
    )


async def restore_rows(
    connection: aiomysql.Connection, project_id: str, saved: dict[str, list]
) -> list[str]:
    """바뀐 표만 떠 둔 행으로 — 한 트랜잭션 → 되돌린 표 이름."""
    now = await snapshot_rows(connection, project_id)
    tree = ("routes", *ROUTE_CHILDREN)
    redo_tree = any(now[table] != saved[table] for table in tree)
    redo_stages = now["project_workflow_stages"] != saved["project_workflow_stages"]
    redo_project = now["projects"] != saved["projects"] and bool(saved["projects"])
    if not (redo_tree or redo_stages or redo_project):
        return []
    await connection.begin()
    try:
        async with connection.cursor() as cursor:
            if redo_tree:
                await cursor.execute("DELETE FROM routes WHERE project_id=%s", (project_id,))
                for table in tree:
                    await _insert(cursor, table, saved[table])
            if redo_stages:
                await cursor.execute(
                    "DELETE FROM project_workflow_stages WHERE project_id=%s", (project_id,)
                )
                await _insert(cursor, "project_workflow_stages", saved["project_workflow_stages"])
            if redo_project:
                row = saved["projects"][0]
                columns = [name for name in row if name != "id"]
                sets = ",".join(f"`{name}`=%s" for name in columns)
                await cursor.execute(
                    f"UPDATE projects SET {sets} WHERE id=%s",
                    [row[name] for name in columns] + [project_id],
                )
        await connection.commit()
    except Exception:
        await connection.rollback()
        raise
    return [
        *(tree if redo_tree else ()),
        *(("project_workflow_stages",) if redo_stages else ()),
        *(("projects",) if redo_project else ()),
    ]