""" 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())