gardel/Mofox-Core
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

数据库重构

Browse Source
This commit is contained in:
雅诺狐
2025-08-16 23:43:45 +08:00
parent 0f0619762b
commit d46d689c43
21 changed files with 834 additions and 1007 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
src/common/database/database.py
View File

@@ -6,7 +6,7 @@ from src.common.logger import get_logger
# SQLAlchemy相关导入
from src.common.database.sqlalchemy_init import initialize_database_compat
from src.common.database.sqlalchemy_models import get_engine, get_session
from src.common.database.sqlalchemy_models import get_engine, get_db_session
install(extra_lines=3)
@@ -18,7 +18,7 @@ logger = get_logger("database")
# 兼容性：为了不破坏现有代码，保留db变量但指向SQLAlchemy
class DatabaseProxy:
"""数据库代理类，提供Peewee到SQLAlchemy的兼容性接口"""
"""数据库代理类"""
def __init__(self):
self._engine = None
@@ -28,56 +28,7 @@ class DatabaseProxy:
"""初始化数据库连接"""
return initialize_database_compat()
def connect(self, reuse_if_open=True):
"""连接数据库（兼容性方法）"""
try:
self._engine = get_engine()
return True
except Exception as e:
logger.error(f"数据库连接失败: {e}")
return False
def is_closed(self):
"""检查数据库是否关闭（兼容性方法）"""
return self._engine is None
def create_tables(self, models, safe=True):
"""创建表（兼容性方法）"""
try:
from src.common.database.sqlalchemy_models import Base
engine = get_engine()
Base.metadata.create_all(bind=engine)
return True
except Exception as e:
logger.error(f"创建表失败: {e}")
return False
def table_exists(self, model):
"""检查表是否存在（兼容性方法）"""
try:
from sqlalchemy import inspect
engine = get_engine()
inspector = inspect(engine)
table_name = getattr(model, '_meta', {}).get('table_name', model.__name__.lower())
return table_name in inspector.get_table_names()
except Exception:
return False
def execute_sql(self, sql):
"""执行SQL（兼容性方法）"""
try:
from sqlalchemy import text
session = get_session()
result = session.execute(text(sql))
session.close()
return result
except Exception as e:
logger.error(f"执行SQL失败: {e}")
raise
def atomic(self):
"""事务上下文管理器（兼容性方法）"""
return SQLAlchemyTransaction()
class SQLAlchemyTransaction:
"""SQLAlchemy事务上下文管理器"""
@@ -86,7 +37,7 @@ class SQLAlchemyTransaction:
self.session = None
def __enter__(self):
self.session = get_session()
self.session = get_db_session()
return self.session
def __exit__(self, exc_type, exc_val, exc_tb):

39
src/common/database/sqlalchemy_database_api.py
View File

