122 lines
4.1 KiB
Python
Executable File
122 lines
4.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
import sys
|
|
|
|
from typing import Annotated
|
|
from uuid import uuid4
|
|
|
|
from fastapi import Depends
|
|
from sqlalchemy import URL, text
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine, AsyncEngine
|
|
from sqlalchemy.orm import Session
|
|
|
|
from backend.common.log import log
|
|
from backend.common.model import MappedBase
|
|
from backend.core.conf import settings, get_db_uri
|
|
|
|
|
|
def create_async_engine_and_session(
|
|
url: str | URL,
|
|
echo: bool = False,
|
|
pool_size: int = 10,
|
|
max_overflow: int = 20,
|
|
pool_timeout: int = 30,
|
|
pool_recycle: int = 3600,
|
|
pool_pre_ping: bool = True,
|
|
application_name: str = "app"
|
|
) -> tuple[create_async_engine, async_sessionmaker[AsyncSession], async_sessionmaker[AsyncSession]]:
|
|
"""
|
|
创建 MySQL 异步引擎和会话工厂
|
|
参数优化说明:
|
|
- pool_size: 建议设置为 (核心数 * 2) + 有效磁盘数
|
|
- max_overflow: 峰值连接缓冲，避免连接风暴
|
|
- pool_recycle: 防止 MySQL 连接超时 (默认为 1 小时)
|
|
- pool_pre_ping: 强烈建议开启，处理连接失效问题
|
|
- application_name: 帮助 DBA 识别连接来源
|
|
"""
|
|
|
|
try:
|
|
# 创建异步引擎 (针对 MySQL 优化)
|
|
engine = create_async_engine(
|
|
url,
|
|
echo=echo,
|
|
echo_pool=echo,
|
|
future=True,
|
|
connect_args={
|
|
"charset": "utf8mb4", # MySQL 特定字符集
|
|
"autocommit": True, # 自动提交
|
|
"connect_timeout": 60, # 连接超时
|
|
},
|
|
pool_size=pool_size,
|
|
max_overflow=max_overflow,
|
|
pool_timeout=pool_timeout,
|
|
pool_recycle=pool_recycle,
|
|
pool_pre_ping=pool_pre_ping,
|
|
pool_use_lifo=True, # 使用 LIFO 提高连接池效率
|
|
# MySQL 特定优化参数
|
|
poolclass=None, # 使用默认 QueuePool
|
|
execution_options={
|
|
"isolation_level": "READ COMMITTED", # MySQL 推荐隔离级别
|
|
"compiled_cache": None # 禁用缓存，避免内存泄漏
|
|
}
|
|
)
|
|
background_engine = create_async_engine(
|
|
url,
|
|
pool_size=5,
|
|
max_overflow=10,
|
|
pool_pre_ping=True,
|
|
pool_recycle=300,
|
|
connect_args={
|
|
"charset": "utf8mb4",
|
|
"autocommit": True,
|
|
"connect_timeout": 60,
|
|
}
|
|
)
|
|
except Exception as e:
|
|
log.error(f'❌ MySQL 数据库连接失败: {e}')
|
|
sys.exit(1)
|
|
else:
|
|
# 创建异步会话工厂 (针对 MySQL 优化)
|
|
db_session = async_sessionmaker(
|
|
bind=engine,
|
|
autoflush=False,
|
|
expire_on_commit=False,
|
|
# MySQL 特定优化
|
|
class_=AsyncSession,
|
|
twophase=False, # 禁用两阶段提交
|
|
enable_baked_queries=False, # 禁用 baked 查询避免内存问题
|
|
info={"app_name": application_name} # 添加应用标识
|
|
)
|
|
background_db_session = async_sessionmaker(
|
|
background_engine,
|
|
expire_on_commit=False,
|
|
autoflush=False
|
|
)
|
|
|
|
log.info(f'✅ MySQL 异步引擎创建成功 | 连接池: [{pool_size}] - [{max_overflow}]')
|
|
return engine, db_session, background_db_session
|
|
|
|
|
|
async def get_db():
|
|
"""session 生成器"""
|
|
async with async_db_session() as session:
|
|
yield session
|
|
|
|
|
|
async def create_table() -> None:
|
|
"""创建数据库表"""
|
|
async with async_engine.begin() as coon:
|
|
await coon.run_sync(MappedBase.metadata.create_all)
|
|
|
|
|
|
def uuid4_str() -> str:
|
|
"""数据库引擎 UUID 类型兼容性解决方案"""
|
|
return str(uuid4())
|
|
|
|
|
|
SQLALCHEMY_DATABASE_URL = get_db_uri(settings)
|
|
|
|
async_engine, async_db_session, background_db_session = create_async_engine_and_session(SQLALCHEMY_DATABASE_URL)
|
|
|
|
# Session Annotated
|
|
CurrentSession = Annotated[AsyncSession, Depends(get_db)] |