youtube-api/database.py

import os
import psycopg2
from psycopg2.extras import RealDictCursor
from typing import Optional, Dict
from dotenv import load_dotenv

load_dotenv()

def get_db_connection():
    return psycopg2.connect(
        host=os.getenv("DB_HOST", "localhost"),
        port=os.getenv("DB_PORT", "5432"),
        database=os.getenv("DB_NAME", "postgres"),
        user=os.getenv("DB_USER", "postgres"),
        password=os.getenv("DB_PASSWORD", ""),
        cursor_factory=RealDictCursor
    )

def get_latest_proxy() -> Optional[Dict]:
    try:
        conn = get_db_connection()
        cursor = conn.cursor()

        cursor.execute("""
            SELECT
                id, ip_address, port, protocol, username, password,
                country_code, country_name, city, is_active, is_anonymous,
                response_time_ms, last_checked_at, last_successful_at,
                failure_count, success_count, usage, source, notes,
                created_at, updated_at
            FROM proxies
            WHERE is_active = TRUE
            ORDER BY
                last_successful_at DESC NULLS LAST,
                response_time_ms ASC NULLS LAST,
                (CASE WHEN success_count + failure_count > 0
                      THEN CAST(success_count AS FLOAT) / (success_count + failure_count)
                      ELSE 0 END) DESC,
                created_at DESC
            LIMIT 1
        """)

        proxy = cursor.fetchone()
        cursor.close()
        conn.close()

        return dict(proxy) if proxy else None
    except Exception:
        return None

def get_all_active_proxies() -> list:
    """
    Retorna todos os proxies ativos do banco, priorizando:
    1. Proxies com sucesso recente
    2. Proxies com baixo failure_count
    3. Proxies com boa taxa de sucesso
    4. Proxies novos ainda não testados
    """
    try:
        conn = get_db_connection()
        cursor = conn.cursor()

        cursor.execute("""
            SELECT
                id, ip_address, port, protocol, username, password,
                country_code, country_name, city, is_active, is_anonymous,
                response_time_ms, last_checked_at, last_successful_at,
                failure_count, success_count, usage, source, notes,
                created_at, updated_at,
                (CASE WHEN success_count + failure_count > 0
                      THEN CAST(success_count AS FLOAT) / (success_count + failure_count)
                      ELSE 0.5 END) as success_rate
            FROM proxies
            WHERE is_active = TRUE
                AND failure_count < 8  -- Ignora proxies com muitas falhas
            ORDER BY
                -- Prioriza proxies com sucesso recente
                last_successful_at DESC NULLS LAST,
                -- Depois por taxa de sucesso
                success_rate DESC,
                -- Depois por menos falhas
                failure_count ASC,
                -- Por último, proxies novos
                created_at DESC
            LIMIT 50  -- Limita a 50 melhores proxies para não perder tempo
        """)

        proxies = cursor.fetchall()
        cursor.close()
        conn.close()

        return [dict(proxy) for proxy in proxies] if proxies else []
    except Exception:
        return []

def delete_proxy(proxy_id: int) -> bool:
    try:
        conn = get_db_connection()
        cursor = conn.cursor()

        cursor.execute("""
            UPDATE proxies
            SET is_active = FALSE,
                failure_count = failure_count + 1,
                last_checked_at = NOW(),
                updated_at = NOW()
            WHERE id = %s
        """, (proxy_id,))
        conn.commit()

        updated = cursor.rowcount > 0
        cursor.close()
        conn.close()

        return updated
    except Exception:
        return False

def format_proxy_url(proxy: Dict) -> str:
    protocol = proxy.get('protocol', 'http').lower()
    ip_address = proxy.get('ip_address')
    port = proxy.get('port')
    username = proxy.get('username')
    password = proxy.get('password')

    if username and password:
        return f"{protocol}://{username}:{password}@{ip_address}:{port}"
    else:
        return f"{protocol}://{ip_address}:{port}"

def mark_proxy_success(proxy_id: int) -> bool:
    try:
        conn = get_db_connection()
        cursor = conn.cursor()

        cursor.execute("""
            UPDATE proxies
            SET success_count = success_count + 1,
                last_successful_at = NOW(),
                last_checked_at = NOW(),
                updated_at = NOW(),
                is_active = TRUE
            WHERE id = %s
        """, (proxy_id,))
        conn.commit()

        updated = cursor.rowcount > 0
        cursor.close()
        conn.close()

        return updated
    except Exception:
        return False

def mark_proxy_failure(proxy_id: int, max_failures: int = 10) -> bool:
    """
    Marca proxy como falha e desativa se atingir max_failures consecutivas.

    Args:
        proxy_id: ID do proxy
        max_failures: Número máximo de falhas antes de desativar (padrão: 10)
    """
    try:
        conn = get_db_connection()
        cursor = conn.cursor()

        # Incrementa failure_count e desativa se atingir o limite
        cursor.execute("""
            UPDATE proxies
            SET failure_count = failure_count + 1,
                last_checked_at = NOW(),
                updated_at = NOW(),
                is_active = CASE
                    WHEN failure_count + 1 >= %s THEN FALSE
                    ELSE is_active
                END
            WHERE id = %s
            RETURNING failure_count, is_active
        """, (max_failures, proxy_id))

        result = cursor.fetchone()
        conn.commit()
        cursor.close()
        conn.close()

        return result is not None

    except Exception:
        return False