File size: 3,588 Bytes
2ed8996 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """
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!")
|