feat: 更新机器人配置并添加数据库迁移脚本
- 将bot_config_template.toml中的版本升级至7.9.0 - 增强数据库配置选项以支持PostgreSQL - 引入一个新脚本,用于在SQLite、MySQL和PostgreSQL之间迁移数据 - 实现一个方言适配器,用于处理特定于数据库的行为和配置
This commit is contained in:
@@ -1,6 +1,11 @@
|
||||
"""数据库引擎管理
|
||||
|
||||
单一职责:创建和管理SQLAlchemy异步引擎
|
||||
|
||||
支持的数据库类型:
|
||||
- SQLite: 轻量级本地数据库,使用 aiosqlite 驱动
|
||||
- MySQL: 高性能关系型数据库,使用 aiomysql 驱动
|
||||
- PostgreSQL: 功能丰富的开源数据库,使用 asyncpg 驱动
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -13,6 +18,7 @@ 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")
|
||||
|
||||
@@ -52,79 +58,27 @@ async def get_engine() -> AsyncEngine:
|
||||
config = global_config.database
|
||||
db_type = config.database_type
|
||||
|
||||
# 初始化方言适配器
|
||||
DialectAdapter.initialize(db_type)
|
||||
|
||||
logger.info(f"正在初始化 {db_type.upper()} 数据库引擎...")
|
||||
|
||||
# 构建数据库URL和引擎参数
|
||||
# 根据数据库类型构建URL和引擎参数
|
||||
if db_type == "mysql":
|
||||
# MySQL配置
|
||||
encoded_user = quote_plus(config.mysql_user)
|
||||
encoded_password = quote_plus(config.mysql_password)
|
||||
|
||||
if config.mysql_unix_socket:
|
||||
# Unix socket连接
|
||||
encoded_socket = quote_plus(config.mysql_unix_socket)
|
||||
url = (
|
||||
f"mysql+aiomysql://{encoded_user}:{encoded_password}"
|
||||
f"@/{config.mysql_database}"
|
||||
f"?unix_socket={encoded_socket}&charset={config.mysql_charset}"
|
||||
)
|
||||
else:
|
||||
# TCP连接
|
||||
url = (
|
||||
f"mysql+aiomysql://{encoded_user}:{encoded_password}"
|
||||
f"@{config.mysql_host}:{config.mysql_port}/{config.mysql_database}"
|
||||
f"?charset={config.mysql_charset}"
|
||||
)
|
||||
|
||||
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,
|
||||
"connect_args": {
|
||||
"autocommit": config.mysql_autocommit,
|
||||
"charset": config.mysql_charset,
|
||||
"connect_timeout": config.connection_timeout,
|
||||
},
|
||||
}
|
||||
|
||||
logger.info(
|
||||
f"MySQL配置: {config.mysql_user}@{config.mysql_host}:{config.mysql_port}/{config.mysql_database}"
|
||||
)
|
||||
|
||||
url, engine_kwargs = _build_mysql_config(config)
|
||||
elif db_type == "postgresql":
|
||||
url, engine_kwargs = _build_postgresql_config(config)
|
||||
else:
|
||||
# SQLite配置
|
||||
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.info(f"SQLite配置: {db_path}")
|
||||
url, engine_kwargs = _build_sqlite_config(config)
|
||||
|
||||
# 创建异步引擎
|
||||
_engine = create_async_engine(url, **engine_kwargs)
|
||||
|
||||
# SQLite特定优化
|
||||
# 数据库特定优化
|
||||
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
|
||||
@@ -134,6 +88,141 @@ async def get_engine() -> AsyncEngine:
|
||||
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.info(f"SQLite配置: {db_path}")
|
||||
return url, engine_kwargs
|
||||
|
||||
|
||||
def _build_mysql_config(config) -> tuple[str, dict]:
|
||||
"""构建 MySQL 配置
|
||||
|
||||
Args:
|
||||
config: 数据库配置对象
|
||||
|
||||
Returns:
|
||||
(url, engine_kwargs) 元组
|
||||
"""
|
||||
encoded_user = quote_plus(config.mysql_user)
|
||||
encoded_password = quote_plus(config.mysql_password)
|
||||
|
||||
if config.mysql_unix_socket:
|
||||
# Unix socket连接
|
||||
encoded_socket = quote_plus(config.mysql_unix_socket)
|
||||
url = (
|
||||
f"mysql+aiomysql://{encoded_user}:{encoded_password}"
|
||||
f"@/{config.mysql_database}"
|
||||
f"?unix_socket={encoded_socket}&charset={config.mysql_charset}"
|
||||
)
|
||||
else:
|
||||
# TCP连接
|
||||
url = (
|
||||
f"mysql+aiomysql://{encoded_user}:{encoded_password}"
|
||||
f"@{config.mysql_host}:{config.mysql_port}/{config.mysql_database}"
|
||||
f"?charset={config.mysql_charset}"
|
||||
)
|
||||
|
||||
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,
|
||||
"connect_args": {
|
||||
"autocommit": config.mysql_autocommit,
|
||||
"charset": config.mysql_charset,
|
||||
"connect_timeout": config.connection_timeout,
|
||||
},
|
||||
}
|
||||
|
||||
logger.info(
|
||||
f"MySQL配置: {config.mysql_user}@{config.mysql_host}:{config.mysql_port}/{config.mysql_database}"
|
||||
)
|
||||
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.info(
|
||||
f"PostgreSQL配置: {config.postgresql_user}@{config.postgresql_host}:{config.postgresql_port}/{config.postgresql_database}"
|
||||
)
|
||||
return url, engine_kwargs
|
||||
|
||||
|
||||
async def close_engine():
|
||||
"""关闭数据库引擎
|
||||
|
||||
@@ -181,6 +270,33 @@ async def _enable_sqlite_optimizations(engine: AsyncEngine):
|
||||
logger.warning(f"⚠️ SQLite性能优化失败: {e},将使用默认配置")
|
||||
|
||||
|
||||
async def _enable_postgresql_optimizations(engine: AsyncEngine):
|
||||
"""启用PostgreSQL性能优化
|
||||
|
||||
优化项:
|
||||
- 设置合适的 work_mem
|
||||
- 启用 JIT 编译(如果可用)
|
||||
- 设置合适的 statement_timeout
|
||||
|
||||
Args:
|
||||
engine: SQLAlchemy异步引擎
|
||||
"""
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
# 设置会话级别的参数
|
||||
# work_mem: 排序和哈希操作的内存(64MB)
|
||||
await conn.execute(text("SET work_mem = '64MB'"))
|
||||
# 设置语句超时(5分钟)
|
||||
await conn.execute(text("SET statement_timeout = '300000'"))
|
||||
# 启用自动 EXPLAIN(可选,用于调试)
|
||||
# await conn.execute(text("SET auto_explain.log_min_duration = '1000'"))
|
||||
|
||||
logger.info("✅ PostgreSQL性能优化已启用")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"⚠️ PostgreSQL性能优化失败: {e},将使用默认配置")
|
||||
|
||||
|
||||
async def get_engine_info() -> dict:
|
||||
"""获取引擎信息(用于监控和调试)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user