Files
FastapiAdmin/backend/app/scripts/initialize.py
T

135 lines
4.9 KiB
Python

# -*- coding: utf-8 -*-
import uuid
import json
from pathlib import Path
from typing import Dict, List
from sqlalchemy import inspect, select, func
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.base_model import MappedBase
from app.core.logger import logger
from app.config.setting import settings
from app.api.v1.module_system.user.model import UserModel, UserRolesModel, UserPositionsModel
from app.api.v1.module_system.role.model import RoleModel, RoleDeptsModel, RoleMenusModel
from app.api.v1.module_system.position.model import PositionModel
from app.api.v1.module_system.dept.model import DeptModel
from app.api.v1.module_system.menu.model import MenuModel
from app.api.v1.module_system.log.model import OperationLogModel
from app.api.v1.module_system.notice.model import NoticeModel
from app.api.v1.module_system.config.model import ConfigModel
from app.api.v1.module_system.dict.model import DictTypeModel, DictDataModel
from app.api.v1.module_monitor.job.model import JobModel, JobLogModel
from app.api.v1.module_example.demo.model import DemoModel
class InitializeData:
"""
初始化数据库和基础数据
"""
def __init__(self) -> None:
# 按照依赖关系排序:先创建基础表,再创建关联表
self.prepare_init_models = [
# 部门表(自引用,需要先创建)
DeptModel,
# 菜单表(自引用,需要先创建)
MenuModel,
# 用户表(依赖部门和角色)
UserModel,
# 角色表(依赖菜单和部门)
RoleModel,
# 岗位表(无外键依赖)
PositionModel,
# 基础表(无外键依赖)
ConfigModel,
DictTypeModel,
NoticeModel,
OperationLogModel,
DictDataModel,
JobModel,
JobLogModel,
DemoModel,
# 关联表(依赖基础表)
UserPositionsModel,
UserRolesModel,
RoleDeptsModel,
RoleMenusModel,
]
self.created_tables = set()
async def __get_existing_tables(self, db: AsyncSession) -> List[str]:
return await db.run_sync(
lambda sync_db: inspect(sync_db.get_bind()).get_table_names()
)
async def __init_model(self, db: AsyncSession) -> None:
"""初始化数据库表结构"""
try:
# 获取所有模型元数据
metadata = MappedBase.metadata
# 只创建不存在的表
for table in metadata.sorted_tables:
if table.name not in await self.__get_existing_tables(db):
await db.run_sync(lambda sync_db: table.create(sync_db.bind))
self.created_tables.add(table.name)
logger.info(f"已创建表: {table.name}")
await self.__init_data(db)
except Exception as e:
logger.error(f"初始化数据库结构失败: {str(e)}")
raise
async def __init_data(self, db: AsyncSession) -> None:
"""初始化基础数据"""
for model in self.prepare_init_models:
table_name = model.__tablename__
# 检查表中是否已经有数据
count_result = await db.execute(select(func.count()).select_from(model))
existing_count = count_result.scalar()
if existing_count > 0:
logger.warning(f"跳过 {table_name} 表数据初始化(表已存在 {existing_count} 条记录)")
continue
data = await self.__get_data(table_name)
if not data:
logger.warning(f"跳过 {table_name} 表,无初始化数据")
continue
try:
# 表为空,直接插入全部数据
objs = [model(**item) for item in data]
db.add_all(objs)
await db.flush()
logger.info(f"已向 {table_name} 表写入 {len(objs)} 条记录")
except Exception as e:
logger.error(f"初始化 {table_name} 表数据失败: {str(e)}")
raise
async def __get_data(self, filename: str) -> List[Dict]:
"""读取初始化数据文件"""
json_path = Path.joinpath(settings.SCRIPT_DIR, f'{filename}.json')
if not json_path.exists():
return []
try:
with open(json_path, 'r', encoding='utf-8') as f:
return json.loads(f.read())
except json.JSONDecodeError as e:
logger.error(f"解析 {json_path} 失败: {str(e)}")
raise
except Exception as e:
logger.error(f"读取 {json_path} 失败: {str(e)}")
raise
async def init_db(self, db: AsyncSession) -> None:
"""
执行完整初始化流程
"""
await self.__init_model(db)