tianfeng_task_modules/app.py

# app.py
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import logging
import time
from utils.logger import setup_logger
from config.component_registry import register_all_components
from config.settings import (
    AppConfig, ServerConfig, ApiConfig, CorsConfig
)
from api.task_api import router as task_router
from api.workflow_api import router as workflow_router
from api.component_api import router as component_router
from core.exceptions import TianfengTaskError
from config.api_config import ApiResponseCode, ApiResponseMessage
from config.component_config import ComponentCategoryConfig
# 导入数据库相关模块
from config.database import DBConfig, CacheConfig, db_session
import data.models  # 导入所有模型以确保它们被注册
from data.models.component import ComponentCategory, ComponentType, Component, ComponentCategoryEnum
# 导入Lifespan
from contextlib import asynccontextmanager

# 设置日志
logger = setup_logger()

# 定义Lifespan上下文管理器
@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    应用生命周期管理
    在应用启动时执行初始化操作，在应用关闭时执行清理操作
    """
    # 应用启动时执行
    logger.info("应用启动")
    
    yield  # 应用运行期间
    
    # 应用关闭时执行
    logger.info("应用关闭")
    DBConfig.shutdown_session()

# 创建FastAPI应用，使用Lifespan
app = FastAPI(
    title=ApiConfig.TITLE,
    description=ApiConfig.DESCRIPTION,
    version=ApiConfig.VERSION,
    lifespan=lifespan
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=CorsConfig.ALLOW_ORIGINS,
    allow_credentials=CorsConfig.ALLOW_CREDENTIALS,
    allow_methods=CorsConfig.ALLOW_METHODS,
    allow_headers=CorsConfig.ALLOW_HEADERS,
)

# 初始化数据库
def init_database():
    """初始化数据库，创建所有表"""
    try:
        logger.info("开始初始化数据库...")
        # 初始化数据库表
        DBConfig.init_db()
        logger.info("数据库表创建成功")
        
        # 初始化基础数据
        logger.info("开始初始化基础数据...")
        init_base_data(db_session)
        logger.info("基础数据初始化成功")
        
        return True
    except Exception as e:
        logger.error(f"数据库初始化失败: {str(e)}")
        raise

def init_base_data(db_session):
    """初始化基础数据"""
    try:
        # 导入需要的模型        
        # 检查是否已存在组件分类
        existing_categories = db_session.query(ComponentCategory).filter(ComponentCategory.is_deleted == False).all()
        if existing_categories:
            logger.info("基础数据已存在，跳过初始化")
            return
        
        # 创建组件分类
        categories = []
        for category_enum in ComponentCategoryEnum:
            category = ComponentCategory(
                name=ComponentCategoryConfig.get_category_name(category_enum),
                code=category_enum,
                description=ComponentCategoryConfig.get_category_description(category_enum),
                icon=f"icon-{category_enum.value}",
                order=ComponentCategoryConfig.get_category_order(category_enum)
            )
            categories.append(category)
            db_session.add(category)
        
        db_session.commit()
        logger.info(f"创建了 {len(categories)} 个组件分类")
        
        # 这里可以添加更多基础数据的初始化，如组件类型、系统组件等
        
    except Exception as e:
        db_session.rollback()
        logger.error(f"基础数据初始化失败: {str(e)}")
        raise

# 初始化数据库
init_database()

# 注册所有组件
register_all_components()

# 注册API路由
app.include_router(task_router, prefix=ApiConfig.PREFIX)
app.include_router(workflow_router, prefix=ApiConfig.PREFIX)
app.include_router(component_router, prefix=ApiConfig.PREFIX)

# 请求中间件
@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    response.headers["X-Process-Time"] = str(process_time)
    return response

# 数据库会话中间件
@app.middleware("http")
async def db_session_middleware(request: Request, call_next):
    try:
        response = await call_next(request)
        return response
    finally:
        db_session.remove()

# 全局异常处理
@app.exception_handler(TianfengTaskError)
async def tianfeng_task_error_handler(request: Request, exc: TianfengTaskError):
    """处理天风任务模块异常"""
    return JSONResponse(
        status_code=ApiResponseCode.BAD_REQUEST,
        content={
            "code": ApiResponseCode.BAD_REQUEST,
            "message": str(exc),
            "data": None
        }
    )

@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    """处理通用异常"""
    logger.exception("未处理的异常")
    return JSONResponse(
        status_code=ApiResponseCode.SERVER_ERROR,
        content={
            "code": ApiResponseCode.SERVER_ERROR,
            "message": f"{ApiResponseMessage.SERVER_ERROR}: {str(exc)}",
            "data": None
        }
    )

# 健康检查接口
@app.get(f"{ApiConfig.PREFIX}/health")
async def health_check():
    """健康检查"""
    return {
        "code": ApiResponseCode.SUCCESS,
        "message": "服务正常",
        "data": {
            "app_name": AppConfig.NAME,
            "app_version": AppConfig.VERSION,
            "status": "healthy"
        }
    }

# 主函数
if __name__ == "__main__":
    import uvicorn
    logger.info(f"启动天风任务模块服务，调试模式: {AppConfig.DEBUG}")
    uvicorn.run("app:app", host=ServerConfig.HOST, port=ServerConfig.PORT, reload=ServerConfig.RELOAD)