diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py index 2226f9cb..842eec47 100644 --- a/backend/apps/db/db.py +++ b/backend/apps/db/db.py @@ -1156,7 +1156,7 @@ def __init__(self, max_pools=500): :param max_pools: max pool """ self.max_pools = max_pools - self._pools = OrderedDict() # 使用有序字典实现 LRU + self._pools = OrderedDict() # datasource_id -> (sessionmaker, Engine) self._lock = threading.Lock() # 保证多线程安全 def get_pool(self, ds: CoreDatasource | AssistantOutDsSchema, **db_config): @@ -1169,18 +1169,18 @@ def get_pool(self, ds: CoreDatasource | AssistantOutDsSchema, **db_config): if ds.id in self._pools: self._pools.move_to_end(ds.id) print(f"[LRU] return: {ds.id}") - return self._pools[ds.id] + return self._pools[ds.id][0] # 2. 如果连接池不存在,检查是否达到上限,若达到则淘汰最久未使用的(字典头部) if len(self._pools) >= self.max_pools: - oldest_id, oldest_pool = self._pools.popitem(last=False) - oldest_pool.close() # 安全关闭被驱逐的连接池 + oldest_id, (_, oldest_engine) = self._pools.popitem(last=False) + oldest_engine.dispose() # 安全关闭被驱逐的连接池 print(f"[LRU] remove oldest: {oldest_id}") # 3. 创建新连接池并放入字典末尾 engine = get_engine(ds, use_pool=True) new_pool = sessionmaker(bind=engine) - self._pools[ds.id] = new_pool + self._pools[ds.id] = (new_pool, engine) print(f"[LRU] create: {ds.id}") return new_pool @@ -1188,9 +1188,9 @@ def remove_pool(self, datasource_id): with self._lock: if datasource_id in self._pools: # 1. 从字典中移除并获取该连接池对象 - pool = self._pools.pop(datasource_id) + _, engine = self._pools.pop(datasource_id) # 2. 安全关闭该连接池,释放底层所有数据库连接和内存 - pool.close() + engine.dispose() print(f"[Manager] Closed pool and remove: {datasource_id}") else: print(f"[Manager] Warning: ds id {datasource_id} not exist in sqlalchemy") @@ -1198,8 +1198,8 @@ def remove_pool(self, datasource_id): def close_all(self): """stop""" with self._lock: - for pool in self._pools.values(): - pool.close() + for _, engine in self._pools.values(): + engine.dispose() self._pools.clear() diff --git a/backend/tests/test_connection_pool_manager.py b/backend/tests/test_connection_pool_manager.py new file mode 100644 index 00000000..981298b5 --- /dev/null +++ b/backend/tests/test_connection_pool_manager.py @@ -0,0 +1,103 @@ +"""Regression tests for SQLAlchemy connection pool lifecycle management.""" + +import ast +import os +import threading +from collections import OrderedDict +from types import SimpleNamespace +from unittest.mock import Mock + +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker + + +_SRC_PATH = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "apps", + "db", + "db.py", +) + + +def _load_connection_pool_manager(): + """Load only ConnectionPoolManager to avoid importing database drivers.""" + with open(_SRC_PATH, encoding="utf-8") as file: + source = file.read() + + tree = ast.parse(source) + class_node = next( + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == "ConnectionPoolManager" + ) + + class CoreDatasource: + pass + + class AssistantOutDsSchema: + pass + + namespace = { + "threading": threading, + "OrderedDict": OrderedDict, + "CoreDatasource": CoreDatasource, + "AssistantOutDsSchema": AssistantOutDsSchema, + "sessionmaker": sessionmaker, + "get_engine": None, + } + module = ast.Module(body=[class_node], type_ignores=[]) + ast.fix_missing_locations(module) + exec(compile(module, _SRC_PATH, "exec"), namespace) + return namespace["ConnectionPoolManager"], namespace + + +ConnectionPoolManager, _namespace = _load_connection_pool_manager() + + +def _engine_factory(created_engines): + def get_engine(ds, use_pool=False): + assert use_pool is True + engine = create_engine("sqlite://") + engine.dispose = Mock(wraps=engine.dispose) + created_engines.append(engine) + return engine + + return get_engine + + +def test_remove_pool_disposes_bound_engine(): + created_engines = [] + _namespace["get_engine"] = _engine_factory(created_engines) + manager = ConnectionPoolManager() + + manager.get_pool(SimpleNamespace(id="ds-1")) + manager.remove_pool("ds-1") + + created_engines[0].dispose.assert_called_once_with() + assert "ds-1" not in manager._pools + + +def test_lru_eviction_disposes_oldest_engine(): + created_engines = [] + _namespace["get_engine"] = _engine_factory(created_engines) + manager = ConnectionPoolManager(max_pools=1) + + manager.get_pool(SimpleNamespace(id="ds-1")) + manager.get_pool(SimpleNamespace(id="ds-2")) + + created_engines[0].dispose.assert_called_once_with() + created_engines[1].dispose.assert_not_called() + + +def test_close_all_disposes_every_engine(): + created_engines = [] + _namespace["get_engine"] = _engine_factory(created_engines) + manager = ConnectionPoolManager(max_pools=2) + + manager.get_pool(SimpleNamespace(id="ds-1")) + manager.get_pool(SimpleNamespace(id="ds-2")) + manager.close_all() + + for engine in created_engines: + engine.dispose.assert_called_once_with() + assert not manager._pools