达梦数据库适配

上代码

1.db_connect.py

python 复制代码

from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker, declarative_base
from sqlalchemy import event, text

from project.core.config import settings
from project.core.config.settings import logger

if settings.DBConfig.CONNECT_METHOD == 'dm':
    # 达梦数据库引擎配置
    engine = create_async_engine(
        settings.DBConfig.DM.DM8_URL,
        echo=settings.DBConfig.DM.ECHO,  # 根据环境配置是否输出SQL
    )
else:
    # 创建mysql异步数据库引擎
    engine = create_async_engine(
        settings.DBConfig.MySql.MYSQL_URL,
        # 连接池配置
        pool_size=settings.DBConfig.MySql.POOL_SIZE,  # 4核 × 2 + 1
        max_overflow=settings.DBConfig.MySql.MAX_OVERFLOW,  # 允许临时增加5个连接
        pool_recycle=settings.DBConfig.MySql.POOL_RECYCLE,  # 30分钟回收连接
        pool_timeout=settings.DBConfig.MySql.POOL_TIMEOUT,  # 获取连接超时30秒
        pool_pre_ping=settings.DBConfig.MySql.POOL_PRE_PING,  # 自动检查连接有效性
        echo=settings.DBConfig.MySql.ECHO,  # 根据环境配置是否输出SQL
        connect_args={
            "connect_timeout": settings.DBConfig.MySql.CONNECT_TIMEOUT,
        }
    )

# 创建异步会话工厂
SessionLocal = sessionmaker(
    bind=engine,
    class_=AsyncSession,
    autocommit=False,
    autoflush=False,
    expire_on_commit=False  # 避免提交后属性访问问题
)

# 声明基类
Base = declarative_base()


async def verify_connection():
    """验证数据库连接是否可用"""
    async with engine.connect() as conn:
        await conn.execute(text("SELECT 1"))
        logger.info("Database connection verified successfully")


async def dispose_engine():
    """关闭连接池并释放所有连接"""
    await engine.dispose()
    logger.info("Database engine disposed successfully")

2.db_session.py

python 复制代码

from contextlib import asynccontextmanager
from fastapi import HTTPException, status
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.exc import SQLAlchemyError

from project.core.config.settings import logger
from project.core.database.db_connect import SessionLocal
import asyncio
from project.core.database.db_connect import verify_connection, dispose_engine


@asynccontextmanager
async def _get_db() -> AsyncSession:
    """内部使用的上下文管理器"""
    session = SessionLocal()
    try:
        yield session
        await session.commit()
    except SQLAlchemyError as e:
        await session.rollback()
        logger.error(f"Rollback failed: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Database operation failed: {str(e)}"
        ) from e
    finally:
        await session.close()


async def get_db() -> AsyncSession:
    """供路由使用的依赖项"""
    async with _get_db() as session:
        yield session


async def test_connection():
    try:
        await verify_connection()
    except Exception as e:
        logger.error(f"Database connection failed: {e}")
    finally:
        await dispose_engine()


asyncio.run(test_connection())

if __name__ == '__main__':
    # 测试数据库连接
    async def test_connection():
        try:
            await verify_connection()
            print("✅ Database connection test successful")
        except Exception as e:
            print(f"❌ Database connection failed: {e}")
        finally:
            await dispose_engine()


    asyncio.run(test_connection())