ProxyPool/repositories/task_repo.py

"""验证任务队列持久化层"""
import aiosqlite
from typing import List, Optional
from models.domain import ProxyRaw
from core.log import logger


class ValidationTaskRepository:
    """验证任务 Repository —— 支持队列持久化"""

    @staticmethod
    async def insert_batch(db: aiosqlite.Connection, proxies: List[ProxyRaw]) -> int:
        if not proxies:
            return 0
        try:
            rows = [(p.ip, p.port, p.protocol) for p in proxies]
            await db.executemany(
                """
                INSERT INTO validation_tasks (ip, port, protocol, status, created_at)
                VALUES (?, ?, ?, 'pending', CURRENT_TIMESTAMP)
                """,
                rows,
            )
            await db.commit()
            return len(rows)
        except Exception as e:
            logger.error(f"insert_batch validation tasks failed: {e}")
            return 0

    @staticmethod
    async def acquire_pending(db: aiosqlite.Connection) -> Optional[dict]:
        """原子性地获取一个 pending 任务并将其标记为 processing"""
        try:
            async with db.execute(
                """
                SELECT id, ip, port, protocol FROM validation_tasks
                WHERE status = 'pending'
                ORDER BY id ASC
                LIMIT 1
                """
            ) as cursor:
                row = await cursor.fetchone()
                if not row:
                    return None
                task_id = row[0]
                await db.execute(
                    "UPDATE validation_tasks SET status = 'processing', updated_at = CURRENT_TIMESTAMP WHERE id = ?",
                    (task_id,),
                )
                await db.commit()
                return {"id": task_id, "ip": row[1], "port": row[2], "protocol": row[3]}
        except Exception as e:
            logger.error(f"acquire_pending failed: {e}")
            return None

    @staticmethod
    async def complete_task(
        db: aiosqlite.Connection,
        task_id: int,
        is_valid: bool,
        response_time_ms: Optional[float] = None,
    ) -> bool:
        try:
            await db.execute(
                """
                UPDATE validation_tasks
                SET status = 'completed',
                    result = ?,
                    response_time_ms = ?,
                    updated_at = CURRENT_TIMESTAMP
                WHERE id = ?
                """,
                ("valid" if is_valid else "invalid", response_time_ms, task_id),
            )
            await db.commit()
            return True
        except Exception as e:
            logger.error(f"complete_task failed: {e}")
            return False

    @staticmethod
    async def fail_task(db: aiosqlite.Connection, task_id: int) -> bool:
        try:
            await db.execute(
                """
                UPDATE validation_tasks
                SET status = 'failed',
                    result = 'invalid',
                    updated_at = CURRENT_TIMESTAMP
                WHERE id = ?
                """,
                (task_id,),
            )
            await db.commit()
            return True
        except Exception as e:
            logger.error(f"fail_task failed: {e}")
            return False

    @staticmethod
    async def get_pending_count(db: aiosqlite.Connection) -> int:
        async with db.execute(
            "SELECT COUNT(*) FROM validation_tasks WHERE status = 'pending'"
        ) as cursor:
            row = await cursor.fetchone()
            return row[0] if row else 0

    @staticmethod
    async def reset_processing(db: aiosqlite.Connection) -> int:
        """将异常中断的 processing 任务重置为 pending，用于启动恢复"""
        try:
            await db.execute(
                """
                UPDATE validation_tasks
                SET status = 'pending', updated_at = CURRENT_TIMESTAMP
                WHERE status = 'processing'
                """
            )
            await db.commit()
            return db.total_changes
        except Exception as e:
            logger.error(f"reset_processing failed: {e}")
            return 0

    @staticmethod
    async def cleanup_old(db: aiosqlite.Connection, days: int = 7) -> int:
        try:
            await db.execute(
                "DELETE FROM validation_tasks WHERE updated_at < datetime('now', '-{} days')".format(days)
            )
            await db.commit()
            return db.total_changes
        except Exception as e:
            logger.error(f"cleanup_old tasks failed: {e}")
            return 0