api-copp/route/init_data_base.py

from fastapi import APIRouter, Depends, HTTPException
import pkgutil
import importlib
from pathlib import Path
from typing import List

from bd import Settings, get_engine
from sqlalchemy import text
from sqlalchemy.exc import SQLAlchemyError
from pydantic import BaseModel

from route.auth_utils import require_admin_key

router = APIRouter(tags=["db"])


def _iter_table_modules() -> List[str]:
    try:
        import bd.tables as tables_pkg
        pkg_paths = getattr(tables_pkg, "__path__", None)
        if not pkg_paths:
            pkg_paths = [str(Path(__file__).resolve().parent.parent / "bd" / "tables")]
    except Exception:
        pkg_paths = [str(Path(__file__).resolve().parent.parent / "bd" / "tables")]

    names = []
    for finder, name, ispkg in pkgutil.iter_modules(pkg_paths):
        names.append(name)
    return names


def collect_metadatas():
    metadatas = []
    for mod_name in _iter_table_modules():
        try:
            module = importlib.import_module(f"bd.tables.{mod_name}")
        except Exception:
            continue
        Base = getattr(module, "Base", None)
        if Base is not None and hasattr(Base, "metadata"):
            metadatas.append(Base.metadata)
    return metadatas


@router.post("/db/create-tables", dependencies=[Depends(require_admin_key)])
def create_tables():
    """Создаёт все таблицы, описанные в модулях `db.tables`.

    Endpoint вызывается по нажатию кнопки в UI (POST).
    """
    engine = get_engine()

    metadatas = collect_metadatas()
    if not metadatas:
        raise HTTPException(status_code=400, detail="No table metadata found in db.tables")

    # deduplicate metadata objects (multiple modules may expose the same Base.metadata)
    unique = []
    seen = set()
    for md in metadatas:
        if id(md) not in seen:
            seen.add(id(md))
            unique.append(md)

    try:
        for md in unique:
            md.create_all(bind=engine)
        return {"status": "ok", "detail": f"Created {len(unique)} metadata groups"}
    except SQLAlchemyError as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/db/migrate-cascade-user", dependencies=[Depends(require_admin_key)])
def migrate_cascade_user():
    """Добавляет FK responses.user_id → users.id ON DELETE CASCADE, если ещё не существует."""
    engine = get_engine()
    try:
        with engine.connect() as conn:
            conn.execute(text("""
                DO $$ BEGIN
                    IF NOT EXISTS (
                        SELECT 1 FROM pg_constraint
                        WHERE conname = 'responses_user_id_fkey'
                    ) THEN
                        ALTER TABLE responses
                        ADD CONSTRAINT responses_user_id_fkey
                        FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE;
                    END IF;
                END $$;
            """))
            conn.commit()
        return {"status": "ok", "detail": "FK responses.user_id → users(id) ON DELETE CASCADE ensured"}
    except SQLAlchemyError as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/db/migrate-users", dependencies=[Depends(require_admin_key)])
def migrate_users():
    """Добавляет колонки username и hashed_password в таблицу users, если они ещё не существуют."""
    engine = get_engine()
    try:
        with engine.connect() as conn:
            conn.execute(text("""
                ALTER TABLE users
                ADD COLUMN IF NOT EXISTS username VARCHAR(150),
                ADD COLUMN IF NOT EXISTS hashed_password VARCHAR(256)
            """))
            conn.commit()
        return {"status": "ok", "detail": "Columns username and hashed_password ensured in users table"}
    except SQLAlchemyError as e:
        raise HTTPException(status_code=500, detail=str(e))


class ClearDBIn(BaseModel):
    confirm: bool


@router.post("/db/migrate", dependencies=[Depends(require_admin_key)])
def migrate_tables():
    """Приводит схему БД в соответствие с моделями: убирает устаревшие колонки, добавляет новые."""
    engine = get_engine()
    migrations = [
        # Убираем старый FK и колонку group.user_id (если остался от прежней схемы)
        'ALTER TABLE "group" DROP COLUMN IF EXISTS user_id',
        # Делаем organization.group_id необязательным
        "ALTER TABLE organization ALTER COLUMN group_id DROP NOT NULL",
        # Добавляем новые колонки в users
        'ALTER TABLE users ADD COLUMN IF NOT EXISTS group_id UUID REFERENCES "group"(id) ON DELETE SET NULL',
        "ALTER TABLE users ADD COLUMN IF NOT EXISTS organization_id UUID REFERENCES organization(id) ON DELETE SET NULL",
        # Добавляем путь к SVG-файлу диаграммы в radar_results
        "ALTER TABLE radar_results ADD COLUMN IF NOT EXISTS image_path VARCHAR(512)",
    ]
    applied = []
    try:
        with engine.begin() as conn:
            for sql in migrations:
                conn.execute(text(sql))
                applied.append(sql)
        return {"status": "ok", "applied": applied}
    except SQLAlchemyError as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/db/clear", dependencies=[Depends(require_admin_key)])
def clear_tables(payload: ClearDBIn):
    """Полная очистка всех таблиц, описанных в `bd.tables`.

    Требуется явное подтверждение: POST с телом {"confirm": true}.
    """
    if not payload.confirm:
        raise HTTPException(status_code=400, detail="Confirmation required")

    engine = get_engine()

    metadatas = collect_metadatas()
    if not metadatas:
        raise HTTPException(status_code=400, detail="No table metadata found in db.tables")

    # deduplicate metadata objects
    unique = []
    seen = set()
    for md in metadatas:
        if id(md) not in seen:
            seen.add(id(md))
            unique.append(md)

    try:
        with engine.begin() as conn:
            conn.execute(text("DROP SCHEMA public CASCADE"))
            conn.execute(text("CREATE SCHEMA public"))
        return {"status": "ok", "detail": "All tables dropped via DROP SCHEMA CASCADE"}
    except SQLAlchemyError as e:
        raise HTTPException(status_code=500, detail=str(e))