x

core/database.py

@@ -0,0 +1,86 @@
+import asyncio
+from contextlib import asynccontextmanager
+from typing import Any, List, Optional
+
+import aiosqlite
+from loguru import logger
+
+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._pool_lock = asyncio.Lock()
+        
+    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)
+        logger.info(f"数据库连接池初始化完成，大小: {self._pool_size}")
+    
+    async def cleanup(self):
+        """清理数据库连接"""
+        for conn in self._pool:
+            await conn.close()
+        self._pool.clear()
+    
+    @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()
+    
+    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
+                
+    async def fetch_one(self, query: str, params: tuple = ()) -> Optional[tuple]:
+        """查询单条记录"""
+        async with self.get_connection() as conn:
+            cursor = await conn.execute(query, params)
+            return await cursor.fetchone()
+                
+    async def fetch_all(self, query: str, params: tuple = ()) -> List[tuple]:
+        """查询多条记录"""
+        async with self.get_connection() as conn:
+            cursor = await conn.execute(query, params)
+            return await cursor.fetchall()