数据库重构

src/common/database/database.py

@@ -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):

src/common/database/sqlalchemy_database_api.py

@@ -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)

src/common/database/sqlalchemy_init.py

@@ -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
     

src/common/database/sqlalchemy_models.py

@@ -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():

src/common/message_repository.py

@@ -8,7 +8,7 @@ from src.config.config import global_config
 
 # from src.common.database.database_model import Messages
 from src.common.database.sqlalchemy_models import Messages
-from src.common.database.sqlalchemy_database_api import get_session
+from src.common.database.sqlalchemy_database_api import get_db_session
 from src.common.logger import get_logger
 
 logger = get_logger(__name__)
@@ -44,92 +44,92 @@ def find_messages(
         消息字典列表，如果出错则返回空列表。
     """
     try:
-        session = get_session()
-        query = select(Messages)
+        with get_db_session() as session:
+            query = select(Messages)
 
-        # 应用过滤器
-        if message_filter:
-            conditions = []
-            for key, value in message_filter.items():
-                if hasattr(Messages, key):
-                    field = getattr(Messages, key)
-                    if isinstance(value, dict):
-                        # 处理 MongoDB 风格的操作符
-                        for op, op_value in value.items():
-                            if op == "$gt":
-                                conditions.append(field > op_value)
-                            elif op == "$lt":
-                                conditions.append(field < op_value)
-                            elif op == "$gte":
-                                conditions.append(field >= op_value)
-                            elif op == "$lte":
-                                conditions.append(field <= op_value)
-                            elif op == "$ne":
-                                conditions.append(field != op_value)
-                            elif op == "$in":
-                                conditions.append(field.in_(op_value))
-                            elif op == "$nin":
-                                conditions.append(field.not_in(op_value))
-                            else:
-                                logger.warning(f"过滤器中遇到未知操作符 '{op}' (字段: '{key}')。将跳过此操作符。")
+            # 应用过滤器
+            if message_filter:
+                conditions = []
+                for key, value in message_filter.items():
+                    if hasattr(Messages, key):
+                        field = getattr(Messages, key)
+                        if isinstance(value, dict):
+                            # 处理 MongoDB 风格的操作符
+                            for op, op_value in value.items():
+                                if op == "$gt":
+                                    conditions.append(field > op_value)
+                                elif op == "$lt":
+                                    conditions.append(field < op_value)
+                                elif op == "$gte":
+                                    conditions.append(field >= op_value)
+                                elif op == "$lte":
+                                    conditions.append(field <= op_value)
+                                elif op == "$ne":
+                                    conditions.append(field != op_value)
+                                elif op == "$in":
+                                    conditions.append(field.in_(op_value))
+                                elif op == "$nin":
+                                    conditions.append(field.not_in(op_value))
+                                else:
+                                    logger.warning(f"过滤器中遇到未知操作符 '{op}' (字段: '{key}')。将跳过此操作符。")
+                        else:
+                            # 直接相等比较
+                            conditions.append(field == value)
                     else:
-                        # 直接相等比较
-                        conditions.append(field == value)
-                else:
-                    logger.warning(f"过滤器键 '{key}' 在 Messages 模型中未找到。将跳过此条件。")
-            if conditions:
-                query = query.where(*conditions)
+                        logger.warning(f"过滤器键 '{key}' 在 Messages 模型中未找到。将跳过此条件。")
+                if conditions:
+                    query = query.where(*conditions)
 
-        if filter_bot:
-            query = query.where(Messages.user_id != global_config.bot.qq_account)
+            if filter_bot:
+                query = query.where(Messages.user_id != global_config.bot.qq_account)
 
-        if filter_command:
-            query = query.where(not_(Messages.is_command))
+            if filter_command:
+                query = query.where(not_(Messages.is_command))
 
-        if limit > 0:
-            # 确保limit是正整数
-            limit = max(1, int(limit))
+            if limit > 0:
+                # 确保limit是正整数
+                limit = max(1, int(limit))
 
-            if limit_mode == "earliest":
-                # 获取时间最早的 limit 条记录，已经是正序
-                query = query.order_by(Messages.time.asc()).limit(limit)
+                if limit_mode == "earliest":
+                    # 获取时间最早的 limit 条记录，已经是正序
+                    query = query.order_by(Messages.time.asc()).limit(limit)
+                    try:
+                        results = session.execute(query).scalars().all()
+                    except Exception as e:
+                        logger.error(f"执行earliest查询失败: {e}")
+                        results = []
+                else:  # 默认为 'latest'
+                    # 获取时间最晚的 limit 条记录
+                    query = query.order_by(Messages.time.desc()).limit(limit)
+                    try:
+                        latest_results = session.execute(query).scalars().all()
+                        # 将结果按时间正序排列
+                        results = sorted(latest_results, key=lambda msg: msg.time)
+                    except Exception as e:
+                        logger.error(f"执行latest查询失败: {e}")
+                        results = []
+            else:
+                # limit 为 0 时，应用传入的 sort 参数
+                if sort:
+                    sort_terms = []
+                    for field_name, direction in sort:
+                        if hasattr(Messages, field_name):
+                            field = getattr(Messages, field_name)
+                            if direction == 1:  # ASC
+                                sort_terms.append(field.asc())
+                            elif direction == -1:  # DESC
+                                sort_terms.append(field.desc())
+                            else:
+                                logger.warning(f"字段 '{field_name}' 的排序方向 '{direction}' 无效。将跳过此排序条件。")
+                        else:
+                            logger.warning(f"排序字段 '{field_name}' 在 Messages 模型中未找到。将跳过此排序条件。")
+                    if sort_terms:
+                        query = query.order_by(*sort_terms)
                 try:
                     results = session.execute(query).scalars().all()
                 except Exception as e:
-                    logger.error(f"执行earliest查询失败: {e}")
+                    logger.error(f"执行无限制查询失败: {e}")
                     results = []
-            else:  # 默认为 'latest'
-                # 获取时间最晚的 limit 条记录
-                query = query.order_by(Messages.time.desc()).limit(limit)
-                try:
-                    latest_results = session.execute(query).scalars().all()
-                    # 将结果按时间正序排列
-                    results = sorted(latest_results, key=lambda msg: msg.time)
-                except Exception as e:
-                    logger.error(f"执行latest查询失败: {e}")
-                    results = []
-        else:
-            # limit 为 0 时，应用传入的 sort 参数
-            if sort:
-                sort_terms = []
-                for field_name, direction in sort:
-                    if hasattr(Messages, field_name):
-                        field = getattr(Messages, field_name)
-                        if direction == 1:  # ASC
-                            sort_terms.append(field.asc())
-                        elif direction == -1:  # DESC
-                            sort_terms.append(field.desc())
-                        else:
-                            logger.warning(f"字段 '{field_name}' 的排序方向 '{direction}' 无效。将跳过此排序条件。")
-                    else:
-                        logger.warning(f"排序字段 '{field_name}' 在 Messages 模型中未找到。将跳过此排序条件。")
-                if sort_terms:
-                    query = query.order_by(*sort_terms)
-            try:
-                results = session.execute(query).scalars().all()
-            except Exception as e:
-                logger.error(f"执行无限制查询失败: {e}")
-                results = []
 
         return [_model_to_dict(msg) for msg in results]
     except Exception as e:
@@ -152,50 +152,50 @@ def count_messages(message_filter: dict[str, Any]) -> int:
         符合条件的消息数量，如果出错则返回 0。
     """
     try:
-        session = get_session()
-        query = select(func.count(Messages.id))
+        with get_db_session() as session:
+            query = select(func.count(Messages.id))
 
-        # 应用过滤器
-        if message_filter:
-            conditions = []
-            for key, value in message_filter.items():
-                if hasattr(Messages, key):
-                    field = getattr(Messages, key)
-                    if isinstance(value, dict):
-                        # 处理 MongoDB 风格的操作符
-                        for op, op_value in value.items():
-                            if op == "$gt":
-                                conditions.append(field > op_value)
-                            elif op == "$lt":
-                                conditions.append(field < op_value)
-                            elif op == "$gte":
-                                conditions.append(field >= op_value)
-                            elif op == "$lte":
-                                conditions.append(field <= op_value)
-                            elif op == "$ne":
-                                conditions.append(field != op_value)
-                            elif op == "$in":
-                                conditions.append(field.in_(op_value))
-                            elif op == "$nin":
-                                conditions.append(field.not_in(op_value))
-                            else:
-                                logger.warning(
-                                    f"计数时，过滤器中遇到未知操作符 '{op}' (字段: '{key}')。将跳过此操作符。"
-                                )
+            # 应用过滤器
+            if message_filter:
+                conditions = []
+                for key, value in message_filter.items():
+                    if hasattr(Messages, key):
+                        field = getattr(Messages, key)
+                        if isinstance(value, dict):
+                            # 处理 MongoDB 风格的操作符
+                            for op, op_value in value.items():
+                                if op == "$gt":
+                                    conditions.append(field > op_value)
+                                elif op == "$lt":
+                                    conditions.append(field < op_value)
+                                elif op == "$gte":
+                                    conditions.append(field >= op_value)
+                                elif op == "$lte":
+                                    conditions.append(field <= op_value)
+                                elif op == "$ne":
+                                    conditions.append(field != op_value)
+                                elif op == "$in":
+                                    conditions.append(field.in_(op_value))
+                                elif op == "$nin":
+                                    conditions.append(field.not_in(op_value))
+                                else:
+                                    logger.warning(
+                                        f"计数时，过滤器中遇到未知操作符 '{op}' (字段: '{key}')。将跳过此操作符。"
+                                    )
+                        else:
+                            # 直接相等比较
+                            conditions.append(field == value)
                     else:
-                        # 直接相等比较
-                        conditions.append(field == value)
-                else:
-                    logger.warning(f"计数时，过滤器键 '{key}' 在 Messages 模型中未找到。将跳过此条件。")
-            if conditions:
-                query = query.where(*conditions)
+                        logger.warning(f"计数时，过滤器键 '{key}' 在 Messages 模型中未找到。将跳过此条件。")
+                if conditions:
+                    query = query.where(*conditions)
 
-        count = session.execute(query).scalar()
-        return count or 0
+            count = session.execute(query).scalar()
+            return count or 0
     except Exception as e:
-        log_message = f"使用 SQLAlchemy 计数消息失败 (message_filter={message_filter}): {e}\n{traceback.format_exc()}"
-        logger.error(log_message)
-        return 0
+            log_message = f"使用 SQLAlchemy 计数消息失败 (message_filter={message_filter}): {e}\n{traceback.format_exc()}"
+            logger.error(log_message)
+            return 0
 
 
 # 你可以在这里添加更多与 messages 集合相关的数据库操作函数，例如 find_one_message, insert_message 等。