backend/tests/create_learning_tables.py

"""
Create learning tables in PostgreSQL for the learning engine.

This creates the necessary tables for storing attack records and pattern metrics
in the main PostgreSQL database instead of SQLite.
"""

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 create_learning_tables():
    """Create learning engine tables in PostgreSQL."""
    
    tables = [
        # Learning attacks table
        """
        CREATE TABLE IF NOT EXISTS learning_attacks (
            attack_id VARCHAR(255) PRIMARY KEY,
            attack_type VARCHAR(100) NOT NULL,
            attack_category VARCHAR(100) NOT NULL,
            target_model VARCHAR(255) NOT NULL,
            dataset VARCHAR(255) NOT NULL,
            prompt TEXT NOT NULL,
            success BOOLEAN NOT NULL,
            safety_score REAL NOT NULL,
            risk_level VARCHAR(50) NOT NULL,
            response_text TEXT,
            response_length INTEGER,
            response_time_ms INTEGER,
            timestamp TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            metadata JSONB,
            created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
            updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
        );
        """,
        
        # Create indexes for performance
        """
        CREATE INDEX IF NOT EXISTS idx_learning_attacks_type ON learning_attacks(attack_type);
        """,
        
        """
        CREATE INDEX IF NOT EXISTS idx_learning_attacks_model ON learning_attacks(target_model);
        """,
        
        """
        CREATE INDEX IF NOT EXISTS idx_learning_attacks_timestamp ON learning_attacks(timestamp DESC);
        """,
        
        """
        CREATE INDEX IF NOT EXISTS idx_learning_attacks_success ON learning_attacks(success);
        """,
        
        # Create trigger for updated_at
        """
        CREATE OR REPLACE FUNCTION update_learning_attacks_updated_at()
        RETURNS TRIGGER AS $$
        BEGIN
            NEW.updated_at = NOW();
            RETURN NEW;
        END;
        $$ language plpgsql;
        """,
        
        """
        DROP TRIGGER IF EXISTS update_learning_attacks_updated_at_trigger ON learning_attacks;
        """,
        
        """
        CREATE TRIGGER update_learning_attacks_updated_at_trigger
            BEFORE UPDATE ON learning_attacks
            FOR EACH ROW
            EXECUTE FUNCTION update_learning_attacks_updated_at();
        """
    ]
    
    async with AsyncSessionLocal() as db:
        try:
            print("🔧 Creating learning engine tables in PostgreSQL...")
            
            for i, table_sql in enumerate(tables):
                print(f"🔧 Executing SQL statement {i+1}/{len(tables)}")
                await db.execute(text(table_sql))
            
            await db.commit()
            print("✅ Learning engine tables created successfully!")
            
            # Verify tables
            result = await db.execute(text("""
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_schema = 'public' 
                AND table_name LIKE 'learning_%'
            """))
            
            tables_created = [row[0] for row in result.fetchall()]
            print(f"📊 Tables created: {tables_created}")
            
            return True
            
        except Exception as e:
            print(f"❌ Failed to create learning tables: {e}")
            await db.rollback()
            return False

if __name__ == "__main__":
    print("🚀 Creating PostgreSQL learning tables...")
    success = asyncio.run(create_learning_tables())
    if success:
        print("🎉 Learning tables created successfully!")
    else:
        print("❌ Learning table creation failed!")