jiateng_ws/utils/sql_utils.py

import sys
import logging
import threading
import time
from utils.config_loader import ConfigLoader

try:
    import psycopg2
except ImportError:
    psycopg2 = None

try:
    import sqlite3
except ImportError:
    sqlite3 = None

try:
    import mysql.connector
except ImportError:
    mysql = None


class SQLUtils:
    # 存储连接池，使用线程ID作为键的一部分
    _connection_pool = {}
    # 连接引用计数
    _connection_refs = {}
    # 最后使用时间记录
    _last_used = {}
    # 轻量级锁，仅用于连接池访问
    _pool_lock = threading.RLock()
    # 空闲连接超时时间（秒）
    _idle_timeout = 300  # 5分钟
    # 初始化清理线程标志
    _cleanup_thread_started = False
    
    @classmethod
    def _start_cleanup_thread(cls):
        """启动清理空闲连接的后台线程"""
        if cls._cleanup_thread_started:
            return
            
        def cleanup_idle_connections():
            """定期清理空闲连接的线程函数"""
            logging.info("数据库连接清理线程已启动")
            while True:
                time.sleep(60)  # 每分钟检查一次
                try:
                    current_time = time.time()
                    
                    with cls._pool_lock:
                        # 复制键列表，避免在迭代过程中修改字典
                        conn_keys = list(cls._connection_pool.keys())
                        
                        for conn_key in conn_keys:
                            # 检查引用计数和最后使用时间
                            if (conn_key in cls._connection_refs and 
                                cls._connection_refs[conn_key] <= 0 and
                                conn_key in cls._last_used and
                                current_time - cls._last_used[conn_key] > cls._idle_timeout):
                                
                                try:
                                    # 获取连接和游标
                                    conn, cursor = cls._connection_pool[conn_key]
                                    
                                    # 关闭资源
                                    if cursor:
                                        cursor.close()
                                    if conn:
                                        conn.close()
                                        
                                    # 从所有集合中移除
                                    cls._connection_pool.pop(conn_key, None)
                                    cls._connection_refs.pop(conn_key, None)
                                    cls._last_used.pop(conn_key, None)
                                    
                                    logging.debug(f"已清理空闲连接: {conn_key}")
                                except Exception as e:
                                    logging.error(f"清理空闲连接时出错: {e}")
                except Exception as e:
                    logging.error(f"连接清理线程执行异常: {e}")
        
        # 创建并启动后台线程
        cleanup_thread = threading.Thread(
            target=cleanup_idle_connections,
            daemon=True,
            name="DB-Connection-Cleanup"
        )
        cleanup_thread.start()
        cls._cleanup_thread_started = True
    
    def __init__(self, db_type=None, source_name=None, **kwargs):
        """初始化SQLUtils对象
        
        Args:
            db_type: 数据库类型 'sqlite', 'postgresql', 'mysql' 之一，如果为None则使用配置文件中的默认数据源
            source_name: 数据源名称，用于从配置中获取特定的数据源，如'sqlite', 'postgresql', 'mysql'
            **kwargs: 连接参数，如果没有提供，则使用配置文件中的参数
        """
        # 确保清理线程已启动
        if not SQLUtils._cleanup_thread_started:
            SQLUtils._start_cleanup_thread()
            
        self.conn = None
        self.cursor = None
        
        # 如果指定了source_name，直接使用该名称的数据源配置
        if source_name:
            config_loader = ConfigLoader.get_instance()
            db_config = config_loader.get_value(f'database.sources.{source_name}', {})
            if not db_config:
                raise ValueError(f"未找到数据源配置: {source_name}")
            
            db_type = source_name  # 数据源名称同时也是类型
            
            if not kwargs:  # 如果没有提供连接参数，则使用配置中的参数
                if source_name == 'sqlite':
                    kwargs = {'database': db_config.get('path', 'db/jtDB.db')}
                else:
                    kwargs = {
                        'host': db_config.get('host', 'localhost'),
                        'user': db_config.get('user', ''),
                        'password': db_config.get('password', ''),
                        'database': db_config.get('name', 'jtDB')
                    }
                    if 'port' in db_config and db_config['port']:
                        kwargs['port'] = int(db_config['port'])
                        
        # 如果没有指定数据库类型和数据源名称，则使用配置中的默认数据源
        elif db_type is None:
            config_loader = ConfigLoader.get_instance()
            default_source = config_loader.get_value('database.default', 'sqlite')
            
            # 如果没有提供连接参数，则从配置文件获取
            if not kwargs:
                db_config = config_loader.get_database_config(default_source)
                if default_source == 'sqlite':
                    kwargs = {'database': db_config.get('path', 'db/jtDB.db')}
                else:
                    kwargs = {
                        'host': db_config.get('host', 'localhost'),
                        'user': db_config.get('user', ''),
                        'password': db_config.get('password', ''),
                        'database': db_config.get('name', 'jtDB')
                    }
                    if 'port' in db_config and db_config['port']:
                        kwargs['port'] = int(db_config['port'])
            
            db_type = default_source
        
        self.db_type = db_type.lower()
        self.kwargs = kwargs
        self.source_name = source_name or self.db_type
        
        # 尝试从连接池获取连接，如果没有则创建新连接
        self._get_connection()
    
    def __enter__(self):
        """上下文管理器入口方法，支持with语句"""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器退出方法，自动关闭游标"""
        self.close()
        return False  # 不抑制异常

    def _get_connection(self):
        """从连接池获取连接，基于线程ID"""
        # 使用线程ID作为连接键的一部分
        thread_id = threading.get_ident()
        conn_key = f"{self.db_type}:{str(self.kwargs)}:{thread_id}"
        
        # 检查连接池中是否已有此线程的连接
        # 只在访问共享资源时使用锁
        with self._pool_lock:
            if conn_key in self._connection_pool:
                try:
                    # 尝试执行简单查询，确认连接有效
                    conn, cursor = self._connection_pool[conn_key]
                    cursor.execute("SELECT 1")
                    # 连接有效，直接使用
                    self.conn = conn
                    self.cursor = cursor
                    # 更新引用计数和最后使用时间
                    self._connection_refs[conn_key] = self._connection_refs.get(conn_key, 0) + 1
                    self._last_used[conn_key] = time.time()
                    return
                except Exception:
                    # 连接已失效，从连接池移除
                    self._cleanup_connection(conn_key)
            
        # 创建新连接 - 这部分不需要锁
        self.connect()
        
        # 将新连接添加到连接池 - 需要锁
        if self.conn and self.cursor:
            with self._pool_lock:
                self._connection_pool[conn_key] = (self.conn, self.cursor)
                self._connection_refs[conn_key] = 1
                self._last_used[conn_key] = time.time()
    
    def _cleanup_connection(self, conn_key):
        """清理指定的连接"""
        try:
            if conn_key in self._connection_pool:
                conn, cursor = self._connection_pool[conn_key]
                if cursor:
                    try:
                        cursor.close()
                    except:
                        pass
                if conn:
                    try:
                        conn.close()
                    except:
                        pass
                # 移除相关引用
                self._connection_pool.pop(conn_key, None)
                self._connection_refs.pop(conn_key, None)
                self._last_used.pop(conn_key, None)
                logging.debug(f"已清理连接: {conn_key}")
        except Exception as e:
            logging.error(f"清理连接失败: {e}")

    def connect(self):
        """连接到数据库 - 不需要全局锁"""
        try:
            if self.db_type in ['pgsql', 'postgresql']:
                if not psycopg2:
                    raise ImportError('psycopg2 is not installed')
                self.conn = psycopg2.connect(**self.kwargs)
            elif self.db_type in ['sqlite', 'sqlite3']:
                if not sqlite3:
                    raise ImportError('sqlite3 is not installed')
                self.conn = sqlite3.connect(self.kwargs.get('database', ':memory:'), check_same_thread=False)
            elif self.db_type == 'mysql':
                if not mysql:
                    raise ImportError('mysql.connector is not installed')
                self.conn = mysql.connector.connect(**self.kwargs)
            else:
                raise ValueError(f'不支持的数据库类型: {self.db_type}')
            
            self.cursor = self.conn.cursor()
            logging.debug(f"成功连接到数据库: {self.db_type}")
        except Exception as e:
            logging.error(f"连接数据库失败: {e}")
            raise

    def execute_query(self, sql, params=None):
        """执行查询 - 不使用全局锁，仅使用单个连接"""
        if params is None:
            params = ()
        
        try:
            # 直接执行查询，因为每个线程有自己的连接
            self.cursor.execute(sql, params)
            # 更新最后使用时间
            thread_id = threading.get_ident()
            conn_key = f"{self.db_type}:{str(self.kwargs)}:{thread_id}"
            with self._pool_lock:
                if conn_key in self._last_used:
                    self._last_used[conn_key] = time.time()
            return self.cursor
        except Exception as e:
            logging.error(f"执行查询失败: {e}, SQL: {sql}, 参数: {params}")
            raise

    def execute_update(self, sql, params=None):
        """执行更新 - 不使用全局锁，仅使用单个连接"""
        try:
            if params is None:
                params = ()
            
            # 直接执行更新
            self.cursor.execute(sql, params)
            self.conn.commit()
            
            # 更新最后使用时间
            thread_id = threading.get_ident()
            conn_key = f"{self.db_type}:{str(self.kwargs)}:{thread_id}"
            with self._pool_lock:
                if conn_key in self._last_used:
                    self._last_used[conn_key] = time.time()
                
            return self.cursor.rowcount
        except Exception as e:
            self.conn.rollback()
            logging.error(f"执行更新失败: {e}, SQL: {sql}, 参数: {params}")
            raise e
    
    def begin_transaction(self) -> None:
        """开始事务 - 不使用全局锁"""
        if self.db_type in ['sqlite', 'sqlite3']:
            self.execute_query('BEGIN TRANSACTION')
        else:
            self.conn.autocommit = False
    
    def commit_transaction(self) -> None:
        """提交事务 - 不使用全局锁"""
        self.conn.commit()
        if self.db_type not in ['sqlite', 'sqlite3']:
            self.conn.autocommit = True
    
    def rollback_transaction(self) -> None:
        """回滚事务 - 不使用全局锁"""
        self.conn.rollback()
        if self.db_type not in ['sqlite', 'sqlite3']:
            self.conn.autocommit = True

    def fetchone(self):
        """获取一行数据 - 不使用全局锁"""
        return self.cursor.fetchone()

    def fetchall(self):
        """获取所有数据 - 不使用全局锁"""
        return self.cursor.fetchall()

    def get_new_cursor(self):
        """获取一个新的游标，用于避免游标递归使用问题
        
        Returns:
            cursor: 数据库游标对象
        """
        if self.conn:
            return self.conn.cursor()
        return None

    def close(self):
        """关闭当前游标，减少引用计数，必要时释放连接"""
        thread_id = threading.get_ident()
        conn_key = f"{self.db_type}:{str(self.kwargs)}:{thread_id}"
        
        with self._pool_lock:
            if conn_key in self._connection_refs:
                # 减少引用计数
                self._connection_refs[conn_key] -= 1
                
                # 如果引用计数为0，关闭连接并从池中移除
                if self._connection_refs[conn_key] <= 0:
                    self._cleanup_connection(conn_key)
    
    def real_close(self):
        """强制关闭连接，无论引用计数"""
        thread_id = threading.get_ident()
        conn_key = f"{self.db_type}:{str(self.kwargs)}:{thread_id}"
        
        with self._pool_lock:
            self._cleanup_connection(conn_key)
    
    @staticmethod
    def close_all_connections():
        """关闭所有连接池中的连接"""
        with SQLUtils._pool_lock:
            conn_keys = list(SQLUtils._connection_pool.keys())
            for conn_key in conn_keys:
                try:
                    conn, cursor = SQLUtils._connection_pool[conn_key]
                    if cursor:
                        cursor.close()
                    if conn:
                        conn.close()
                except Exception as e:
                    logging.error(f"关闭数据库连接失败: {e}")
            
            # 清空所有字典
            SQLUtils._connection_pool.clear()
            SQLUtils._connection_refs.clear()
            SQLUtils._last_used.clear()
            logging.info("已关闭所有数据库连接")
    
    @staticmethod
    def get_sqlite_connection():
        """获取SQLite连接"""
        return SQLUtils(source_name='sqlite')
    
    @staticmethod
    def get_postgresql_connection():
        """获取PostgreSQL连接"""
        return SQLUtils(source_name='postgresql')
    
    @staticmethod
    def get_mysql_connection():
        """获取MySQL连接"""
        return SQLUtils(source_name='mysql')
        
    @classmethod
    def get_connection_pool_stats(cls):
        """获取连接池统计信息
        
        Returns:
            dict: 包含连接池统计信息的字典
        """
        with cls._pool_lock:
            stats = {
                'active_connections': len(cls._connection_pool),
                'connection_details': [],
                'connection_count_by_type': {},
                'active_threads': {},
            }
            
            # 统计不同类型连接数量
            for conn_key in cls._connection_pool:
                parts = conn_key.split(':')
                if len(parts) > 0:
                    db_type = parts[0]
                    stats['connection_count_by_type'][db_type] = stats['connection_count_by_type'].get(db_type, 0) + 1
                
                # 获取线程ID
                if len(parts) > 2:
                    thread_id = parts[2]
                    stats['active_threads'][thread_id] = stats['active_threads'].get(thread_id, 0) + 1
                
                # 连接详情
                refs = cls._connection_refs.get(conn_key, 0)
                last_used = cls._last_used.get(conn_key, 0)
                idle_time = time.time() - last_used if last_used else 0
                
                stats['connection_details'].append({
                    'key': conn_key,
                    'references': refs,
                    'idle_time_seconds': int(idle_time),
                    'is_idle': refs <= 0
                })
            
            return stats
    
    @classmethod
    def log_connection_pool_status(cls):
        """记录当前连接池状态到日志"""
        stats = cls.get_connection_pool_stats()
        logging.info(f"数据库连接池状态: 活动连接数={stats['active_connections']}")
        
        # 记录每种数据库类型的连接数
        for db_type, count in stats['connection_count_by_type'].items():
            logging.info(f"  - {db_type}: {count}个连接")
        
        # 记录空闲连接
        idle_connections = [d for d in stats['connection_details'] if d['is_idle']]
        if idle_connections:
            logging.info(f"  - 空闲连接: {len(idle_connections)}个")
            for conn in idle_connections:
                logging.debug(f"    * {conn['key']} (空闲{conn['idle_time_seconds']}秒)")
        
        return stats