backend/database/connection.py

import threading
import mysql.connector
from mysql.connector import pooling
from backend.config import DB_HOST, DB_PORT, DB_USER, DB_PASSWORD, DB_NAME

# ── Connection Pool (Lazy Initialization) ─────────────────────────────────────
# We initialize the pool lazily on first access.
# This prevents the application from crashing on startup if the database is 
# temporarily down, slow, or if the DNS name is not yet resolvable.
_pool = None
_pool_lock = threading.Lock()

def get_pool():
    global _pool
    if _pool is None:
        with _pool_lock:
            if _pool is None:
                print(f"[Database] Initializing MySQLConnectionPool: host={DB_HOST}, port={DB_PORT}, db={DB_NAME}")
                _pool = pooling.MySQLConnectionPool(
                    pool_name="rag_pool",
                    pool_size=5,
                    host=DB_HOST,
                    port=DB_PORT,
                    user=DB_USER,
                    password=DB_PASSWORD,
                    database=DB_NAME,
                    charset="utf8mb4",
                    collation="utf8mb4_unicode_ci",
                    autocommit=False,
                    ssl_disabled=False,  # Enable SSL for Aiven
                )
    return _pool


def get_connection():
    """
    Get a connection from the pool.
    Always use this inside a 'with' block so it auto-returns to pool.

    Usage:
        with get_connection() as conn:
            cursor = conn.cursor(dictionary=True)
            cursor.execute("SELECT * FROM sources")
            rows = cursor.fetchall()
            conn.commit()
    """
    pool = get_pool()
    return pool.get_connection()


def check_and_update_schema():
    """
    Ensures that the sources table has the progress_percentage and error_message columns.
    This guarantees that cloud persistence works even if the DB hasn't been migrated manually.
    """
    try:
        conn = get_connection()
        cursor = conn.cursor()
        
        # Check if progress_percentage column exists
        try:
            cursor.execute("SELECT progress_percentage FROM sources LIMIT 1")
            cursor.fetchall()
            print("[Schema] progress_percentage column already exists.")
        except Exception:
            print("[Schema] Adding progress_percentage column to sources table...")
            try:
                cursor.execute("ALTER TABLE sources ADD COLUMN progress_percentage INT DEFAULT 0")
                conn.commit()
                print("[Schema] Successfully added progress_percentage column.")
            except Exception as ex:
                print(f"[Schema] Error adding progress_percentage: {ex}")
            
        # Check if error_message column exists
        try:
            cursor.execute("SELECT error_message FROM sources LIMIT 1")
            cursor.fetchall()
            print("[Schema] error_message column already exists.")
        except Exception:
            print("[Schema] Adding error_message column to sources table...")
            try:
                cursor.execute("ALTER TABLE sources ADD COLUMN error_message TEXT DEFAULT NULL")
                conn.commit()
                print("[Schema] Successfully added error_message column.")
            except Exception as ex:
                print(f"[Schema] Error adding error_message: {ex}")
            
        cursor.close()
        conn.close()
        print("[Schema] Sources table check completed successfully.")
    except Exception as e:
        print(f"[Schema] WARNING: Failed to check/update sources table schema: {e}")


def update_source_progress(source_id: str, status: str, progress_percentage: int, error_message: str = None, chunk_count: int = None):
    """
    Updates the status, progress_percentage, and optionally error_message or chunk_count
    for a given source in the MySQL database.
    """
    try:
        conn = get_connection()
        cursor = conn.cursor()
        if error_message is not None:
            cursor.execute("""
                UPDATE sources 
                SET status = %s, progress_percentage = %s, error_message = %s
                WHERE id = %s
            """, (status, progress_percentage, error_message, source_id))
        elif chunk_count is not None:
            cursor.execute("""
                UPDATE sources 
                SET status = %s, progress_percentage = %s, chunk_count = %s
                WHERE id = %s
            """, (status, progress_percentage, chunk_count, source_id))
        else:
            cursor.execute("""
                UPDATE sources 
                SET status = %s, progress_percentage = %s
                WHERE id = %s
            """, (status, progress_percentage, source_id))
        conn.commit()
        cursor.close()
        conn.close()
        print(f"[Progress] Source {source_id[:8]} updated: {status} ({progress_percentage}%)")
    except Exception as e:
        print(f"[Progress] Failed to update progress for source {source_id}: {e}")