server/scripts/migrate_v2.py

import os
import sys
from sqlalchemy import create_engine, text

# Add parent directory to path
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.append(parent_dir)

from app.core.config import settings
from app.core.database import Base, engine
# Import all models to register them
from app.models import *

def migrate_v2():
    print("🚀 Starting V2 Migration...")
    
    # 1. Create new tables (Tenants, AuditLogs, EmergencyAlerts)
    print("Creating new tables...")
    Base.metadata.create_all(bind=engine)
    print("✅ New tables created (if they didn't exist)")
    
    with engine.connect() as conn:
        # 2. Add tenant_id to Users
        print("Checking Users table...")
        try:
            conn.execute(text("ALTER TABLE users ADD COLUMN tenant_id INTEGER REFERENCES tenants(id)"))
            print("✅ Added tenant_id to users")
        except Exception as e:
            print(f"⚠️  tenant_id in users might already exist: {e}")
            
        # 3. Add tenant_id to Websites
        print("Checking Websites table...")
        try:
            conn.execute(text("ALTER TABLE websites ADD COLUMN tenant_id INTEGER REFERENCES tenants(id)"))
            print("✅ Added tenant_id to websites")
        except Exception as e:
            print(f"⚠️  tenant_id in websites might already exist: {e}")

        # 4. Enable RLS on ChatSessions (Postgres only)
        # Note: ensuring chat_sessions table exists first (Base.metadata.create_all should have handled it if it was new, but it's existing)
        # Verify db is postgres
        if 'postgres' in settings.DATABASE_URL:
            print("Enabling Row-Level Security on chat_sessions...")
            try:
                conn.execute(text("ALTER TABLE chat_sessions ENABLE ROW LEVEL SECURITY"))
                
                # Create policy
                # Note: current_setting('app.current_tenant_id') needs to be set in middleware
                # For now, just enabling RLS. Policy creation might fail if policy exists.
                conn.execute(text("""
                    DO $$
                    BEGIN
                        IF NOT EXISTS (
                            SELECT 1 FROM pg_policies 
                            WHERE tablename = 'chat_sessions' 
                            AND policyname = 'tenant_isolation_policy'
                        ) THEN
                            CREATE POLICY tenant_isolation_policy ON chat_sessions
                            USING (website_id IN (
                                SELECT id FROM websites 
                                WHERE tenant_id = current_setting('app.current_tenant_id', true)::integer
                            ));
                        END IF;
                    END
                    $$;
                """))
                print("✅ RLS enabled and policy created")
            except Exception as e:
                print(f"⚠️  RLS setup failed (might be SQLite?): {e}")
        else:
            print("Skipping RLS (Not PostgreSQL)")

        conn.commit()
    
    print("✅ Migration Complete!")

if __name__ == "__main__":
    migrate_v2()