Files
ProxyPool/app/core/db.py
祀梦 38bd66128b 重构: 迁移后端代码到 app 目录,前端移动到 WebUI,添加完整测试套件
主要变更:
- 后端代码从根目录迁移到 app/ 目录
- 前端代码从 frontend/ 重命名为 WebUI/
- 更新所有导入路径以适配新结构
- 提取公共 API 响应函数到 app/api/common.py
- 精简验证器服务代码
- 更新启动脚本和文档

测试:
- 新增完整测试套件 (tests/)
- 单元测试: 模型、仓库层
- 集成测试: 覆盖所有 22+ API 端点
- E2E 测试: 4个完整工作流场景
- 添加 pytest 配置和测试运行脚本
2026-04-04 13:32:36 +08:00

121 lines
4.6 KiB
Python

"""数据库连接管理 - 使用上下文管理器,避免全局单例连接泄漏"""
import os
import aiosqlite
from contextlib import asynccontextmanager
from typing import AsyncIterator
from app.core.config import settings
from app.core.log import logger
DB_PATH = os.path.join(settings.base_dir, settings.db_path)
def ensure_db_dir():
db_dir = os.path.dirname(DB_PATH)
if db_dir and not os.path.exists(db_dir):
os.makedirs(db_dir, exist_ok=True)
async def init_db():
"""初始化数据库表结构(支持迁移)"""
ensure_db_dir()
async with aiosqlite.connect(DB_PATH) as db:
await db.execute("PRAGMA journal_mode=WAL")
await db.execute("PRAGMA synchronous=NORMAL")
await db.execute("PRAGMA cache_size=-64000")
await db.execute("PRAGMA temp_store=MEMORY")
await db.execute("""
CREATE TABLE IF NOT EXISTS proxies (
id INTEGER PRIMARY KEY AUTOINCREMENT,
ip TEXT NOT NULL,
port INTEGER NOT NULL,
protocol TEXT DEFAULT 'http',
score INTEGER DEFAULT 10,
response_time_ms REAL,
last_check TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
UNIQUE(ip, port)
)
""")
# 迁移:如果旧表缺少 response_time_ms 列,则添加
try:
await db.execute("SELECT response_time_ms FROM proxies LIMIT 1")
except Exception:
await db.execute("ALTER TABLE proxies ADD COLUMN response_time_ms REAL")
logger.info("Migrated: added response_time_ms column")
# 迁移:如果旧表缺少 created_at 列,则添加
try:
await db.execute("SELECT created_at FROM proxies LIMIT 1")
except Exception:
await db.execute("ALTER TABLE proxies ADD COLUMN created_at TIMESTAMP")
await db.execute("UPDATE proxies SET created_at = CURRENT_TIMESTAMP WHERE created_at IS NULL")
logger.info("Migrated: added created_at column")
await db.execute("CREATE INDEX IF NOT EXISTS idx_score ON proxies(score)")
await db.execute("CREATE INDEX IF NOT EXISTS idx_protocol ON proxies(protocol)")
await db.execute("CREATE INDEX IF NOT EXISTS idx_last_check ON proxies(last_check)")
await db.execute("CREATE INDEX IF NOT EXISTS idx_ip_port ON proxies(ip, port)")
# 插件设置表
await db.execute("""
CREATE TABLE IF NOT EXISTS plugin_settings (
plugin_id TEXT PRIMARY KEY,
enabled INTEGER DEFAULT 1,
config_json TEXT DEFAULT '{}',
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
""")
# 迁移:为旧版 plugin_settings 表增加 config_json 列
try:
await db.execute("SELECT config_json FROM plugin_settings LIMIT 1")
except Exception:
await db.execute("ALTER TABLE plugin_settings ADD COLUMN config_json TEXT DEFAULT '{}'")
logger.info("Migrated: added config_json column to plugin_settings")
# 验证任务队列表
await db.execute("""
CREATE TABLE IF NOT EXISTS validation_tasks (
id INTEGER PRIMARY KEY AUTOINCREMENT,
ip TEXT NOT NULL,
port INTEGER NOT NULL,
protocol TEXT DEFAULT 'http',
status TEXT DEFAULT 'pending',
result TEXT,
response_time_ms REAL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
""")
await db.execute("CREATE INDEX IF NOT EXISTS idx_validation_status ON validation_tasks(status)")
await db.execute("CREATE INDEX IF NOT EXISTS idx_validation_created ON validation_tasks(created_at)")
# 系统设置表
await db.execute("""
CREATE TABLE IF NOT EXISTS settings (
key TEXT PRIMARY KEY,
value TEXT NOT NULL,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
""")
await db.commit()
logger.info("Database initialized")
@asynccontextmanager
async def get_db() -> AsyncIterator[aiosqlite.Connection]:
"""获取数据库连接的异步上下文管理器"""
ensure_db_dir()
db = await aiosqlite.connect(DB_PATH)
try:
await db.execute("PRAGMA journal_mode=WAL")
await db.execute("PRAGMA synchronous=NORMAL")
yield db
finally:
await db.close()