import asyncio import json from contextlib import asynccontextmanager from typing import Any, Dict, List, Optional, Tuple, Union import aiomysql from loguru import logger # 使用条件导入替代直接导入 REDIS_AVAILABLE = False try: # 尝试导入redis.asyncio (Redis-py 4.2.0+) import redis.asyncio as redis_asyncio REDIS_AVAILABLE = True REDIS_TYPE = "redis-py" except ImportError: try: # 尝试导入aioredis (旧版本) import aioredis REDIS_AVAILABLE = True REDIS_TYPE = "aioredis" except (ImportError, TypeError): REDIS_AVAILABLE = False REDIS_TYPE = None from core.config import Config class DatabaseManager: def __init__(self, config: Config): # 数据库配置 self.db_config = config.database_config self._pool_size = self.db_config.pool_size self._pool = None # 连接池 self._pool_lock = asyncio.Lock() # Redis配置 self.use_redis = self.db_config.use_redis self.redis_config = config.redis_config if hasattr(config, 'redis_config') else None self.redis = None async def initialize(self): """初始化数据库连接池""" logger.info("初始化数据库连接池") # 创建MySQL连接池 try: logger.info(f"连接MySQL: {self.db_config.host}:{self.db_config.port}, 用户: {self.db_config.username}, 数据库: {self.db_config.database}") self._pool = await aiomysql.create_pool( host=self.db_config.host, port=self.db_config.port, user=self.db_config.username, password=self.db_config.password, db=self.db_config.database, maxsize=self._pool_size, autocommit=True, charset='utf8mb4' ) logger.info("MySQL连接池创建成功") except Exception as e: logger.error(f"MySQL连接池创建失败: {str(e)}") logger.error("请检查MySQL配置是否正确，以及MySQL服务是否已启动") logger.info(f"您可能需要创建MySQL用户和数据库：") logger.info(f" CREATE USER '{self.db_config.username}'@'localhost' IDENTIFIED BY '{self.db_config.password}';") logger.info(f" CREATE DATABASE {self.db_config.database} CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;") logger.info(f" GRANT ALL PRIVILEGES ON {self.db_config.database}.* TO '{self.db_config.username}'@'localhost';") logger.info(f" FLUSH PRIVILEGES;") raise # 初始化表结构 async with self.get_connection() as conn: async with conn.cursor() as cursor: await cursor.execute(''' CREATE TABLE IF NOT EXISTS email_accounts ( id INT AUTO_INCREMENT PRIMARY KEY, email VARCHAR(255) UNIQUE NOT NULL, password VARCHAR(255) NOT NULL, client_id VARCHAR(255) NOT NULL, refresh_token TEXT NOT NULL, in_use BOOLEAN DEFAULT 0, cursor_password VARCHAR(255), cursor_cookie TEXT, cursor_token TEXT, sold BOOLEAN DEFAULT 0, status VARCHAR(20) DEFAULT 'pending', extracted BOOLEAN DEFAULT 0, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, INDEX idx_status_inuse_sold (status, in_use, sold), INDEX idx_extracted (extracted, status, sold) ) ''') # 初始化Redis连接（如果配置了） if self.use_redis and REDIS_AVAILABLE and self.redis_config: try: # 根据检测到的Redis库类型创建连接 if REDIS_TYPE == "redis-py": # 使用redis.asyncio创建连接 logger.info(f"使用redis-py连接Redis: {self.redis_config.host}:{self.redis_config.port}") self.redis = redis_asyncio.Redis( host=self.redis_config.host, port=self.redis_config.port, password=self.redis_config.password or None, db=self.redis_config.db, decode_responses=True ) # 测试连接 await self.redis.ping() elif REDIS_TYPE == "aioredis": # 使用旧版aioredis创建连接 logger.info(f"使用aioredis连接Redis: {self.redis_config.host}:{self.redis_config.port}") self.redis = await aioredis.from_url( f"redis://{self.redis_config.host}:{self.redis_config.port}", password=self.redis_config.password or None, db=self.redis_config.db, encoding="utf-8", decode_responses=True ) logger.info("Redis连接初始化成功") except Exception as e: logger.error(f"Redis连接初始化失败: {e}") logger.info("Redis缓存将被禁用") self.redis = None logger.info(f"数据库连接池初始化完成，大小: {self._pool_size}") async def cleanup(self): """清理数据库连接""" if self._pool: self._pool.close() await self._pool.wait_closed() if self.redis: if REDIS_TYPE == "redis-py": await self.redis.close() else: await self.redis.close() logger.info("数据库连接已清理") @asynccontextmanager async def get_connection(self): """获取数据库连接""" if self._pool is None: raise Exception("数据库连接池未初始化") async with self._pool.acquire() as conn: try: yield conn finally: pass # 连接会自动返回池中 async def execute(self, query: str, params: tuple = ()) -> Any: """执行SQL语句""" logger.debug(f"执行SQL: {query}, 参数: {params}") try: async with self.get_connection() as conn: async with conn.cursor() as cursor: await cursor.execute(query, params) # 对于INSERT语句，返回最后插入的ID if query.strip().upper().startswith("INSERT"): return cursor.lastrowid # 对于UPDATE/DELETE语句，返回影响的行数 return cursor.rowcount except Exception as e: logger.error(f"SQL执行失败: {query}, 参数: {params}, 错误: {str(e)}") raise async def fetch_one(self, query: str, params: tuple = ()) -> Optional[Dict]: """查询单条记录""" logger.debug(f"查询单条: {query}, 参数: {params}") # 尝试从Redis获取缓存 cache_key = f"db:{self._make_cache_key(query, params)}" cached_result = await self._get_from_cache(cache_key) if cached_result is not None: return cached_result try: async with self.get_connection() as conn: async with conn.cursor(aiomysql.DictCursor) as cursor: await cursor.execute(query, params) result = await cursor.fetchone() # 缓存结果 if result and self.redis: await self._store_in_cache(cache_key, result) logger.debug(f"查询结果: {result}") return result except Exception as e: logger.error(f"查询单条失败: {query}, 参数: {params}, 错误: {str(e)}") raise async def fetch_all(self, query: str, params: tuple = ()) -> List[Dict]: """查询多条记录""" logger.debug(f"查询多条: {query}, 参数: {params}") # 尝试从Redis获取缓存 cache_key = f"db:{self._make_cache_key(query, params)}" cached_result = await self._get_from_cache(cache_key) if cached_result is not None: return cached_result try: async with self.get_connection() as conn: async with conn.cursor(aiomysql.DictCursor) as cursor: await cursor.execute(query, params) results = await cursor.fetchall() # 缓存结果 if results and self.redis: await self._store_in_cache(cache_key, results) logger.debug(f"查询结果数量: {len(results)}") return results except Exception as e: logger.error(f"查询多条失败: {query}, 参数: {params}, 错误: {str(e)}") raise async def _get_from_cache(self, key: str) -> Optional[Union[Dict, List[Dict]]]: """从Redis缓存获取数据""" if not self.redis: return None try: cached_data = await self.redis.get(key) if cached_data: return json.loads(cached_data) except Exception as e: logger.error(f"从缓存获取数据失败: {e}") return None async def _store_in_cache(self, key: str, data: Union[Dict, List[Dict]], ttl: int = 300) -> bool: """存储数据到Redis缓存""" if not self.redis: return False try: json_data = json.dumps(data) if REDIS_TYPE == "redis-py": await self.redis.setex(key, ttl, json_data) else: await self.redis.setex(key, ttl, json_data) return True except Exception as e: logger.error(f"存储数据到缓存失败: {e}") return False async def clear_cache(self, pattern: str = "db:*") -> int: """清除缓存""" if not self.redis: return 0 try: if REDIS_TYPE == "redis-py": keys = await self.redis.keys(pattern) if not keys: return 0 return await self.redis.delete(*keys) else: keys = await self.redis.keys(pattern) if not keys: return 0 return await self.redis.delete(*keys) except Exception as e: logger.error(f"清除缓存失败: {e}") return 0 def _make_cache_key(self, query: str, params: tuple) -> str: """生成缓存键""" return f"{query}:{hash(params)}"