backend/tests/setup_database.py

"""
Simple Database Setup for AegisLM

Creates all necessary tables using existing working patterns.
"""

import asyncio
import sys
import os

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

async def setup_database():
    """Setup database tables."""
    try:
        from core.database import AsyncSessionLocal
        from db_models import User, Evaluation, ApiKey, Role, Permission
        
        print("🔧 Setting up database tables...")
        
        # Create tables using the same pattern as main.py
        async with AsyncSessionLocal() as db:
            # Create tables one by one to avoid issues
            from sqlalchemy import text
            
            # Create roles table first
            try:
                await db.execute(text("""
                    CREATE TABLE IF NOT EXISTS roles (
                        id SERIAL PRIMARY KEY,
                        name VARCHAR(100) UNIQUE NOT NULL,
                        slug VARCHAR(50) UNIQUE NOT NULL,
                        description TEXT,
                        level INTEGER DEFAULT 1,
                        is_active BOOLEAN DEFAULT TRUE,
                        is_system BOOLEAN DEFAULT FALSE,
                        created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                        updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                    )
                """))
                print("✅ Created roles table")
            except Exception as e:
                print(f"⚠️ Roles table already exists: {str(e)}")
            
            # Create users table enhancements
            try:
                await db.execute(text("""
                    ALTER TABLE users 
                    ADD COLUMN IF NOT EXISTS verified_at TIMESTAMP WITH TIME ZONE,
                    ADD COLUMN IF NOT EXISTS is_superuser BOOLEAN DEFAULT FALSE
                """))
                print("✅ Enhanced users table")
            except Exception as e:
                print(f"⚠️ Users table enhancement failed: {str(e)}")
            
            # Create api_keys table
            try:
                await db.execute(text("""
                    CREATE TABLE IF NOT EXISTS api_keys (
                        id SERIAL PRIMARY KEY,
                        key_name VARCHAR(100) NOT NULL,
                        key_prefix VARCHAR(10) NOT NULL,
                        key_hash VARCHAR(255) NOT NULL,
                        user_id INTEGER REFERENCES users(id),
                        is_active BOOLEAN DEFAULT TRUE,
                        expires_at TIMESTAMP WITH TIME ZONE,
                        last_used_at TIMESTAMP WITH TIME ZONE,
                        usage_count INTEGER DEFAULT 0,
                        rate_limit_per_minute INTEGER DEFAULT 60,
                        permissions TEXT,
                        scope VARCHAR(50) DEFAULT 'evaluation',
                        created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                        updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                    )
                """))
                print("✅ Created api_keys table")
            except Exception as e:
                print(f"⚠️ API keys table already exists: {str(e)}")
            
            # Create evaluations table
            try:
                await db.execute(text("""
                    CREATE TABLE IF NOT EXISTS evaluations (
                        id SERIAL PRIMARY KEY,
                        job_id VARCHAR(255) UNIQUE NOT NULL,
                        user_id INTEGER REFERENCES users(id),
                        model_name VARCHAR(100) NOT NULL,
                        model_type VARCHAR(50) NOT NULL,
                        api_endpoint VARCHAR(255),
                        api_key VARCHAR(255),
                        dataset_name VARCHAR(100),
                        status VARCHAR(50) DEFAULT 'pending',
                        progress INTEGER DEFAULT 0,
                        current_phase VARCHAR(100),
                        config_snapshot JSONB,
                        result_summary JSONB,
                        prompt_count INTEGER DEFAULT 0,
                        created_by VARCHAR(255),
                        tags TEXT,
                        created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                        updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                        started_at TIMESTAMP WITH TIME ZONE,
                        completed_at TIMESTAMP WITH TIME ZONE
                    )
                """))
                print("✅ Created evaluations table")
            except Exception as e:
                print(f"⚠️ Evaluations table already exists: {str(e)}")
                
            print("✅ Database setup completed successfully!")
                
    except Exception as e:
        print(f"❌ Database setup failed: {str(e)}")
        return False
    
    return True

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