db.py

from mysql.connector import pooling, Error
import os
import time

# Global pool variable
cnxpool = None

def create_pool(retries=3, delay=2):
    """
    Attempt to create the connection pool with retries.
    This prevents immediate crash if the DB is not ready yet.
    """
    global cnxpool
    
    dbconfig = {
        "host": os.getenv("MYSQL_HOST", "localhost"),
        "port": int(os.getenv("MYSQL_PORT", 3306)),
        "user": os.getenv("MYSQL_USER", "root"),
        "password": os.getenv("MYSQL_PASSWORD", ""),
        "database": os.getenv("MYSQL_DB", "ragdb")
    }

    for i in range(retries):
        try:
            print(f"Attempting to connect to MySQL (Attempt {i+1}/{retries})...")
            cnxpool = pooling.MySQLConnectionPool(pool_name="rag_pool",
                                                  pool_size=5,
                                                  **dbconfig)
            
            # Initialize Tables immediately to verify connection works
            conn = cnxpool.get_connection()
            cursor = conn.cursor()
            cursor.execute("""
            CREATE TABLE IF NOT EXISTS chunks (
                id INT AUTO_INCREMENT PRIMARY KEY,
                document VARCHAR(255),
                chunk_id INT,
                text LONGTEXT,
                embedding LONGTEXT
            )
            """)
            conn.commit()
            cursor.close()
            conn.close()
            print("Database initialized successfully.")
            return
            
        except Error as e:
            print(f"Database connection failed: {e}")
            if i < retries - 1:
                time.sleep(delay)
            else:
                print("Could not connect to database after retries. Proceeding without DB (App will start but DB features will fail).")
                cnxpool = None

# Initialize pool on module load (but safely)
create_pool()

def get_db_connection():
    global cnxpool
    if cnxpool is None:
        # Try one more time? Or just fail.
        # Let's try to reconnect if it's None (maybe DB came up later)
        create_pool(retries=1, delay=0)
        
    if cnxpool is None:
        raise Exception("Database connection is not available. Check MYSQL_* environment variables.")
        
    return cnxpool.get_connection()