scripts/setup_db.py

import mysql.connector
import sys
import os
from pathlib import Path

# Add the project root to sys.path so we can import backend
sys.path.append(str(Path(__file__).resolve().parent.parent))

from backend.database.connection import get_connection
from backend.config import DB_NAME

def run_setup():
    print("Starting Database Setup...")
    
    schema_path = Path(__file__).resolve().parent.parent / "backend" / "database" / "schema.sql"
    if not schema_path.exists():
        print(f"Error: Schema file not found: {schema_path}")
        return

    try:
        conn = get_connection()
        cursor = conn.cursor()

        # 1. Read and Execute schema.sql
        print(f"Reading {schema_path.name}...")
        with open(schema_path, 'r', encoding='utf-8') as f:
            schema_sql = f.read()

        # Split by semicolon to execute one by one
        statements = schema_sql.split(';')
        for statement in statements:
            stmt = statement.strip()
            if stmt and not stmt.startswith('--'):
                try:
                    cursor.execute(stmt)
                except mysql.connector.Error as err:
                    if "already exists" in str(err).lower() or "Duplicate" in str(err):
                        continue
                    # Ignore USE error on Aiven
                    if "Unknown database" in str(err):
                        continue
                    print(f"Warning on stmt: {stmt[:50]}... \n Error: {err}")

        # 2. Run additional safety checks/migrations
        print("Running additional safety checks...")
        
        # Ensure conv_type column exists in conversations
        try:
            cursor.execute("ALTER TABLE conversations ADD COLUMN conv_type VARCHAR(20) DEFAULT 'general'")
            print("Column 'conv_type' added to 'conversations'.")
        except mysql.connector.Error as err:
            if "Duplicate column name" not in str(err):
                print(f"Note (conv_type): {err}")

        # Ensure url_ref column exists in chunks
        try:
            cursor.execute("ALTER TABLE chunks ADD COLUMN url_ref TEXT DEFAULT NULL")
            print("Column 'url_ref' added to 'chunks'.")
        except mysql.connector.Error as err:
            if "Duplicate column name" not in str(err):
                print(f"Note (url_ref): {err}")

        # Ensure conversation_id column in chat_history
        try:
            cursor.execute("ALTER TABLE chat_history ADD COLUMN conversation_id VARCHAR(36) DEFAULT NULL")
            print("Column 'conversation_id' added to 'chat_history'.")
        except mysql.connector.Error as err:
            if "Duplicate column name" not in str(err):
                 print(f"Note (conversation_id): {err}")

        # Modify image_jobs.source_id to be NULLable
        try:
            cursor.execute("ALTER TABLE image_jobs MODIFY COLUMN source_id VARCHAR(36) NULL")
            print("Column 'image_jobs.source_id' modified to NULLable.")
        except mysql.connector.Error as err:
            pass

        conn.commit()
        cursor.close()
        conn.close()
        print("\nDatabase setup complete!")

    except Exception as e:
        print(f"Critical error during setup: {e}")

if __name__ == "__main__":
    run_setup()