migrate_database.py

#!/usr/bin/env python3
"""
Database migration script to add missing columns to users table
Run this once to update your Neon database schema
"""

import os
from sqlalchemy import create_engine, text
from src.core.config import settings

def migrate_database():
    """Add missing columns to users table"""
    print("🔄 Starting database migration...")
    print(f"Database URL: {settings.SQLALCHEMY_DATABASE_URI[:50]}...")
    
    engine = create_engine(settings.SQLALCHEMY_DATABASE_URI)
    
    migrations = [
        # Add role column if it doesn't exist
        """
        DO $$ 
        BEGIN
            IF NOT EXISTS (SELECT 1 FROM information_schema.columns 
                          WHERE table_name='users' AND column_name='role') THEN
                ALTER TABLE users ADD COLUMN role VARCHAR(20) DEFAULT 'user';
                UPDATE users SET role = 'user' WHERE role IS NULL;
                RAISE NOTICE 'Added role column';
            ELSE
                RAISE NOTICE 'role column already exists';
            END IF;
        END $$;
        """,
        
        # Add is_active column if it doesn't exist
        """
        DO $$ 
        BEGIN
            IF NOT EXISTS (SELECT 1 FROM information_schema.columns 
                          WHERE table_name='users' AND column_name='is_active') THEN
                ALTER TABLE users ADD COLUMN is_active BOOLEAN DEFAULT TRUE;
                UPDATE users SET is_active = TRUE WHERE is_active IS NULL;
                RAISE NOTICE 'Added is_active column';
            ELSE
                RAISE NOTICE 'is_active column already exists';
            END IF;
        END $$;
        """,
        
        # Add full_name column if it doesn't exist
        """
        DO $$ 
        BEGIN
            IF NOT EXISTS (SELECT 1 FROM information_schema.columns 
                          WHERE table_name='users' AND column_name='full_name') THEN
                ALTER TABLE users ADD COLUMN full_name VARCHAR(255);
                RAISE NOTICE 'Added full_name column';
            ELSE
                RAISE NOTICE 'full_name column already exists';
            END IF;
        END $$;
        """,
        
        # Create refresh_tokens table if it doesn't exist
        """
        CREATE TABLE IF NOT EXISTS refresh_tokens (
            id SERIAL PRIMARY KEY,
            user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
            token VARCHAR(500) NOT NULL UNIQUE,
            expires_at TIMESTAMP NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            revoked BOOLEAN DEFAULT FALSE
        );
        """,
        
        # Create index on refresh_tokens
        """
        CREATE INDEX IF NOT EXISTS idx_refresh_tokens_user_id ON refresh_tokens(user_id);
        CREATE INDEX IF NOT EXISTS idx_refresh_tokens_token ON refresh_tokens(token);
        """,
    ]
    
    try:
        with engine.connect() as conn:
            for i, migration in enumerate(migrations, 1):
                print(f"\n📝 Running migration {i}/{len(migrations)}...")
                conn.execute(text(migration))
                conn.commit()
                print(f"✅ Migration {i} completed")
        
        print("\n✅ All migrations completed successfully!")
        print("\n🎉 Database schema is now up to date")
        return True
        
    except Exception as e:
        print(f"\n❌ Migration failed: {e}")
        return False
    finally:
        engine.dispose()

if __name__ == "__main__":
    success = migrate_database()
    exit(0 if success else 1)