444 lines
17 KiB
Python
444 lines
17 KiB
Python
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 |