jiateng_ws/utils/sql_utils.py

import sys
import logging
import threading
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:
    # 存储连接池，避免重复创建连接
    _connection_pool = {}
    # 添加线程锁，用于防止多线程同时使用同一个cursor
    _lock = threading.RLock()
    
    def __init__(self, db_type=None, source_name=None, **kwargs):
        """初始化SQLUtils对象
        
        Args:
            db_type: 数据库类型 'sqlite', 'postgresql', 'mysql' 之一，如果为None则使用配置文件中的默认数据源
            source_name: 数据源名称，用于从配置中获取特定的数据源，如'sqlite', 'postgresql', 'mysql'
            **kwargs: 连接参数，如果没有提供，则使用配置文件中的参数
        """
        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):
        """从连接池获取连接，如果没有则创建新连接"""
        # 创建连接键，包含数据库类型和连接参数
        conn_key = f"{self.db_type}:{str(self.kwargs)}"
        
        # 检查连接池中是否已有此连接
        with SQLUtils._lock:
            if conn_key in SQLUtils._connection_pool:
                try:
                    # 尝试执行简单查询，确认连接有效
                    conn, cursor = SQLUtils._connection_pool[conn_key]
                    cursor.execute("SELECT 1")
                    # 连接有效，直接使用
                    self.conn = conn
                    self.cursor = cursor
                    return
                except Exception:
                    # 连接已失效，从连接池移除
                    del SQLUtils._connection_pool[conn_key]
            
            # 创建新连接
            self.connect()
            
            # 将新连接添加到连接池
            if self.conn and self.cursor:
                SQLUtils._connection_pool[conn_key] = (self.conn, self.cursor)

    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 = ()
        with SQLUtils._lock:
            try:
                self.cursor.execute(sql, params)
                return self.cursor
            except Exception as e:
                logging.error(f"执行查询失败: {e}, SQL: {sql}, 参数: {params}")
                raise

    def execute_update(self, sql, params=None):
        try:
            with SQLUtils._lock:
                self.cursor.execute(sql, params)
                self.conn.commit()
                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:
        """开始事务"""
        with SQLUtils._lock:
            if self.db_type in ['sqlite', 'sqlite3']:
                self.execute_query('BEGIN TRANSACTION')
            else:
                self.conn.autocommit = False
        
    def commit_transaction(self) -> None:
        """提交事务"""
        with SQLUtils._lock:
            self.conn.commit()
            if self.db_type not in ['sqlite', 'sqlite3']:
                self.conn.autocommit = True
        
    def rollback_transaction(self) -> None:
        """回滚事务"""
        with SQLUtils._lock:
            self.conn.rollback()
            if self.db_type not in ['sqlite', 'sqlite3']:
                self.conn.autocommit = True

    def fetchone(self):
        with SQLUtils._lock:
            return self.cursor.fetchone()

    def fetchall(self):
        with SQLUtils._lock:
            return self.cursor.fetchall()

    def close(self):
        """关闭当前游标，但保留连接在连接池中"""
        # 不再关闭连接，只关闭游标，减少重复创建连接的开销
        # 连接仍然保留在连接池中供后续使用
        pass
    
    def real_close(self):
        """真正关闭连接，从连接池中移除"""
        conn_key = f"{self.db_type}:{str(self.kwargs)}"
        with SQLUtils._lock:
            if conn_key in SQLUtils._connection_pool:
                try:
                    conn, cursor = SQLUtils._connection_pool[conn_key]
                    if cursor:
                        cursor.close()
                    if conn:
                        conn.close()
                    del SQLUtils._connection_pool[conn_key]
                    logging.debug(f"已关闭并移除连接: {conn_key}")
                except Exception as e:
                    logging.error(f"关闭连接失败: {e}")
    
    @staticmethod
    def close_all_connections():
        """关闭所有连接池中的连接"""
        with SQLUtils._lock:
            for conn, cursor in SQLUtils._connection_pool.values():
                try:
                    if cursor:
                        cursor.close()
                    if conn:
                        conn.close()
                except Exception as e:
                    logging.error(f"关闭数据库连接失败: {e}")
            
            SQLUtils._connection_pool.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')