fix(datasource): dispose SQLAlchemy engine when removing pool

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,37 +1169,37 @@ 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
 
     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")
 
     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()
 
 

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