backend/tests/fix_user_schema.py

"""
User Schema Migration Script for AegisLM

This script adds missing columns to users 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_user_migration():
    """Add missing columns to users table."""
    
    missing_columns = [
        # User verification timestamp
        ("verified_at", "TIMESTAMP WITH TIME ZONE", True, "CREATE INDEX IF NOT EXISTS idx_users_verified_at ON users(verified_at)"),
        
        # User superuser flag
        ("is_superuser", "BOOLEAN DEFAULT FALSE", False, None),
    ]
    
    try:
        async with AsyncSessionLocal() as db:
            # Check current columns
            result = await db.execute(text("""
                SELECT column_name 
                FROM information_schema.columns 
                WHERE table_name = 'users'
            """))
            existing_columns = {row[0] for row in result.fetchall()}
            
            print(f"📊 Current columns: {len(existing_columns)}")
            print(f"🔧 Missing user 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:
                    print(f"➕ Adding column: {column_name}")
                    try:
                        # Add column with proper NULL/NOT NULL handling
                        null_clause = "NULL" if nullable else "NOT NULL"
                        await db.execute(text(f"""
                            ALTER TABLE users 
                            ADD COLUMN {column_name} {column_type} {null_clause}
                        """))
                        
                        # Add index if specified
                        if index_sql:
                            await db.execute(text(index_sql))
                            
                        print(f"✅ Added column: {column_name}")
                        
                    except Exception as e:
                        print(f"❌ Failed to add column {column_name}: {str(e)}")
                else:
                    print(f"✅ Column already exists: {column_name}")
            
            await db.commit()
            print("✅ User schema migration completed successfully")
            
    except Exception as e:
        print(f"⚠️ Migration failed: {str(e)}")
        print("💡 This is expected if PostgreSQL is not running locally")
        print("💡 Migration will work when database is available")
        # Don't raise exception - continue with other fixes

if __name__ == "__main__":
    asyncio.run(run_user_migration())