准备创建mysqlv1分支的提交

core/database.py

@@ -1,86 +1,274 @@
 import asyncio
+import json
 from contextlib import asynccontextmanager
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional, Tuple, Union
 
-import aiosqlite
+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_path = config.database_config.path
-        self._pool_size = config.database_config.pool_size
-        self._pool: List[aiosqlite.Connection] = []
+        # 数据库配置
+        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("初始化数据库连接池")
-        async with aiosqlite.connect(self.db_path) as db:
-            await db.execute('''
-                CREATE TABLE IF NOT EXISTS email_accounts (
-                    id INTEGER PRIMARY KEY AUTOINCREMENT,
-                    email TEXT UNIQUE NOT NULL,
-                    password TEXT NOT NULL,
-                    client_id TEXT NOT NULL,
-                    refresh_token TEXT NOT NULL,
-                    in_use BOOLEAN DEFAULT 0,
-                    cursor_password TEXT,
-                    cursor_cookie TEXT,
-                    sold BOOLEAN DEFAULT 0,
-                    status TEXT DEFAULT 'pending',
-                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-                )
-            ''')
-            await db.commit()
-            
-        # 初始化连接池
-        for i in range(self._pool_size):
-            conn = await aiosqlite.connect(self.db_path)
-            self._pool.append(conn)
+        
+        # 创建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',
+                        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)
+                    )
+                ''')
+                
+        # 初始化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):
         """清理数据库连接"""
-        for conn in self._pool:
-            await conn.close()
-        self._pool.clear()
+        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):
         """获取数据库连接"""
-        async with self._pool_lock:
-            if not self._pool:
-                conn = await aiosqlite.connect(self.db_path)
-            else:
-                conn = self._pool.pop()
-                
-        try:
-            yield conn
-        finally:
-            if len(self._pool) < self._pool_size:
-                self._pool.append(conn)
-            else:
-                await conn.close()
+        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语句"""
-        async with self.get_connection() as conn:
-            cursor = await conn.execute(query, params)
-            await conn.commit()
-            return cursor.lastrowid
+        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[tuple]:
+    async def fetch_one(self, query: str, params: tuple = ()) -> Optional[Dict]:
         """查询单条记录"""
-        async with self.get_connection() as conn:
-            cursor = await conn.execute(query, params)
-            return await cursor.fetchone()
+        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[tuple]:
+    async def fetch_all(self, query: str, params: tuple = ()) -> List[Dict]:
         """查询多条记录"""
-        async with self.get_connection() as conn:
-            cursor = await conn.execute(query, params)
-            return await cursor.fetchall() 
+        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)}"