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