logan
4c9427e546
- P0: 实例访问控制全线接入权限组 (829测试) - P1: permissions字段激活 + 操作级权限检查 (843测试) - 前端: 删除数据权限页面, 权限组合并为唯一入口 - 新增: CodeMirror编辑器 + 查询模板 + 慢查询分析 - 新增: MongoDB管理页面 + Dashboard增强 - 修复: Pydantic V2迁移 + 前端TS错误修复
229 lines
7.4 KiB
Python
229 lines
7.4 KiB
Python
"""数据库连接池管理 — 复用数据库连接，提高性能。
|
|
|
|
使用异步连接池:
|
|
- MySQL: aiomysql.Pool
|
|
- PostgreSQL: asyncpg.Pool
|
|
|
|
连接池按实例 ID 管理，每个实例维护一个连接池。
|
|
"""
|
|
import asyncio
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Dict, Optional
|
|
|
|
from app.models.instance import Instance
|
|
from app.core.crypto import decrypt_password
|
|
|
|
|
|
@dataclass
|
|
class PoolStats:
|
|
"""连接池统计信息。"""
|
|
instance_id: int
|
|
db_type: str
|
|
active_connections: int = 0
|
|
idle_connections: int = 0
|
|
max_connections: int = 5
|
|
total_queries: int = 0
|
|
failed_queries: int = 0
|
|
avg_response_ms: int = 0
|
|
pool_hit_rate: float = 1.0
|
|
created_at: float = field(default_factory=time.time)
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"instance_id": self.instance_id,
|
|
"db_type": self.db_type,
|
|
"active_connections": self.active_connections,
|
|
"idle_connections": self.idle_connections,
|
|
"max_connections": self.max_connections,
|
|
"total_queries": self.total_queries,
|
|
"failed_queries": self.failed_queries,
|
|
"avg_response_ms": self.avg_response_ms,
|
|
"pool_hit_rate": round(self.pool_hit_rate, 4),
|
|
}
|
|
|
|
|
|
class ConnectionPoolManager:
|
|
"""数据库连接池管理器。"""
|
|
|
|
def __init__(self):
|
|
# 连接池缓存: {instance_id: pool}
|
|
self._pools: Dict[int, Any] = {}
|
|
# 连接池统计
|
|
self._stats: Dict[int, PoolStats] = {}
|
|
# 创建锁，防止并发创建重复连接池
|
|
self._locks: Dict[int, asyncio.Lock] = {}
|
|
# 总体锁
|
|
self._global_lock = asyncio.Lock()
|
|
|
|
async def get_pool(self, instance: Instance):
|
|
"""获取或创建连接池。"""
|
|
instance_id = instance.id
|
|
|
|
# 检查现有连接池
|
|
if instance_id in self._pools:
|
|
pool = self._pools[instance_id]
|
|
# 简单的健康检查
|
|
try:
|
|
if instance.db_type == "mysql":
|
|
pool._free # 检查 aiomysql pool 是否活跃
|
|
return pool
|
|
except Exception:
|
|
# 连接池已失效，重新创建
|
|
await self.close_pool(instance_id)
|
|
|
|
# 获取或创建实例锁
|
|
if instance_id not in self._locks:
|
|
self._locks[instance_id] = asyncio.Lock()
|
|
|
|
async with self._locks[instance_id]:
|
|
# 双重检查
|
|
if instance_id in self._pools:
|
|
return self._pools[instance_id]
|
|
|
|
pool = await self._create_pool(instance)
|
|
self._pools[instance_id] = pool
|
|
self._stats[instance_id] = PoolStats(
|
|
instance_id=instance_id,
|
|
db_type=instance.db_type,
|
|
max_connections=instance.connection_pool_size or 5,
|
|
)
|
|
return pool
|
|
|
|
async def _create_pool(self, instance: Instance):
|
|
"""创建连接池。"""
|
|
password = decrypt_password(instance.password)
|
|
pool_size = instance.connection_pool_size or 5
|
|
|
|
if instance.db_type == "mysql":
|
|
import aiomysql
|
|
return await aiomysql.create_pool(
|
|
host=instance.host,
|
|
port=instance.port,
|
|
user=instance.username,
|
|
password=password,
|
|
db=instance.database,
|
|
minsize=1,
|
|
maxsize=pool_size,
|
|
pool_recycle=3600, # 1小时回收连接
|
|
connect_timeout=10,
|
|
autocommit=True,
|
|
)
|
|
elif instance.db_type in ("postgres", "postgresql"):
|
|
import asyncpg
|
|
dsn = (
|
|
f"postgresql://{instance.username}:{password}"
|
|
f"@{instance.host}:{instance.port}"
|
|
f"/{instance.database or 'postgres'}"
|
|
)
|
|
return await asyncpg.create_pool(
|
|
dsn=dsn,
|
|
min_size=1,
|
|
max_size=pool_size,
|
|
command_timeout=60,
|
|
max_inactive_connection_lifetime=3600,
|
|
)
|
|
else:
|
|
raise ValueError(f"不支持的连接池类型: {instance.db_type}")
|
|
|
|
async def close_pool(self, instance_id: int) -> None:
|
|
"""关闭指定实例的连接池。"""
|
|
pool = self._pools.pop(instance_id, None)
|
|
if pool:
|
|
try:
|
|
pool.close()
|
|
await pool.wait_closed()
|
|
except Exception:
|
|
pass
|
|
self._stats.pop(instance_id, None)
|
|
|
|
async def close_all(self) -> None:
|
|
"""关闭所有连接池。"""
|
|
for instance_id in list(self._pools.keys()):
|
|
await self.close_pool(instance_id)
|
|
self._pools.clear()
|
|
self._stats.clear()
|
|
|
|
async def get_connection(self, instance: Instance):
|
|
"""从连接池获取一个连接（上下文管理器）。"""
|
|
pool = await self.get_pool(instance)
|
|
return pool.acquire()
|
|
|
|
def get_pool_stats(self, instance_id: int) -> Optional[PoolStats]:
|
|
"""获取连接池统计。"""
|
|
return self._stats.get(instance_id)
|
|
|
|
def get_all_stats(self) -> Dict[int, PoolStats]:
|
|
"""获取所有连接池统计。"""
|
|
return dict(self._stats)
|
|
|
|
async def health_check(self, instance_id: int) -> bool:
|
|
"""检查连接池健康状态。"""
|
|
pool = self._pools.get(instance_id)
|
|
if not pool:
|
|
return False
|
|
|
|
try:
|
|
if hasattr(pool, '_free'):
|
|
# aiomysql
|
|
conn = await pool.acquire()
|
|
try:
|
|
async with conn.cursor() as cur:
|
|
await cur.execute("SELECT 1")
|
|
await cur.fetchone()
|
|
return True
|
|
finally:
|
|
pool.release(conn)
|
|
else:
|
|
# asyncpg
|
|
conn = await pool.acquire()
|
|
try:
|
|
await conn.fetchval("SELECT 1")
|
|
return True
|
|
finally:
|
|
await pool.release(conn)
|
|
except Exception:
|
|
return False
|
|
|
|
async def reset_pool(self, instance_id: int) -> None:
|
|
"""重置连接池（关闭并重新创建）。"""
|
|
await self.close_pool(instance_id)
|
|
|
|
def update_stats(self, instance_id: int, query_success: bool = True,
|
|
response_ms: int = 0) -> None:
|
|
"""更新查询统计。"""
|
|
stats = self._stats.get(instance_id)
|
|
if not stats:
|
|
return
|
|
|
|
stats.total_queries += 1
|
|
if not query_success:
|
|
stats.failed_queries += 1
|
|
|
|
# 移动平均响应时间
|
|
if response_ms > 0:
|
|
if stats.avg_response_ms == 0:
|
|
stats.avg_response_ms = response_ms
|
|
else:
|
|
stats.avg_response_ms = (stats.avg_response_ms * 9 + response_ms) // 10
|
|
|
|
|
|
# 全局连接池管理器实例
|
|
_pool_manager: Optional[ConnectionPoolManager] = None
|
|
|
|
|
|
def get_pool_manager() -> ConnectionPoolManager:
|
|
"""获取全局连接池管理器。"""
|
|
global _pool_manager
|
|
if _pool_manager is None:
|
|
_pool_manager = ConnectionPoolManager()
|
|
return _pool_manager
|
|
|
|
|
|
async def close_all_pools():
|
|
"""关闭所有连接池（应用关闭时调用）。"""
|
|
global _pool_manager
|
|
if _pool_manager:
|
|
await _pool_manager.close_all()
|
|
_pool_manager = None
|