first commit
This commit is contained in:
祀梦
334
core/sqlite.py
Normal file
334
core/sqlite.py
Normal file
@@ -0,0 +1,334 @@
|
||||
import aiosqlite
|
||||
import os
|
||||
import asyncio
|
||||
from core.log import logger
|
||||
|
||||
VALID_PROTOCOLS = ['http', 'https', 'socks4', 'socks5']
|
||||
|
||||
class SQLiteManager:
|
||||
_instance = None
|
||||
_connection = None
|
||||
_lock = asyncio.Lock()
|
||||
|
||||
def __new__(cls, *args, **kwargs):
|
||||
if cls._instance is None:
|
||||
cls._instance = super(SQLiteManager, cls).__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self, db_path=None):
|
||||
if hasattr(self, 'initialized') and self.initialized:
|
||||
return
|
||||
|
||||
if db_path is None:
|
||||
base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
db_dir = os.path.join(base_dir, 'db')
|
||||
if not os.path.exists(db_dir):
|
||||
os.makedirs(db_dir)
|
||||
self.db_path = os.path.join(db_dir, 'proxies.sqlite')
|
||||
else:
|
||||
self.db_path = db_path
|
||||
|
||||
self.initialized = True
|
||||
|
||||
async def get_connection(self):
|
||||
async with self._lock:
|
||||
if self._connection is None:
|
||||
self._connection = await aiosqlite.connect(self.db_path)
|
||||
await self._connection.execute("PRAGMA journal_mode=WAL")
|
||||
await self._connection.execute("PRAGMA synchronous=NORMAL")
|
||||
await self._connection.execute("PRAGMA cache_size=-64000")
|
||||
await self._connection.execute("PRAGMA temp_store=MEMORY")
|
||||
return self._connection
|
||||
|
||||
async def close_connection(self):
|
||||
async with self._lock:
|
||||
if self._connection is not None:
|
||||
await self._connection.close()
|
||||
self._connection = None
|
||||
|
||||
async def init_db(self):
|
||||
"""初始化数据库和表结构"""
|
||||
db = await self.get_connection()
|
||||
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,
|
||||
last_check TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
UNIQUE(ip, port)
|
||||
)
|
||||
''')
|
||||
|
||||
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.commit()
|
||||
|
||||
async def insert_proxy(self, ip, port, protocol='http', score=10):
|
||||
"""异步插入或更新代理"""
|
||||
try:
|
||||
# 验证协议类型
|
||||
if protocol not in VALID_PROTOCOLS:
|
||||
protocol = 'http'
|
||||
logger.warning(f"无效的协议类型 {protocol}，默认使用 http")
|
||||
|
||||
db = await self.get_connection()
|
||||
# 先检查是否存在
|
||||
async with db.execute('SELECT score FROM proxies WHERE ip = ? AND port = ?', (ip, port)) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
if row:
|
||||
# 如果存在，则更新最后检查时间和分数
|
||||
await db.execute('''
|
||||
UPDATE proxies SET last_check = CURRENT_TIMESTAMP, score = ?, protocol = ? WHERE ip = ? AND port = ?
|
||||
''', (score, protocol, ip, port))
|
||||
else:
|
||||
# 如果不存在，则插入新记录
|
||||
await db.execute('''
|
||||
INSERT INTO proxies (ip, port, protocol, score, last_check)
|
||||
VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP)
|
||||
''', (ip, port, protocol, score))
|
||||
await db.commit()
|
||||
return True
|
||||
except aiosqlite.IntegrityError as e:
|
||||
# 处理唯一性约束冲突
|
||||
if "UNIQUE" in str(e):
|
||||
# 代理已存在，更新它
|
||||
if protocol not in VALID_PROTOCOLS:
|
||||
protocol = 'http'
|
||||
db = await self.get_connection()
|
||||
await db.execute('''
|
||||
UPDATE proxies SET last_check = CURRENT_TIMESTAMP, score = ?, protocol = ? WHERE ip = ? AND port = ?
|
||||
''', (score, protocol, ip, port))
|
||||
await db.commit()
|
||||
return True
|
||||
else:
|
||||
logger.error(f"数据库完整性错误: {e}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"插入代理失败 {ip}:{port} - {e}")
|
||||
return False
|
||||
|
||||
async def get_all_proxies(self):
|
||||
"""异步获取所有代理"""
|
||||
db = await self.get_connection()
|
||||
async with db.execute('SELECT ip, port, protocol, score, last_check FROM proxies') as cursor:
|
||||
return await cursor.fetchall()
|
||||
|
||||
async def get_random_proxy(self):
|
||||
"""异步随机获取一个高分代理"""
|
||||
db = await self.get_connection()
|
||||
async with db.execute('SELECT ip, port, protocol, score, last_check FROM proxies WHERE score > 0 ORDER BY RANDOM() LIMIT 1') as cursor:
|
||||
return await cursor.fetchone()
|
||||
|
||||
async def update_score(self, ip, port, delta, min_score=0, max_score=100):
|
||||
"""异步更新代理分数（增量更新，带分数限制）"""
|
||||
try:
|
||||
db = await self.get_connection()
|
||||
# 获取当前分数
|
||||
async with db.execute('SELECT score FROM proxies WHERE ip = ? AND port = ?', (ip, port)) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
if row:
|
||||
current_score = row[0]
|
||||
new_score = max(min_score, min(max_score, current_score + delta))
|
||||
await db.execute('''
|
||||
UPDATE proxies SET score = ?, last_check = CURRENT_TIMESTAMP WHERE ip = ? AND port = ?
|
||||
''', (new_score, ip, port))
|
||||
if new_score <= 0:
|
||||
await db.execute('DELETE FROM proxies WHERE score <= 0')
|
||||
await db.commit()
|
||||
return True
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"更新代理分数失败 {ip}:{port} - {e}")
|
||||
return False
|
||||
|
||||
async def delete_proxy(self, ip, port):
|
||||
"""异步删除指定代理"""
|
||||
db = await self.get_connection()
|
||||
await db.execute('DELETE FROM proxies WHERE ip = ? AND port = ?', (ip, port))
|
||||
await db.commit()
|
||||
|
||||
async def count_proxies(self):
|
||||
"""异步统计代理数量"""
|
||||
db = await self.get_connection()
|
||||
async with db.execute('SELECT COUNT(*) FROM proxies') as cursor:
|
||||
row = await cursor.fetchone()
|
||||
return row[0] if row else 0
|
||||
|
||||
async def get_proxies_paginated_with_total(self, page: int = 1, page_size: int = 20,
|
||||
protocol: str = None, min_score: int = 0,
|
||||
max_score: int = None,
|
||||
sort_by: str = 'last_check',
|
||||
sort_order: str = 'DESC'):
|
||||
"""分页获取代理列表（一次查询返回数据和总数）"""
|
||||
db = await self.get_connection()
|
||||
conditions = ['score >= ?']
|
||||
params = [min_score]
|
||||
|
||||
if protocol:
|
||||
conditions.append('protocol = ?')
|
||||
params.append(protocol)
|
||||
|
||||
if max_score is not None:
|
||||
conditions.append('score <= ?')
|
||||
params.append(max_score)
|
||||
|
||||
where_clause = ' AND '.join(conditions)
|
||||
|
||||
order_by_clause = f'{sort_by} {sort_order}'
|
||||
|
||||
offset = (page - 1) * page_size
|
||||
query = f'''
|
||||
SELECT ip, port, protocol, score, last_check,
|
||||
COUNT(*) OVER() as total_count
|
||||
FROM proxies
|
||||
WHERE {where_clause}
|
||||
ORDER BY {order_by_clause}
|
||||
LIMIT ? OFFSET ?
|
||||
'''
|
||||
params.extend([page_size, offset])
|
||||
|
||||
async with db.execute(query, params) as cursor:
|
||||
rows = await cursor.fetchall()
|
||||
total = rows[0][5] if rows else 0
|
||||
proxies = [(row[0], row[1], row[2], row[3], row[4]) for row in rows]
|
||||
return proxies, total
|
||||
|
||||
async def get_proxies_paginated(self, page: int = 1, page_size: int = 20,
|
||||
protocol: str = None, min_score: int = 0,
|
||||
max_score: int = None,
|
||||
sort_by: str = 'last_check',
|
||||
sort_order: str = 'DESC'):
|
||||
"""分页获取代理列表"""
|
||||
db = await self.get_connection()
|
||||
conditions = ['score >= ?']
|
||||
params = [min_score]
|
||||
|
||||
if protocol:
|
||||
conditions.append('protocol = ?')
|
||||
params.append(protocol)
|
||||
|
||||
if max_score is not None:
|
||||
conditions.append('score <= ?')
|
||||
params.append(max_score)
|
||||
|
||||
where_clause = ' AND '.join(conditions)
|
||||
|
||||
order_by_clause = f'{sort_by} {sort_order}'
|
||||
|
||||
offset = (page - 1) * page_size
|
||||
query = f'''
|
||||
SELECT ip, port, protocol, score, last_check
|
||||
FROM proxies
|
||||
WHERE {where_clause}
|
||||
ORDER BY {order_by_clause}
|
||||
LIMIT ? OFFSET ?
|
||||
'''
|
||||
params.extend([page_size, offset])
|
||||
|
||||
async with db.execute(query, params) as cursor:
|
||||
return await cursor.fetchall()
|
||||
|
||||
async def get_proxies_total(self, protocol: str = None, min_score: int = 0, max_score: int = None):
|
||||
"""获取符合条件的代理总数"""
|
||||
db = await self.get_connection()
|
||||
conditions = ['score >= ?']
|
||||
params = [min_score]
|
||||
|
||||
if protocol:
|
||||
conditions.append('protocol = ?')
|
||||
params.append(protocol)
|
||||
|
||||
if max_score is not None:
|
||||
conditions.append('score <= ?')
|
||||
params.append(max_score)
|
||||
|
||||
where_clause = ' AND '.join(conditions)
|
||||
|
||||
query = f'SELECT COUNT(*) FROM proxies WHERE {where_clause}'
|
||||
|
||||
async with db.execute(query, params) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
return row[0] if row else 0
|
||||
|
||||
async def get_proxy_detail(self, ip: str, port: int):
|
||||
"""获取单个代理的详细信息"""
|
||||
db = await self.get_connection()
|
||||
async with db.execute(
|
||||
'SELECT ip, port, protocol, score, last_check FROM proxies WHERE ip = ? AND port = ?',
|
||||
(ip, port)
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
return row
|
||||
|
||||
async def batch_delete_proxies(self, proxy_list: list):
|
||||
"""批量删除代理，返回实际删除的数量"""
|
||||
deleted_count = 0
|
||||
db = await self.get_connection()
|
||||
for ip, port in proxy_list:
|
||||
cursor = await db.execute('DELETE FROM proxies WHERE ip = ? AND port = ?', (ip, port))
|
||||
deleted_count += cursor.rowcount
|
||||
await db.commit()
|
||||
return deleted_count
|
||||
|
||||
async def get_stats(self):
|
||||
"""获取统计信息"""
|
||||
db = await self.get_connection()
|
||||
stats = {}
|
||||
|
||||
async with db.execute('SELECT COUNT(*) FROM proxies') as cursor:
|
||||
row = await cursor.fetchone()
|
||||
stats['total'] = row[0] if row else 0
|
||||
|
||||
async with db.execute('SELECT COUNT(*) FROM proxies WHERE score > 0') as cursor:
|
||||
row = await cursor.fetchone()
|
||||
stats['available'] = row[0] if row else 0
|
||||
|
||||
async with db.execute('SELECT COUNT(*) FROM proxies WHERE protocol = "http"') as cursor:
|
||||
row = await cursor.fetchone()
|
||||
stats['http_count'] = row[0] if row else 0
|
||||
|
||||
async with db.execute('SELECT COUNT(*) FROM proxies WHERE protocol = "https"') as cursor:
|
||||
row = await cursor.fetchone()
|
||||
stats['https_count'] = row[0] if row else 0
|
||||
|
||||
async with db.execute('SELECT COUNT(*) FROM proxies WHERE protocol = "socks4"') as cursor:
|
||||
row = await cursor.fetchone()
|
||||
stats['socks4_count'] = row[0] if row else 0
|
||||
|
||||
async with db.execute('SELECT COUNT(*) FROM proxies WHERE protocol = "socks5"') as cursor:
|
||||
row = await cursor.fetchone()
|
||||
stats['socks5_count'] = row[0] if row else 0
|
||||
|
||||
async with db.execute('SELECT AVG(score) FROM proxies') as cursor:
|
||||
row = await cursor.fetchone()
|
||||
stats['avg_score'] = row[0] if row and row[0] else 0
|
||||
|
||||
return stats
|
||||
|
||||
async def get_today_new_count(self):
|
||||
"""获取今日新增代理数量"""
|
||||
try:
|
||||
db = await self.get_connection()
|
||||
query = '''
|
||||
SELECT COUNT(*) FROM proxies
|
||||
WHERE DATE(last_check) = DATE('now', 'localtime')
|
||||
'''
|
||||
async with db.execute(query) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
return row[0] if row else 0
|
||||
except Exception as e:
|
||||
logger.error(f"获取今日新增数量失败: {e}")
|
||||
return 0
|
||||
|
||||
async def clean_invalid_proxies(self):
|
||||
"""清理无效代理（分数<=0）"""
|
||||
db = await self.get_connection()
|
||||
async with db.execute('DELETE FROM proxies WHERE score <= 0') as cursor:
|
||||
deleted_count = cursor.rowcount
|
||||
await db.commit()
|
||||
return deleted_count
|
||||
Reference in New Issue
Block a user