diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index acd7522e..11720a78 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/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 { diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py index 738dd052..8f27432c 100644 --- a/backend/apps/db/db.py +++ b/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: