refactor: Optimize connection pool #1256

Author: XiaJunjie2020

backend/apps/datasource/crud/datasource.py

@@ -24,6 +24,7 @@
 from ..crud.table import delete_table_by_ds_id, update_table
 from ..models.datasource import CoreDatasource, CreateDatasource, CoreTable, CoreField, ColumnSchema, TableObj, \
     DatasourceConf, TableAndFields
+from apps.db.db import pool_manager, driver_pool_manager
 
 
 def get_datasource_list(session: SessionDep, user: CurrentUser, oid: Optional[int] = None) -> List[CoreDatasource]:
@@ -109,6 +110,10 @@ def update_ds(session: SessionDep, trans: Trans, user: CurrentUser, ds: CoreData
     session.add(record)
     session.commit()
 
+    # update pool
+    pool_manager.remove_pool(ds.id)
+    driver_pool_manager.remove_pool(ds.id)
+
     run_save_ds_embeddings([ds.id])
     return ds
 
@@ -135,6 +140,11 @@ async def delete_ds(session: SessionDep, id: int):
     session.commit()
     delete_table_by_ds_id(session, id)
     delete_field_by_ds_id(session, id)
+
+    # update pool
+    pool_manager.remove_pool(id)
+    driver_pool_manager.remove_pool(id)
+
     if term:
         await clear_ws_ds_cache(term.oid)
     return {

backend/apps/db/db.py

@@ -237,7 +237,8 @@ def get_driver_connection(ds: CoreDatasource | AssistantOutDsSchema, db_config:
             )
     elif equals_ignore_case(ds.type, 'redshift'):
         if not use_pool:
-            conn = redshift_connector.connect(host=conf.host, port=conf.port, database=conf.database, user=conf.username,
+            conn = redshift_connector.connect(host=conf.host, port=conf.port, database=conf.database,
+                                              user=conf.username,
                                               password=conf.password,
                                               timeout=conf.timeout, **conn_conf)
         else:
@@ -1175,6 +1176,17 @@ def get_pool(self, ds: CoreDatasource | AssistantOutDsSchema, **db_config):
             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)
+                # 2. 安全关闭该连接池，释放底层所有数据库连接和内存
+                pool.close()
+                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:
@@ -1220,6 +1232,17 @@ def get_pool(self, ds: CoreDatasource | AssistantOutDsSchema, db_config):
             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)
+                # 2. 安全关闭该连接池，释放底层所有数据库连接和内存
+                pool.close()
+                print(f"[Manager] Closed pool and remove: {datasource_id}")
+            else:
+                print(f"[Manager] Warning: ds id {datasource_id} not exist in dbutils")
+
     def close_all(self):
         """stop"""
         with self._lock: