350 lines
11 KiB
Python
350 lines
11 KiB
Python
import shutil
|
||
from pathlib import Path
|
||
|
||
from aerich import Command
|
||
from fastapi import FastAPI
|
||
from fastapi.middleware import Middleware
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from tortoise.expressions import Q
|
||
|
||
from app.api import api_router
|
||
from app.controllers.api import api_controller
|
||
from app.controllers.user import UserCreate, user_controller
|
||
from app.core.exceptions import (
|
||
DoesNotExist,
|
||
DoesNotExistHandle,
|
||
HTTPException,
|
||
HttpExcHandle,
|
||
IntegrityError,
|
||
IntegrityHandle,
|
||
RequestValidationError,
|
||
RequestValidationHandle,
|
||
ResponseValidationError,
|
||
ResponseValidationHandle,
|
||
)
|
||
from app.log import logger
|
||
from app.models.admin import Api, Menu, Role
|
||
from app.schemas.menus import MenuType
|
||
from app.settings.config import settings
|
||
from app.core.clickhouse import clickhouse_manager
|
||
from app.core.clickhouse_init import ClickHouseInitializer
|
||
from app.services.ingest.remote_push import close_http_client
|
||
|
||
from .middlewares import BackGroundTaskMiddleware
|
||
from .ip_tracking import IpTrackingMiddleware
|
||
|
||
|
||
def make_middlewares():
|
||
middleware = [
|
||
Middleware(
|
||
CORSMiddleware,
|
||
allow_origins=settings.CORS_ORIGINS,
|
||
allow_credentials=settings.CORS_ALLOW_CREDENTIALS,
|
||
allow_methods=settings.CORS_ALLOW_METHODS,
|
||
allow_headers=settings.CORS_ALLOW_HEADERS,
|
||
),
|
||
Middleware(BackGroundTaskMiddleware),
|
||
Middleware(IpTrackingMiddleware),
|
||
]
|
||
return middleware
|
||
|
||
|
||
def register_exceptions(app: FastAPI):
|
||
app.add_exception_handler(DoesNotExist, DoesNotExistHandle)
|
||
app.add_exception_handler(HTTPException, HttpExcHandle)
|
||
app.add_exception_handler(IntegrityError, IntegrityHandle)
|
||
app.add_exception_handler(RequestValidationError, RequestValidationHandle)
|
||
app.add_exception_handler(ResponseValidationError, ResponseValidationHandle)
|
||
|
||
|
||
def register_routers(app: FastAPI, prefix: str = "/api"):
|
||
app.include_router(api_router, prefix=prefix)
|
||
|
||
|
||
async def init_superuser():
|
||
user = await user_controller.model.exists()
|
||
if not user:
|
||
await user_controller.create_user(
|
||
UserCreate(
|
||
username="admin",
|
||
email="admin@admin.com",
|
||
password="123456",
|
||
is_active=True,
|
||
is_superuser=True,
|
||
)
|
||
)
|
||
|
||
|
||
async def init_menus():
|
||
menus = await Menu.exists()
|
||
if not menus:
|
||
parent_menu = await Menu.create(
|
||
menu_type=MenuType.CATALOG,
|
||
name="系统管理",
|
||
path="/system",
|
||
order=1,
|
||
parent_id=0,
|
||
icon="carbon:gui-management",
|
||
is_hidden=False,
|
||
component="Layout",
|
||
keepalive=False,
|
||
redirect="/system/user",
|
||
)
|
||
children_menu = [
|
||
Menu(
|
||
menu_type=MenuType.MENU,
|
||
name="用户管理",
|
||
path="user",
|
||
order=1,
|
||
parent_id=parent_menu.id,
|
||
icon="material-symbols:person-outline-rounded",
|
||
is_hidden=False,
|
||
component="/system/user",
|
||
keepalive=False,
|
||
),
|
||
Menu(
|
||
menu_type=MenuType.MENU,
|
||
name="角色管理",
|
||
path="role",
|
||
order=2,
|
||
parent_id=parent_menu.id,
|
||
icon="carbon:user-role",
|
||
is_hidden=False,
|
||
component="/system/role",
|
||
keepalive=False,
|
||
),
|
||
Menu(
|
||
menu_type=MenuType.MENU,
|
||
name="菜单管理",
|
||
path="menu",
|
||
order=3,
|
||
parent_id=parent_menu.id,
|
||
icon="material-symbols:list-alt-outline",
|
||
is_hidden=False,
|
||
component="/system/menu",
|
||
keepalive=False,
|
||
),
|
||
Menu(
|
||
menu_type=MenuType.MENU,
|
||
name="API管理",
|
||
path="api",
|
||
order=4,
|
||
parent_id=parent_menu.id,
|
||
icon="ant-design:api-outlined",
|
||
is_hidden=False,
|
||
component="/system/api",
|
||
keepalive=False,
|
||
),
|
||
Menu(
|
||
menu_type=MenuType.MENU,
|
||
name="部门管理",
|
||
path="dept",
|
||
order=5,
|
||
parent_id=parent_menu.id,
|
||
icon="mingcute:department-line",
|
||
is_hidden=False,
|
||
component="/system/dept",
|
||
keepalive=False,
|
||
),
|
||
Menu(
|
||
menu_type=MenuType.MENU,
|
||
name="审计日志",
|
||
path="auditlog",
|
||
order=6,
|
||
parent_id=parent_menu.id,
|
||
icon="ph:clipboard-text-bold",
|
||
is_hidden=False,
|
||
component="/system/auditlog",
|
||
keepalive=False,
|
||
),
|
||
]
|
||
await Menu.bulk_create(children_menu)
|
||
|
||
# 创建招聘数据管理菜单
|
||
recruitment_menu = await Menu.create(
|
||
menu_type=MenuType.CATALOG,
|
||
name="招聘数据管理",
|
||
path="/recruitment",
|
||
order=2,
|
||
parent_id=0,
|
||
icon="mdi:briefcase-search",
|
||
is_hidden=False,
|
||
component="Layout",
|
||
keepalive=False,
|
||
redirect="/recruitment/qcwy",
|
||
)
|
||
recruitment_children = [
|
||
Menu(
|
||
menu_type=MenuType.MENU,
|
||
name="前程无忧",
|
||
path="qcwy",
|
||
order=1,
|
||
parent_id=recruitment_menu.id,
|
||
icon="mdi:alpha-q-box",
|
||
is_hidden=False,
|
||
component="/recruitment/qcwy",
|
||
keepalive=True,
|
||
),
|
||
Menu(
|
||
menu_type=MenuType.MENU,
|
||
name="智联招聘",
|
||
path="zhilian",
|
||
order=2,
|
||
parent_id=recruitment_menu.id,
|
||
icon="mdi:alpha-z-box",
|
||
is_hidden=False,
|
||
component="/recruitment/zhilian",
|
||
keepalive=True,
|
||
),
|
||
Menu(
|
||
menu_type=MenuType.MENU,
|
||
name="Boss直聘",
|
||
path="boss",
|
||
order=3,
|
||
parent_id=recruitment_menu.id,
|
||
icon="mdi:alpha-b-box",
|
||
is_hidden=False,
|
||
component="/recruitment/boss",
|
||
keepalive=True,
|
||
),
|
||
]
|
||
await Menu.bulk_create(recruitment_children)
|
||
|
||
# 创建数据清理菜单
|
||
cleaning_menu = await Menu.create(
|
||
menu_type=MenuType.CATALOG,
|
||
name="数据清理",
|
||
path="/cleaning",
|
||
order=3,
|
||
parent_id=0,
|
||
icon="mdi:database-refresh",
|
||
is_hidden=False,
|
||
component="Layout",
|
||
keepalive=False,
|
||
redirect="/cleaning/targeted",
|
||
)
|
||
cleaning_children = [
|
||
Menu(
|
||
menu_type=MenuType.MENU,
|
||
name="定向数据",
|
||
path="targeted",
|
||
order=1,
|
||
parent_id=cleaning_menu.id,
|
||
icon="mdi:filter-target",
|
||
is_hidden=False,
|
||
component="/cleaning/index",
|
||
keepalive=True,
|
||
),
|
||
Menu(
|
||
menu_type=MenuType.MENU,
|
||
name="清洗监控",
|
||
path="monitor",
|
||
order=2,
|
||
parent_id=cleaning_menu.id,
|
||
icon="mdi:monitor-dashboard",
|
||
is_hidden=False,
|
||
component="/cleaning/monitor",
|
||
keepalive=True,
|
||
),
|
||
]
|
||
await Menu.bulk_create(cleaning_children)
|
||
|
||
|
||
async def init_apis():
|
||
apis = await api_controller.model.exists()
|
||
if not apis:
|
||
await api_controller.refresh_api()
|
||
|
||
|
||
async def init_db():
|
||
"""执行数据库迁移（受环境开关与并发保护控制）"""
|
||
command = Command(tortoise_config=settings.TORTOISE_ORM)
|
||
await command.init()
|
||
migration_dir = Path("migrations") / "models"
|
||
if not migration_dir.exists():
|
||
await command.init_db(safe=True)
|
||
return
|
||
try:
|
||
await command.migrate()
|
||
except FileExistsError as e:
|
||
logger.info(f"跳过重复迁移文件生成: {e}")
|
||
except AttributeError:
|
||
logger.warning("unable to retrieve model history from database, model history will be created from scratch")
|
||
shutil.rmtree("migrations")
|
||
await command.init_db(safe=True)
|
||
await command.upgrade(run_in_transaction=True)
|
||
|
||
|
||
async def init_roles():
|
||
roles = await Role.exists()
|
||
if not roles:
|
||
admin_role = await Role.create(
|
||
name="管理员",
|
||
desc="管理员角色",
|
||
)
|
||
user_role = await Role.create(
|
||
name="普通用户",
|
||
desc="普通用户角色",
|
||
)
|
||
|
||
# 分配所有API给管理员角色
|
||
all_apis = await Api.all()
|
||
await admin_role.apis.add(*all_apis)
|
||
# 分配所有菜单给管理员和普通用户
|
||
all_menus = await Menu.all()
|
||
await admin_role.menus.add(*all_menus)
|
||
await user_role.menus.add(*all_menus)
|
||
|
||
# 为普通用户分配基本API
|
||
basic_apis = await Api.filter(Q(method__in=["GET"]) | Q(tags="基础模块"))
|
||
await user_role.apis.add(*basic_apis)
|
||
|
||
|
||
async def init_clickhouse():
|
||
"""初始化ClickHouse数据库（若未配置则跳过）"""
|
||
host = settings.CLICKHOUSE_HOST or ""
|
||
if not host:
|
||
return
|
||
try:
|
||
client = await clickhouse_manager.get_client()
|
||
initializer = ClickHouseInitializer(client)
|
||
await initializer.initialize_all_tables()
|
||
logger.info("ClickHouse初始化完成")
|
||
except Exception as e:
|
||
logger.error(f"ClickHouse初始化失败: {e}")
|
||
|
||
|
||
async def init_data():
|
||
"""应用启动数据初始化：受环境变量控制并在多进程下只执行一次"""
|
||
should_migrate = settings.RUN_MIGRATIONS_ON_STARTUP
|
||
should_seed = settings.INITIALIZE_SEED_DATA_ON_STARTUP
|
||
|
||
lock_dir = ".startup_lock"
|
||
acquired = False
|
||
try:
|
||
# 简单文件锁，避免多 worker 并发执行
|
||
import os
|
||
os.mkdir(lock_dir)
|
||
acquired = True
|
||
except Exception:
|
||
acquired = False
|
||
|
||
if should_migrate and acquired:
|
||
await init_db()
|
||
|
||
if should_seed and acquired:
|
||
await init_superuser()
|
||
await init_menus()
|
||
await init_apis()
|
||
await init_roles()
|
||
|
||
# ClickHouse 初始化为可选，且不影响主应用
|
||
await init_clickhouse()
|
||
|
||
if acquired:
|
||
try:
|
||
import os
|
||
os.rmdir(lock_dir)
|
||
except Exception:
|
||
pass
|