backend/tests/migrate_database.py

"""
Database Migration Script for AegisLM

This script adds missing columns to the evaluations table
to align the database schema with the model definitions.
"""

import asyncio
import sys
import os

# Add the backend directory to Python path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from core.database import AsyncSessionLocal
from sqlalchemy import text

async def run_migration():
    """Add missing columns to evaluations table."""
    
    missing_columns = [
        # Experiment tracking
        ("experiment_run_id", "VARCHAR(255)", True, "CREATE INDEX IF NOT EXISTS idx_evaluations_experiment_run_id ON evaluations(experiment_run_id)"),
        
        # Audit tracking
        ("config_hash", "VARCHAR(64)", True, "CREATE INDEX IF NOT EXISTS idx_evaluations_config_hash ON evaluations(config_hash)"),
        ("result_checksum", "VARCHAR(64)", True, "CREATE INDEX IF NOT EXISTS idx_evaluations_result_checksum ON evaluations(result_checksum)"),
        ("audit_trail_id", "VARCHAR(255)", True, "CREATE INDEX IF NOT EXISTS idx_evaluations_audit_trail_id ON evaluations(audit_trail_id)"),
        
        # Performance metrics
        ("queue_delay_ms", "INTEGER", True, None),
        ("total_latency_ms", "INTEGER", True, None),
    ]
    
    async with AsyncSessionLocal() as db:
        try:
            # Check current columns
            result = await db.execute(text("""
                SELECT column_name 
                FROM information_schema.columns 
                WHERE table_name = 'evaluations'
            """))
            existing_columns = {row[0] for row in result.fetchall()}
            
            print(f"📊 Current columns: {len(existing_columns)}")
            print(f"🔧 Missing columns: {len(missing_columns)}")
            
            # Add missing columns
            for column_name, column_type, nullable, index_sql in missing_columns:
                if column_name not in existing_columns:
                    nullable_str = "NULL" if nullable else "NOT NULL"
                    alter_sql = f"""
                        ALTER TABLE evaluations 
                        ADD COLUMN {column_name} {column_type} {nullable_str}
                    """
                    
                    print(f"🔧 Adding column: {column_name}")
                    await db.execute(text(alter_sql))
                    
                    # Add index if specified
                    if index_sql:
                        print(f"🔧 Adding index for: {column_name}")
                        await db.execute(text(index_sql))
                else:
                    print(f"✅ Column already exists: {column_name}")
            
            await db.commit()
            print("✅ Migration completed successfully!")
            
            # Verify the migration
            result = await db.execute(text("""
                SELECT column_name 
                FROM information_schema.columns 
                WHERE table_name = 'evaluations'
            """))
            new_columns = {row[0] for row in result.fetchall()}
            print(f"📊 New total columns: {len(new_columns)}")
            
            return True
            
        except Exception as e:
            print(f"❌ Migration failed: {e}")
            await db.rollback()
            return False

if __name__ == "__main__":
    print("🚀 Starting database migration...")
    success = asyncio.run(run_migration())
    if success:
        print("🎉 Migration completed successfully!")
    else:
        print("❌ Migration failed!")