ProxyPool/repositories/proxy_repo.py

"""代理数据访问层 - 所有 SQL 操作收敛于此"""
import aiosqlite
from datetime import datetime, timedelta
from typing import List, Optional, Tuple, Union
from models.domain import Proxy
from 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}")
            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:
            async with db.execute(
                "SELECT score FROM proxies WHERE ip = ? AND port = ?", (ip, port)
            ) as cursor:
                row = await cursor.fetchone()
                if not row:
                    return False
                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
        except Exception as e:
            logger.error(f"update_score failed: {e}")
            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}")
            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,
    ) -> 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 ?"
        params.append(limit)

        async with db.execute(query, params) as cursor:
            rows = await cursor.fetchall()
            return [_row_to_proxy(row) for row in rows]

    @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)
        order_clause = f"{sort_by} {sort_order}"
        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(last_check) = 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}")
            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')".format(days)
            )
            await db.commit()
            return db.total_changes
        except Exception as e:
            logger.error(f"clean_expired failed: {e}")
            return 0