"""代理数据访问层 - 所有 SQL 操作收敛于此""" import aiosqlite from datetime import datetime, timedelta from typing import List, Optional, Tuple, Union from app.models.domain import Proxy from app.core.log import logger VALID_PROTOCOLS = ("http", "https", "socks4", "socks5") def _to_datetime(value: Union[str, datetime, None]) -> Optional[datetime]: if value is None: return None if isinstance(value, datetime): return value if isinstance(value, str): for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M:%S.%f", "%Y-%m-%dT%H:%M:%S", "%Y-%m-%dT%H:%M:%S.%f"): try: return datetime.strptime(value, fmt) except ValueError: continue return None def _row_to_proxy(row: Tuple) -> Proxy: return Proxy( ip=row[0], port=row[1], protocol=row[2], score=row[3], response_time_ms=row[4], last_check=_to_datetime(row[5]), created_at=_to_datetime(row[6]), ) class ProxyRepository: """代理 Repository""" @staticmethod async def insert_or_update( db: aiosqlite.Connection, ip: str, port: int, protocol: str = "http", score: int = 10, ) -> bool: if protocol not in VALID_PROTOCOLS: protocol = "http" try: await db.execute( """ INSERT INTO proxies (ip, port, protocol, score, last_check, created_at) VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) ON CONFLICT(ip, port) DO UPDATE SET protocol = excluded.protocol, score = excluded.score, last_check = CURRENT_TIMESTAMP """, (ip, port, protocol, score), ) await db.commit() return True except Exception as e: logger.error(f"insert_or_update proxy failed: {e}", exc_info=True) return False @staticmethod async def update_score( db: aiosqlite.Connection, ip: str, port: int, delta: int, min_score: int = 0, max_score: int = 100, ) -> bool: try: # 原子更新:计算新分数并直接更新 await db.execute( """ UPDATE proxies SET score = MAX(?, MIN(?, score + ?)), last_check = CURRENT_TIMESTAMP WHERE ip = ? AND port = ? """, (min_score, max_score, delta, ip, port), ) # 删除分数已降至 0 及以下的代理 await db.execute( "DELETE FROM proxies WHERE ip = ? AND port = ? AND score <= ?", (ip, port, min_score), ) await db.commit() return db.total_changes > 0 except Exception as e: logger.error(f"update_score failed: {e}", exc_info=True) return False @staticmethod async def update_response_time( db: aiosqlite.Connection, ip: str, port: int, response_time_ms: float, ) -> bool: try: await db.execute( "UPDATE proxies SET response_time_ms = ? WHERE ip = ? AND port = ?", (response_time_ms, ip, port), ) await db.commit() return True except Exception as e: logger.error(f"update_response_time failed: {e}", exc_info=True) return False @staticmethod async def delete(db: aiosqlite.Connection, ip: str, port: int) -> None: await db.execute("DELETE FROM proxies WHERE ip = ? AND port = ?", (ip, port)) await db.commit() @staticmethod async def batch_delete(db: aiosqlite.Connection, proxies: List[Tuple[str, int]]) -> int: if not proxies: return 0 await db.executemany("DELETE FROM proxies WHERE ip = ? AND port = ?", proxies) await db.commit() return len(proxies) @staticmethod async def get_by_ip_port( db: aiosqlite.Connection, ip: str, port: int ) -> Optional[Proxy]: async with db.execute( "SELECT ip, port, protocol, score, response_time_ms, last_check, created_at FROM proxies WHERE ip = ? AND port = ?", (ip, port), ) as cursor: row = await cursor.fetchone() if row: return _row_to_proxy(row) return None @staticmethod async def get_random(db: aiosqlite.Connection) -> Optional[Proxy]: async with db.execute( "SELECT ip, port, protocol, score, response_time_ms, last_check, created_at FROM proxies WHERE score > 0 ORDER BY RANDOM() LIMIT 1" ) as cursor: row = await cursor.fetchone() if row: return _row_to_proxy(row) return None @staticmethod async def list_all( db: aiosqlite.Connection, protocol: Optional[str] = None, limit: int = 100000, offset: int = 0, ) -> List[Proxy]: query = "SELECT ip, port, protocol, score, response_time_ms, last_check, created_at FROM proxies" params: List = [] if protocol: query += " WHERE protocol = ?" params.append(protocol.lower()) query += " LIMIT ? OFFSET ?" params.extend([limit, offset]) async with db.execute(query, params) as cursor: rows = await cursor.fetchall() return [_row_to_proxy(row) for row in rows] @staticmethod async def iter_batches( db: aiosqlite.Connection, protocol: Optional[str] = None, batch_size: int = 1000, ): """流式分批读取代理,避免一次性加载大量数据到内存""" offset = 0 while True: batch = await ProxyRepository.list_all(db, protocol, batch_size, offset) if not batch: break yield batch offset += batch_size @staticmethod async def list_paginated( db: aiosqlite.Connection, page: int = 1, page_size: int = 20, protocol: Optional[str] = None, min_score: int = 0, max_score: Optional[int] = None, sort_by: str = "last_check", sort_order: str = "DESC", ) -> Tuple[List[Proxy], int]: conditions = ["score >= ?"] params: List = [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) allowed_sort_by = {"ip", "port", "protocol", "score", "last_check"} allowed_sort_order = {"ASC", "DESC"} if sort_by not in allowed_sort_by or sort_order.upper() not in allowed_sort_order: order_clause = "last_check DESC" else: order_clause = f"{sort_by} {sort_order.upper()}" offset = (page - 1) * page_size count_query = f"SELECT COUNT(*) FROM proxies WHERE {where_clause}" async with db.execute(count_query, list(params)) as cursor: row = await cursor.fetchone() total = row[0] if row else 0 data_query = f""" SELECT ip, port, protocol, score, response_time_ms, last_check, created_at FROM proxies WHERE {where_clause} ORDER BY {order_clause} LIMIT ? OFFSET ? """ params.extend([page_size, offset]) async with db.execute(data_query, params) as cursor: rows = await cursor.fetchall() proxies = [_row_to_proxy(row) for row in rows] return proxies, total @staticmethod async def get_stats(db: aiosqlite.Connection) -> dict: query = """ SELECT COUNT(*) as total, COUNT(CASE WHEN score > 0 THEN 1 END) as available, AVG(score) as avg_score, COUNT(CASE WHEN protocol = 'http' THEN 1 END) as http_count, COUNT(CASE WHEN protocol = 'https' THEN 1 END) as https_count, COUNT(CASE WHEN protocol = 'socks4' THEN 1 END) as socks4_count, COUNT(CASE WHEN protocol = 'socks5' THEN 1 END) as socks5_count FROM proxies """ async with db.execute(query) as cursor: row = await cursor.fetchone() if row: return { "total": row[0] or 0, "available": row[1] or 0, "avg_score": round(row[2], 2) if row[2] else 0, "http_count": row[3] or 0, "https_count": row[4] or 0, "socks4_count": row[5] or 0, "socks5_count": row[6] or 0, } return { "total": 0, "available": 0, "avg_score": 0, "http_count": 0, "https_count": 0, "socks4_count": 0, "socks5_count": 0, } @staticmethod async def get_today_new_count(db: aiosqlite.Connection) -> int: try: async with db.execute( "SELECT COUNT(*) FROM proxies WHERE DATE(created_at) = DATE('now', 'localtime')" ) as cursor: row = await cursor.fetchone() return row[0] if row else 0 except Exception as e: logger.error(f"get_today_new_count failed: {e}", exc_info=True) return 0 @staticmethod async def clean_invalid(db: aiosqlite.Connection) -> int: await db.execute("DELETE FROM proxies WHERE score <= 0") await db.commit() return db.total_changes @staticmethod async def clean_expired(db: aiosqlite.Connection, days: int) -> int: try: await db.execute( "DELETE FROM proxies WHERE last_check < datetime('now', '-' || ? || ' days')", (days,), ) await db.commit() return db.total_changes except Exception as e: logger.error(f"clean_expired failed: {e}", exc_info=True) return 0