from contextlib import asynccontextmanager import os import asyncio import sqlparse from pathlib import Path from typing import List, Dict, Any, Optional, Union, TYPE_CHECKING import aiosqlite if TYPE_CHECKING: from . import DatabaseManager # 全局数据库管理器实例 _global_db_manager: Optional["DatabaseManager"] = None def get_global_db_manager() -> "DatabaseManager": """获取全局数据库管理器实例""" global _global_db_manager if _global_db_manager is None: from . import DatabaseManager _global_db_manager = DatabaseManager() return _global_db_manager def close_global_db_manager() -> None: """关闭全局数据库管理器实例""" global _global_db_manager if _global_db_manager is not None: # 注意：这个函数应该在async环境中调用close_all_connections _global_db_manager = None class DatabaseManager: """异步数据库管理器""" def __init__(self, db_path: Optional[Union[str, Path]] = None, pool_size: int = 5): """ 初始化数据库管理器 Args: db_path: 数据库文件路径，支持str和Path类型 pool_size: 连接池大小 """ if db_path is None: self.db_path = os.environ.get("DATABASE_PATH", "./data/database.db") else: self.db_path = str(db_path) if isinstance(db_path, Path) else db_path # 连接池 self._connection_pool = [] self._pool_size = pool_size self._lock = asyncio.Lock() self._in_use = set() # 跟踪正在使用的连接 async def _get_connection(self) -> aiosqlite.Connection: """从连接池获取连接""" async with self._lock: # 尝试从池中获取现有连接 while self._connection_pool: conn = self._connection_pool.pop() # 检查连接是否仍然有效 try: await conn.execute("SELECT 1") self._in_use.add(conn) return conn except: # 连接已失效，关闭它 try: await conn.close() except: pass # 如果连接池为空，创建新连接 conn = await aiosqlite.connect(self.db_path) await conn.execute("PRAGMA foreign_keys = ON") self._in_use.add(conn) return conn async def _return_connection(self, conn: aiosqlite.Connection) -> None: """将连接返回到连接池""" async with self._lock: self._in_use.discard(conn) if len(self._connection_pool) < self._pool_size: self._connection_pool.append(conn) else: # 池已满，直接关闭连接 try: await conn.close() except: pass @asynccontextmanager async def get_conn(self): conn = await self._get_connection() yield conn await self._return_connection(conn) async def query( self, query: str, params: Optional[tuple] = None ) -> List[Dict[str, Any]]: """执行查询语句并返回结果""" conn = await self._get_connection() try: cursor = await conn.execute(query, params or ()) columns = [description[0] for description in cursor.description] rows = await cursor.fetchall() results = [dict(zip(columns, row)) for row in rows] await cursor.close() return results except Exception as e: # 记录错误但重新抛出，让调用者处理 raise Exception(f"数据库查询失败: {str(e)}") from e finally: await self._return_connection(conn) async def query_by_sql_file( self, file_path: Union[str, Path], params: Optional[tuple] = None ) -> List[Dict[str, Any]]: """从 SQL 文件中读取查询语句并执行""" path = str(file_path) if isinstance(file_path, Path) else file_path with open(path, "r", encoding="utf-8") as f: query = f.read() return await self.query(query, params) async def execute(self, command: str, params: Optional[tuple] = None) -> None: """执行非查询语句""" conn = await self._get_connection() try: await conn.execute(command, params or ()) await conn.commit() except Exception as e: await conn.rollback() raise Exception(f"数据库执行失败: {str(e)}") from e finally: await self._return_connection(conn) async def execute_script(self, script: str) -> None: """执行SQL脚本""" conn = await self._get_connection() try: await conn.executescript(script) await conn.commit() except Exception as e: await conn.rollback() raise Exception(f"数据库脚本执行失败: {str(e)}") from e finally: await self._return_connection(conn) def _parse_sql_statements(self, script: str) -> List[str]: """解析SQL脚本，分割成独立的语句""" # 使用sqlparse库更准确地分割SQL语句 parsed = sqlparse.split(script) statements = [] for statement in parsed: statement = statement.strip() if statement: statements.append(statement) return statements async def execute_by_sql_file( self, file_path: Union[str, Path], params: Optional[Union[tuple, List[tuple]]] = None, ) -> None: """从 SQL 文件中读取非查询语句并执行""" path = str(file_path) if isinstance(file_path, Path) else file_path with open(path, "r", encoding="utf-8") as f: script = f.read() # 如果有参数且是元组，使用execute执行整个脚本 if params is not None and isinstance(params, tuple): await self.execute(script, params) # 如果有参数且是列表，分别执行每个语句 elif params is not None and isinstance(params, list): # 使用sqlparse准确分割SQL语句 statements = self._parse_sql_statements(script) if len(statements) != len(params): raise ValueError( f"语句数量({len(statements)})与参数组数量({len(params)})不匹配" ) for statement, stmt_params in zip(statements, params): if statement: await self.execute(statement, stmt_params) # 如果无参数，使用executescript else: await self.execute_script(script) async def execute_many(self, command: str, seq_of_params: List[tuple]) -> None: """执行多条非查询语句""" conn = await self._get_connection() try: await conn.executemany(command, seq_of_params) await conn.commit() except Exception as e: await conn.rollback() raise Exception(f"数据库批量执行失败: {str(e)}") from e finally: await self._return_connection(conn) async def execute_many_values_by_sql_file( self, file_path: Union[str, Path], seq_of_params: List[tuple] ) -> None: """从 SQL 文件中读取一条语句，但是被不同值同时执行""" path = str(file_path) if isinstance(file_path, Path) else file_path with open(path, "r", encoding="utf-8") as f: command = f.read() await self.execute_many(command, seq_of_params) async def close_all_connections(self) -> None: """关闭所有连接""" async with self._lock: # 关闭池中的连接 for conn in self._connection_pool: try: await conn.close() except: pass self._connection_pool.clear() # 关闭正在使用的连接 for conn in self._in_use.copy(): try: await conn.close() except: pass self._in_use.clear()