# P2P Network Communication - Database Module """ 数据库模块，提供MySQL数据库连接池管理和表初始化功能 使用aiomysql异步驱动实现高性能数据库操作 """ import asyncio import logging from contextlib import asynccontextmanager from typing import Optional, AsyncGenerator import aiomysql from config import ServerConfig, load_server_config logger = logging.getLogger(__name__) class DatabaseManager: """数据库管理器，负责连接池管理和数据库初始化""" def __init__(self, config: Optional[ServerConfig] = None): """ 初始化数据库管理器 Args: config: 服务器配置，如果为None则从环境变量加载 """ self.config = config or load_server_config() self._pool: Optional[aiomysql.Pool] = None self._initialized = False async def initialize(self) -> None: """初始化数据库连接池和表结构""" if self._initialized: return await self._create_pool() await self._create_tables() self._initialized = True logger.info("Database initialized successfully") async def _create_pool(self) -> None: """创建数据库连接池""" try: self._pool = await aiomysql.create_pool( host=self.config.db_host, port=self.config.db_port, user=self.config.db_user, password=self.config.db_password, db=self.config.db_name, minsize=1, maxsize=self.config.db_pool_size, charset='utf8mb4', autocommit=True, echo=False, ) logger.info(f"Database connection pool created: {self.config.db_host}:{self.config.db_port}") except Exception as e: logger.error(f"Failed to create database pool: {e}") raise async def _create_tables(self) -> None: """创建数据库表结构""" create_users_table = """ CREATE TABLE IF NOT EXISTS users ( user_id VARCHAR(64) PRIMARY KEY, username VARCHAR(100) UNIQUE NOT NULL, display_name VARCHAR(200) NOT NULL, password_hash VARCHAR(255) NOT NULL DEFAULT '', public_key BLOB, status VARCHAR(20) DEFAULT 'offline', ip_address VARCHAR(45) DEFAULT '', port INT DEFAULT 0, last_seen TIMESTAMP NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, INDEX idx_username (username), INDEX idx_status (status) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; """ create_messages_table = """ CREATE TABLE IF NOT EXISTS messages ( message_id VARCHAR(64) PRIMARY KEY, sender_id VARCHAR(64) NOT NULL, receiver_id VARCHAR(64) NOT NULL, content_type VARCHAR(50) NOT NULL, content TEXT, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP, is_read BOOLEAN DEFAULT FALSE, is_delivered BOOLEAN DEFAULT FALSE, INDEX idx_sender (sender_id), INDEX idx_receiver (receiver_id), INDEX idx_timestamp (timestamp), INDEX idx_conversation (sender_id, receiver_id, timestamp) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; """ create_file_transfers_table = """ CREATE TABLE IF NOT EXISTS file_transfers ( transfer_id VARCHAR(64) PRIMARY KEY, file_name VARCHAR(500) NOT NULL, file_size BIGINT NOT NULL, file_hash VARCHAR(128) NOT NULL, sender_id VARCHAR(64) NOT NULL, receiver_id VARCHAR(64) NOT NULL, status VARCHAR(20) NOT NULL DEFAULT 'pending', progress DECIMAL(5,2) DEFAULT 0, start_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP, end_time TIMESTAMP NULL, INDEX idx_sender (sender_id), INDEX idx_receiver (receiver_id), INDEX idx_status (status) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; """ create_offline_messages_table = """ CREATE TABLE IF NOT EXISTS offline_messages ( id BIGINT AUTO_INCREMENT PRIMARY KEY, user_id VARCHAR(64) NOT NULL, message_data BLOB NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, INDEX idx_user_id (user_id), INDEX idx_created_at (created_at) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; """ tables = [ ("users", create_users_table), ("messages", create_messages_table), ("file_transfers", create_file_transfers_table), ("offline_messages", create_offline_messages_table), ] async with self.acquire() as conn: async with conn.cursor() as cursor: for table_name, create_sql in tables: try: await cursor.execute(create_sql) logger.info(f"Table '{table_name}' created or already exists") except Exception as e: logger.error(f"Failed to create table '{table_name}': {e}") raise @asynccontextmanager async def acquire(self) -> AsyncGenerator[aiomysql.Connection, None]: """ 获取数据库连接的上下文管理器 Yields: 数据库连接对象 """ if not self._pool: raise RuntimeError("Database pool not initialized. Call initialize() first.") conn = await self._pool.acquire() try: yield conn finally: self._pool.release(conn) async def execute(self, query: str, args: tuple = ()) -> int: """ 执行SQL语句（INSERT, UPDATE, DELETE） Args: query: SQL语句 args: 参数元组 Returns: 受影响的行数 """ async with self.acquire() as conn: async with conn.cursor() as cursor: await cursor.execute(query, args) return cursor.rowcount async def fetch_one(self, query: str, args: tuple = ()) -> Optional[tuple]: """ 查询单条记录 Args: query: SQL语句 args: 参数元组 Returns: 查询结果元组，如果没有结果返回None """ async with self.acquire() as conn: async with conn.cursor() as cursor: await cursor.execute(query, args) return await cursor.fetchone() async def fetch_all(self, query: str, args: tuple = ()) -> list: """ 查询多条记录 Args: query: SQL语句 args: 参数元组 Returns: 查询结果列表 """ async with self.acquire() as conn: async with conn.cursor() as cursor: await cursor.execute(query, args) return await cursor.fetchall() async def fetch_dict(self, query: str, args: tuple = ()) -> Optional[dict]: """ 查询单条记录并返回字典 Args: query: SQL语句 args: 参数元组 Returns: 查询结果字典，如果没有结果返回None """ async with self.acquire() as conn: async with conn.cursor(aiomysql.DictCursor) as cursor: await cursor.execute(query, args) return await cursor.fetchone() async def fetch_all_dict(self, query: str, args: tuple = ()) -> list: """ 查询多条记录并返回字典列表 Args: query: SQL语句 args: 参数元组 Returns: 查询结果字典列表 """ async with self.acquire() as conn: async with conn.cursor(aiomysql.DictCursor) as cursor: await cursor.execute(query, args) return await cursor.fetchall() async def close(self) -> None: """关闭数据库连接池""" if self._pool: self._pool.close() await self._pool.wait_closed() self._pool = None self._initialized = False logger.info("Database connection pool closed") @property def is_initialized(self) -> bool: """检查数据库是否已初始化""" return self._initialized @property def pool(self) -> Optional[aiomysql.Pool]: """获取连接池对象""" return self._pool # 全局数据库管理器实例 _db_manager: Optional[DatabaseManager] = None async def get_database() -> DatabaseManager: """ 获取全局数据库管理器实例 Returns: DatabaseManager实例 """ global _db_manager if _db_manager is None: _db_manager = DatabaseManager() await _db_manager.initialize() return _db_manager async def close_database() -> None: """关闭全局数据库连接""" global _db_manager if _db_manager: await _db_manager.close() _db_manager = None