287 lines
8.9 KiB
Python
287 lines
8.9 KiB
Python
"""数据库引擎管理
|
||
|
||
单一职责：创建和管理SQLAlchemy异步引擎
|
||
|
||
支持的数据库类型：
|
||
- SQLite: 轻量级本地数据库，使用 aiosqlite 驱动
|
||
- PostgreSQL: 功能丰富的开源数据库，使用 asyncpg 驱动
|
||
"""
|
||
|
||
import asyncio
|
||
import os
|
||
from urllib.parse import quote_plus
|
||
|
||
from sqlalchemy import text
|
||
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
|
||
|
||
from src.common.logger import get_logger
|
||
|
||
from ..utils.exceptions import DatabaseInitializationError
|
||
from .dialect_adapter import DialectAdapter
|
||
|
||
logger = get_logger("database.engine")
|
||
|
||
# 全局引擎实例
|
||
_engine: AsyncEngine | None = None
|
||
_engine_lock: asyncio.Lock | None = None
|
||
|
||
|
||
async def get_engine() -> AsyncEngine:
|
||
"""获取全局数据库引擎（单例模式）
|
||
|
||
Returns:
|
||
AsyncEngine: SQLAlchemy异步引擎
|
||
|
||
Raises:
|
||
DatabaseInitializationError: 引擎初始化失败
|
||
"""
|
||
global _engine, _engine_lock
|
||
|
||
# 快速路径：引擎已初始化
|
||
if _engine is not None:
|
||
return _engine
|
||
|
||
# 延迟创建锁（避免在导入时创建）
|
||
if _engine_lock is None:
|
||
_engine_lock = asyncio.Lock()
|
||
|
||
assert _engine_lock is not None
|
||
# 使用锁保护初始化过程
|
||
async with _engine_lock:
|
||
# 双重检查锁定模式
|
||
if _engine is not None:
|
||
return _engine
|
||
|
||
try:
|
||
from src.config.config import global_config
|
||
|
||
assert global_config is not None
|
||
config = global_config.database
|
||
db_type = config.database_type
|
||
|
||
# 初始化方言适配器
|
||
DialectAdapter.initialize(db_type)
|
||
|
||
logger.info(f"正在初始化 {db_type.upper()} 数据库引擎...")
|
||
|
||
# 根据数据库类型构建URL和引擎参数
|
||
if db_type == "postgresql":
|
||
url, engine_kwargs = _build_postgresql_config(config)
|
||
else:
|
||
url, engine_kwargs = _build_sqlite_config(config)
|
||
|
||
# 创建异步引擎
|
||
_engine = create_async_engine(url, **engine_kwargs)
|
||
|
||
# 数据库特定优化
|
||
if db_type == "sqlite":
|
||
await _enable_sqlite_optimizations(_engine)
|
||
elif db_type == "postgresql":
|
||
await _enable_postgresql_optimizations(_engine)
|
||
|
||
logger.info(f"{db_type.upper()} 数据库引擎初始化成功")
|
||
return _engine
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 数据库引擎初始化失败: {e}")
|
||
raise DatabaseInitializationError(f"引擎初始化失败: {e}") from e
|
||
|
||
|
||
def _build_sqlite_config(config) -> tuple[str, dict]:
|
||
"""构建 SQLite 配置
|
||
|
||
Args:
|
||
config: 数据库配置对象
|
||
|
||
Returns:
|
||
(url, engine_kwargs) 元组
|
||
"""
|
||
if not os.path.isabs(config.sqlite_path):
|
||
ROOT_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", ".."))
|
||
db_path = os.path.join(ROOT_PATH, config.sqlite_path)
|
||
else:
|
||
db_path = config.sqlite_path
|
||
|
||
# 确保数据库目录存在
|
||
os.makedirs(os.path.dirname(db_path), exist_ok=True)
|
||
|
||
url = f"sqlite+aiosqlite:///{db_path}"
|
||
|
||
engine_kwargs = {
|
||
"echo": False,
|
||
"future": True,
|
||
"connect_args": {
|
||
"check_same_thread": False,
|
||
"timeout": 60,
|
||
},
|
||
}
|
||
|
||
logger.debug(f"SQLite配置: {db_path}")
|
||
return url, engine_kwargs
|
||
|
||
|
||
def _build_postgresql_config(config) -> tuple[str, dict]:
|
||
"""构建 PostgreSQL 配置
|
||
|
||
Args:
|
||
config: 数据库配置对象
|
||
|
||
Returns:
|
||
(url, engine_kwargs) 元组
|
||
"""
|
||
encoded_user = quote_plus(config.postgresql_user)
|
||
encoded_password = quote_plus(config.postgresql_password)
|
||
|
||
# 构建基本 URL
|
||
url = (
|
||
f"postgresql+asyncpg://{encoded_user}:{encoded_password}"
|
||
f"@{config.postgresql_host}:{config.postgresql_port}/{config.postgresql_database}"
|
||
)
|
||
|
||
# SSL 配置
|
||
connect_args = {}
|
||
if config.postgresql_ssl_mode != "disable":
|
||
ssl_config = {"ssl": config.postgresql_ssl_mode}
|
||
if config.postgresql_ssl_ca:
|
||
ssl_config["ssl_ca"] = config.postgresql_ssl_ca
|
||
if config.postgresql_ssl_cert:
|
||
ssl_config["ssl_cert"] = config.postgresql_ssl_cert
|
||
if config.postgresql_ssl_key:
|
||
ssl_config["ssl_key"] = config.postgresql_ssl_key
|
||
connect_args.update(ssl_config)
|
||
|
||
# 设置 schema（如果不是 public）
|
||
if config.postgresql_schema and config.postgresql_schema != "public":
|
||
connect_args["server_settings"] = {"search_path": config.postgresql_schema}
|
||
|
||
engine_kwargs = {
|
||
"echo": False,
|
||
"future": True,
|
||
"pool_size": config.connection_pool_size,
|
||
"max_overflow": config.connection_pool_size * 2,
|
||
"pool_timeout": config.connection_timeout,
|
||
"pool_recycle": 3600,
|
||
"pool_pre_ping": True,
|
||
}
|
||
|
||
if connect_args:
|
||
engine_kwargs["connect_args"] = connect_args
|
||
|
||
logger.debug(
|
||
f"PostgreSQL配置: {config.postgresql_user}@{config.postgresql_host}:{config.postgresql_port}/{config.postgresql_database}"
|
||
)
|
||
return url, engine_kwargs
|
||
|
||
|
||
async def close_engine():
|
||
"""关闭数据库引擎
|
||
|
||
释放所有连接池资源
|
||
"""
|
||
global _engine
|
||
|
||
if _engine is not None:
|
||
logger.info("正在关闭数据库引擎...")
|
||
await _engine.dispose()
|
||
_engine = None
|
||
logger.info("数据库引擎已关闭")
|
||
|
||
|
||
async def _enable_sqlite_optimizations(engine: AsyncEngine):
|
||
"""启用SQLite性能优化
|
||
|
||
优化项：
|
||
- WAL模式：提高并发性能
|
||
- NORMAL同步：平衡性能和安全性
|
||
- 启用外键约束
|
||
- 设置busy_timeout：避免锁定错误
|
||
|
||
Args:
|
||
engine: SQLAlchemy异步引擎
|
||
"""
|
||
try:
|
||
async with engine.begin() as conn:
|
||
# 启用WAL模式
|
||
await conn.execute(text("PRAGMA journal_mode = WAL"))
|
||
# 设置适中的同步级别
|
||
await conn.execute(text("PRAGMA synchronous = NORMAL"))
|
||
# 启用外键约束
|
||
await conn.execute(text("PRAGMA foreign_keys = ON"))
|
||
# 设置busy_timeout，避免锁定错误
|
||
await conn.execute(text("PRAGMA busy_timeout = 10000"))
|
||
# 设置缓存大小（10MB）
|
||
await conn.execute(text("PRAGMA cache_size = -10000"))
|
||
# 临时存储使用内存
|
||
await conn.execute(text("PRAGMA temp_store = MEMORY"))
|
||
|
||
except Exception as e:
|
||
logger.warning(f"⚠️ SQLite性能优化失败: {e}，将使用默认配置")
|
||
|
||
|
||
async def _enable_postgresql_optimizations(engine: AsyncEngine) -> None:
|
||
"""启用 PostgreSQL 会话级性能优化。
|
||
|
||
包含：
|
||
- work_mem：单次排序/哈希可用内存
|
||
- statement_timeout：单条语句超时
|
||
- synchronous_commit：提交同步级别
|
||
- jit：是否启用 JIT
|
||
- idle_in_transaction_session_timeout：事务空闲超时
|
||
- lock_timeout：等待锁超时
|
||
"""
|
||
try:
|
||
async with engine.begin() as conn:
|
||
# 排序/哈希内存，注意这是“每个排序操作”的上限，不要乱调太大
|
||
await conn.execute(text("SET work_mem = '64MB'"))
|
||
|
||
# 单条语句超时（1 分钟）
|
||
await conn.execute(text("SET statement_timeout = '60000'"))
|
||
|
||
# 提交同步级别：'local' 性能好一些，崩溃时可能丢几 ms 的数据
|
||
await conn.execute(text("SET synchronous_commit = 'local'"))
|
||
|
||
# 对大量短小查询，通常关掉 JIT 更省 CPU
|
||
await conn.execute(text("SET jit = 'off'"))
|
||
|
||
# 事务空闲超时，避免长时间占用锁
|
||
await conn.execute(
|
||
text("SET idle_in_transaction_session_timeout = '60000'")
|
||
)
|
||
|
||
# 等锁超过 5s 就报错，避免全部堆死
|
||
await conn.execute(text("SET lock_timeout = '5000'"))
|
||
|
||
# 可选：根据你预估的缓存情况设置
|
||
# await conn.execute(text("SET random_page_cost = 1.1"))
|
||
# await conn.execute(text("SET effective_cache_size = '2GB'"))
|
||
|
||
except Exception as e:
|
||
logger.warning(
|
||
f"⚠️ PostgreSQL 会话级性能优化失败: {e}，将使用默认配置"
|
||
)
|
||
|
||
|
||
async def get_engine_info() -> dict:
|
||
"""获取引擎信息（用于监控和调试）
|
||
|
||
Returns:
|
||
dict: 引擎信息字典
|
||
"""
|
||
try:
|
||
engine = await get_engine()
|
||
|
||
info = {
|
||
"name": engine.name,
|
||
"driver": engine.driver,
|
||
"url": str(engine.url).replace(str(engine.url.password or ""), "***"),
|
||
"pool_size": getattr(engine.pool, "size", lambda: None)(),
|
||
"pool_checked_out": getattr(engine.pool, "checked_out", lambda: 0)(),
|
||
"pool_overflow": getattr(engine.pool, "overflow", lambda: 0)(),
|
||
}
|
||
|
||
return info
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取引擎信息失败: {e}")
|
||
return {}
|