@@ -15,7 +15,7 @@ from src.common.logger import get_logger
from src.common.database.sqlalchemy_models import (
Base, get_db_session, Messages, ActionRecords, PersonInfo, ChatStreams,
LLMUsage, Emoji, Images, ImageDescriptions, OnlineTime, Memory,
Expression, ThinkingLog, GraphNodes, GraphEdges, Schedule, MaiZoneScheduleStatus, get_session
Expression, ThinkingLog, GraphNodes, GraphEdges, Schedule, MaiZoneScheduleStatus
)
logger = get_logger("sqlalchemy_database_api")
@@ -41,38 +41,9 @@ MODEL_MAPPING = {
}
@contextmanager
def get_db_session():
"""数据库会话上下文管理器，自动处理事务和连接错误"""
session = None
max_retries = 3
retry_delay = 1.0
for attempt in range(max_retries):
try:
session = get_session()
yield session
session.commit()
break
except (DisconnectionError, OperationalError) as e:
logger.warning(f"数据库连接错误 (尝试 {attempt + 1}/{max_retries}): {e}")
if session:
session.rollback()
session.close()
if attempt < max_retries - 1:
time.sleep(retry_delay * (attempt + 1))
else:
raise
except Exception:
if session:
session.rollback()
raise
finally:
if session:
session.close()
def build_filters(session: Session, model_class: Type[Base], filters: Dict[str, Any]):
def build_filters(session, model_class: Type[Base], filters: Dict[str, Any]):
"""构建查询过滤条件"""
conditions = []
@@ -296,6 +267,7 @@ async def db_save(
# 创建新记录
new_record = model_class(**data)
session.add(new_record)
session.commit()
session.flush()
# 转换为字典格式返回
@@ -415,8 +387,3 @@ async def store_action_info(
traceback.print_exc()
return None
# 兼容性函数，方便从Peewee迁移
def get_model_class(model_name: str) -> Optional[Type[Base]]:
"""根据模型名称获取模型类"""
return MODEL_MAPPING.get(model_name)

35
src/common/database/sqlalchemy_init.py
View File

@@ -8,7 +8,7 @@ from typing import Optional
from sqlalchemy.exc import SQLAlchemyError
from src.common.logger import get_logger
from src.common.database.sqlalchemy_models import (
Base, get_engine, get_session, initialize_database
Base, get_engine, initialize_database
)
logger = get_logger("sqlalchemy_init")
@@ -72,36 +72,6 @@ def create_all_tables() -> bool:
return False
def check_database_connection() -> bool:
"""
检查数据库连接是否正常
Returns:
bool: 连接是否正常
"""
try:
session = get_session()
if session is None:
logger.error("无法获取数据库会话")
return False
# 检查会话是否可用（如果能获取到会话说明连接正常）
if session is None:
logger.error("数据库会话无效")
return False
session.close()
logger.info("数据库连接检查通过")
return True
except SQLAlchemyError as e:
logger.error(f"数据库连接检查失败: {e}")
return False
except Exception as e:
logger.error(f"数据库连接检查过程中发生未知错误: {e}")
return False
def get_database_info() -> Optional[dict]:
"""
@@ -149,9 +119,6 @@ def initialize_database_compat() -> bool:
if success:
success = create_all_tables()
if success:
success = check_database_connection()
if success:
_database_initialized = True

115
src/common/database/sqlalchemy_models.py
View File

@@ -29,102 +29,6 @@ def get_string_field(max_length=255, **kwargs):
return String(max_length, **kwargs)
else:
return Text(**kwargs)
class SessionProxy:
"""线程安全的Session代理类，自动管理session生命周期"""
def __init__(self):
self._local = threading.local()
def _get_current_session(self):
"""获取当前线程的session，如果没有则创建新的"""
if not hasattr(self._local, 'session') or self._local.session is None:
_, SessionLocal = initialize_database()
self._local.session = SessionLocal()
return self._local.session
def _close_current_session(self):
"""关闭当前线程的session"""
if hasattr(self._local, 'session') and self._local.session is not None:
try:
self._local.session.close()
except:
pass
finally:
self._local.session = None
def __getattr__(self, name):
"""代理所有session方法"""
session = self._get_current_session()
attr = getattr(session, name)
# 如果是方法，需要特殊处理一些关键方法
if callable(attr):
if name in ['commit', 'rollback']:
def wrapper(*args, **kwargs):
try:
result = attr(*args, **kwargs)
if name == 'commit':
# commit后不要清除session，只是刷新状态
pass # 保持session活跃
return result
except Exception:
try:
if session and hasattr(session, 'rollback'):
session.rollback()
except:
pass
# 发生错误时重新创建session
self._close_current_session()
raise
return wrapper
elif name == 'close':
def wrapper(*args, **kwargs):
result = attr(*args, **kwargs)
self._close_current_session()
return result
return wrapper
elif name in ['execute', 'query', 'add', 'delete', 'merge']:
def wrapper(*args, **kwargs):
try:
return attr(*args, **kwargs)
except Exception as e:
# 如果是连接相关错误，重新创建session再试一次
if "not bound to a Session" in str(e) or "provisioning a new connection" in str(e):
logger.warning(f"Session问题，重新创建session: {e}")
self._close_current_session()
new_session = self._get_current_session()
new_attr = getattr(new_session, name)
return new_attr(*args, **kwargs)
raise
return wrapper
return attr
def new_session(self):
"""强制创建新的session（关闭当前的，创建新的）"""
self._close_current_session()
return self._get_current_session()
def ensure_fresh_session(self):
"""确保使用新鲜的session（如果当前session有问题则重新创建）"""
if hasattr(self._local, 'session') and self._local.session is not None:
try:
# 测试session是否还可用
self._local.session.execute("SELECT 1")
except Exception:
# session有问题，重新创建
self._close_current_session()
return self._get_current_session()
# 创建全局session代理实例
_global_session_proxy = SessionProxy()
def get_session():
"""返回线程安全的session代理，自动管理生命周期"""
return _global_session_proxy
class ChatStreams(Base):
@@ -482,6 +386,22 @@ class MaiZoneScheduleStatus(Base):
)
class BanUser(Base):
"""被禁用用户模型"""
__tablename__ = 'ban_users'
id = Column(Integer, primary_key=True, autoincrement=True)
user_id = Column(get_string_field(50), nullable=False, index=True)
violation_num = Column(Integer, nullable=False, default=0)
reason = Column(Text, nullable=False)
created_at = Column(DateTime, nullable=False, default=datetime.datetime.now)
__table_args__ = (
Index('idx_violation_num', 'violation_num'),
Index('idx_banuser_user_id', 'user_id'),
)
# 数据库引擎和会话管理
_engine = None
_SessionLocal = None
@@ -593,7 +513,7 @@ def get_db_session():
_, SessionLocal = initialize_database()
session = SessionLocal()
yield session
session.commit()
# session.commit()
except Exception:
if session:
session.rollback()
@@ -601,6 +521,7 @@ def get_db_session():
finally:
if session:
session.close()
def get_engine():