dpb/backend/app/scripts/initialize.py

import json
import secrets
from typing import Any

from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.api.v1.module_system.dept.model import DeptModel
from app.api.v1.module_system.dict.model import DictDataModel, DictTypeModel
from app.api.v1.module_system.menu.model import MenuModel
from app.api.v1.module_system.params.model import ParamsModel
from app.api.v1.module_system.role.model import RoleModel
from app.api.v1.module_bre.germplasm.model import BreedingGermplasmModel
from app.api.v1.module_system.user.model import UserModel, UserRolesModel
from app.api.v1.module_system.versions.model import VersionModel
from app.common.enums import EnvironmentEnum
from app.config.path_conf import LOG_DIR, SCRIPT_DIR
from app.config.setting import settings
from app.core.database import async_db_session, check_db, create_tables
from app.core.logger import logger
from app.utils.password_util import PwdUtil


class InitializeData:
    """初始化数据库和基础数据"""

    # 按依赖关系排序：先基础表，再关联表
    prepare_init_models: list[type] = [
        MenuModel,
        DeptModel,
        ParamsModel,
        RoleModel,
        DictTypeModel,
        DictDataModel,
        UserModel,
        UserRolesModel,
        VersionModel,
        BreedingGermplasmModel,
    ]

    # 树形模型：JSON 含嵌套 children，需递归创建对象
    _RECURSIVE_TABLES: set[str] = {"sys_menu", "sys_dept"}

    async def init_db(self) -> None:
        """建表并导入种子数据"""
        await check_db()
        # await drop_tables()
        await create_tables()

        async with async_db_session() as session, session.begin():
            await self.__init_data(session)

    async def __init_data(self, db: AsyncSession) -> None:
        """按依赖顺序初始化各表种子数据"""
        dict_type_mapping: dict[str, Any] = {}

        for model in self.prepare_init_models:
            table_name = model.__tablename__

            data = await self.__load_json(table_name)
            if not data:
                logger.info(f"⏭️  跳过 {table_name} 表，无初始化数据")
                continue

            # 已有数据则跳过
            count = await db.execute(select(func.count()).select_from(model))
            if count.scalar():
                logger.info(f"⏭️  跳过 {table_name} 表数据初始化（表已有数据）")
                continue

            try:
                if table_name in self._RECURSIVE_TABLES:
                    objs = self.__create_objects_with_children(data, model)
                elif table_name == "sys_dict_type":
                    objs = []
                    for item in data:
                        obj = model(**item)
                        objs.append(obj)
                        dict_type_mapping[item["dict_type"]] = obj
                elif table_name == "sys_dict_data":
                    objs = []
                    for item in data:
                        dict_type_str = item.get("dict_type")
                        if dict_type_str not in dict_type_mapping:
                            logger.warning(f"⚠️  未找到字典类型 {dict_type_str}，跳过")
                            continue
                        item["dict_type_id"] = dict_type_mapping[dict_type_str].id
                        objs.append(model(**item))
                else:
                    # 生产首次初始化：把仓库内置的固定种子密码（如 123456）随机化，防止已知密码登录
                    if settings.ENVIRONMENT == EnvironmentEnum.PROD and table_name == "sys_user":
                        self._randomize_seed_passwords(data)
                    objs = [model(**item) for item in data]

                if objs:
                    db.add_all(objs)
                    await db.flush()
                    logger.info(f"✅️ 已向 {table_name} 写入初始化数据")
                else:
                    logger.info(f"⏭️  跳过 {table_name} 表数据初始化（无有效数据）")

            except Exception:
                logger.error(f"❌️ 初始化 {table_name} 表数据失败")
                raise

    def _randomize_seed_passwords(self, data: list[dict]) -> None:
        """生产环境首次初始化时，把种子账号的固定弱密码替换为随机值。

        仓库内 sys_user.json 携带对所有部署相同的初始密码（如 123456），
        直接使用会让任何未改密的实例可被已知密码直接登录。生产初始化时随机化，
        一次性初始密码写入日志目录文件，供运维首次登录使用；登录修改后务必删除该文件。
        """
        creds: list[str] = []
        for item in data:
            if not item.get("password"):
                continue
            raw = secrets.token_urlsafe(18)
            item["password"] = PwdUtil.hash_password(raw)
            creds.append(f"{item.get('username')}={raw}")
        if not creds:
            return
        try:
            LOG_DIR.mkdir(parents=True, exist_ok=True)
            note = LOG_DIR / "prod_initial_passwords.txt"
            note.write_text(
                "生产环境首次初始化种子账号随机初始密码（登录后请修改并删除本文件）：\n"
                + "\n".join(creds)
                + "\n",
                encoding="utf-8",
            )
            logger.warning(
                "⚠️ 生产初始化已为种子账号生成随机初始密码，已写入 {}，请立即登录修改：{}",
                note,
                "；".join(creds),
            )
        except Exception as e:
            logger.error("生成种子初始密码失败: {}", e)

    @staticmethod
    def __create_objects_with_children(data: list[dict], model_class: type) -> list:
        """递归创建树形模型实例，处理嵌套 children 并注入 parent_id"""

        def _create(obj_data: dict) -> Any:
            children_data = obj_data.pop("children", [])
            obj = model_class(**obj_data)

            # 子节点通过 relationship 自动设置 parent_id
            if children_data:
                obj.children = [_create(child) for child in children_data]

            return obj

        return [_create(item) for item in data]

    async def __load_json(self, filename: str) -> list[dict]:
        """读取并解析种子数据 JSON 文件"""
        json_path = SCRIPT_DIR / f"{filename}.json"
        if not json_path.exists():
            return []

        try:
            with open(json_path, encoding="utf-8") as f:
                return json.load(f)
        except json.JSONDecodeError as e:
            logger.error(f"❌️ 解析 {json_path} 失败: {e!s}")
            raise
        except Exception as e:
            logger.error(f"❌️ 读取 {json_path} 失败: {e!s}")
            raise