170 lines
6.8 KiB
Python
170 lines
6.8 KiB
Python
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
|