#!/usr/bin/env python3 """ One-time data migration script to populate demo_session_id for existing virtual tenants. This script fixes existing demo sessions created before the demo_session_id fix was implemented. It links tenants to their sessions using DemoSession.virtual_tenant_id and session_metadata.child_tenant_ids. Usage: python3 scripts/fix_existing_demo_sessions.py Requirements: - Both demo_session and tenant services must be accessible - Database credentials must be available via environment variables """ import asyncio import asyncpg import json import os import sys from datetime import datetime from typing import List, Dict, Any from uuid import UUID # Database connection URLs DEMO_SESSION_DB_URL = os.getenv( "DEMO_SESSION_DATABASE_URL", "postgresql://demo_session_user:demo_password@localhost:5432/demo_session_db" ) TENANT_DB_URL = os.getenv( "TENANT_DATABASE_URL", "postgresql://tenant_user:T0uJnXs0r4TUmxSQeQ2DuQGP6HU0LEba@localhost:5432/tenant_db" ) async def get_all_demo_sessions(demo_session_conn) -> List[Dict[str, Any]]: """Fetch all demo sessions from demo_session database""" query = """ SELECT id, session_id, virtual_tenant_id, demo_account_type, session_metadata, status, created_at FROM demo_sessions WHERE status IN ('ready', 'active', 'partial') ORDER BY created_at DESC """ rows = await demo_session_conn.fetch(query) sessions = [] for row in rows: sessions.append({ "id": row["id"], "session_id": row["session_id"], "virtual_tenant_id": row["virtual_tenant_id"], "demo_account_type": row["demo_account_type"], "session_metadata": row["session_metadata"], "status": row["status"], "created_at": row["created_at"] }) return sessions async def check_tenant_exists(tenant_conn, tenant_id: UUID) -> bool: """Check if a tenant exists in the tenant database""" query = """ SELECT id FROM tenants WHERE id = $1 AND is_demo = true """ result = await tenant_conn.fetchrow(query, tenant_id) return result is not None async def update_tenant_session_id(tenant_conn, tenant_id: UUID, session_id: str): """Update a tenant's demo_session_id""" query = """ UPDATE tenants SET demo_session_id = $2 WHERE id = $1 AND is_demo = true """ await tenant_conn.execute(query, tenant_id, session_id) async def get_tenant_session_id(tenant_conn, tenant_id: UUID) -> str: """Get the current demo_session_id for a tenant""" query = """ SELECT demo_session_id FROM tenants WHERE id = $1 AND is_demo = true """ result = await tenant_conn.fetchrow(query, tenant_id) return result["demo_session_id"] if result else None async def migrate_demo_sessions(): """Main migration function""" print("=" * 80) print("Demo Session Migration Script") print("=" * 80) print(f"Started at: {datetime.now()}") print() # Connect to both databases print("Connecting to databases...") demo_session_conn = await asyncpg.connect(DEMO_SESSION_DB_URL) tenant_conn = await asyncpg.connect(TENANT_DB_URL) print("✓ Connected to both databases") print() try: # Fetch all demo sessions print("Fetching demo sessions...") sessions = await get_all_demo_sessions(demo_session_conn) print(f"✓ Found {len(sessions)} demo sessions") print() # Statistics stats = { "sessions_processed": 0, "tenants_updated": 0, "tenants_already_set": 0, "tenants_not_found": 0, "errors": 0 } # Process each session for session in sessions: session_id = session["session_id"] virtual_tenant_id = session["virtual_tenant_id"] demo_account_type = session["demo_account_type"] session_metadata = session["session_metadata"] or {} print(f"Processing session: {session_id}") print(f" Type: {demo_account_type}") print(f" Main tenant: {virtual_tenant_id}") tenant_ids_to_update = [virtual_tenant_id] # For enterprise sessions, also get child tenant IDs if demo_account_type in ["enterprise_chain", "enterprise_parent"]: child_tenant_ids = session_metadata.get("child_tenant_ids", []) if child_tenant_ids: # Convert string UUIDs to UUID objects child_uuids = [UUID(tid) if isinstance(tid, str) else tid for tid in child_tenant_ids] tenant_ids_to_update.extend(child_uuids) print(f" Child tenants: {len(child_uuids)}") # Update each tenant session_tenants_updated = 0 for tenant_id in tenant_ids_to_update: try: # Check if tenant exists exists = await check_tenant_exists(tenant_conn, tenant_id) if not exists: print(f" ⚠ Tenant {tenant_id} not found - skipping") stats["tenants_not_found"] += 1 continue # Check current session_id current_session_id = await get_tenant_session_id(tenant_conn, tenant_id) if current_session_id == session_id: print(f" ✓ Tenant {tenant_id} already has session_id set") stats["tenants_already_set"] += 1 continue # Update the tenant await update_tenant_session_id(tenant_conn, tenant_id, session_id) print(f" ✓ Updated tenant {tenant_id}") stats["tenants_updated"] += 1 session_tenants_updated += 1 except Exception as e: print(f" ✗ Error updating tenant {tenant_id}: {e}") stats["errors"] += 1 stats["sessions_processed"] += 1 print(f" Session complete: {session_tenants_updated} tenant(s) updated") print() # Print summary print("=" * 80) print("Migration Complete!") print("=" * 80) print(f"Sessions processed: {stats['sessions_processed']}") print(f"Tenants updated: {stats['tenants_updated']}") print(f"Tenants already set: {stats['tenants_already_set']}") print(f"Tenants not found: {stats['tenants_not_found']}") print(f"Errors: {stats['errors']}") print() print(f"Finished at: {datetime.now()}") print("=" * 80) # Return success status return stats["errors"] == 0 except Exception as e: print(f"✗ Migration failed with error: {e}") import traceback traceback.print_exc() return False finally: # Close connections await demo_session_conn.close() await tenant_conn.close() print("Database connections closed") async def verify_migration(): """Verify that the migration was successful""" print() print("=" * 80) print("Verification Check") print("=" * 80) tenant_conn = await asyncpg.connect(TENANT_DB_URL) try: # Count tenants without session_id query = """ SELECT COUNT(*) as count FROM tenants WHERE is_demo = true AND demo_session_id IS NULL """ result = await tenant_conn.fetchrow(query) null_count = result["count"] if null_count == 0: print("✓ All demo tenants have demo_session_id set") else: print(f"⚠ {null_count} demo tenant(s) still have NULL demo_session_id") print(" These may be template tenants or orphaned records") # Count tenants with session_id query2 = """ SELECT COUNT(*) as count FROM tenants WHERE is_demo = true AND demo_session_id IS NOT NULL """ result2 = await tenant_conn.fetchrow(query2) set_count = result2["count"] print(f"✓ {set_count} demo tenant(s) have demo_session_id set") print("=" * 80) print() finally: await tenant_conn.close() if __name__ == "__main__": # Run migration success = asyncio.run(migrate_demo_sessions()) # Run verification if success: asyncio.run(verify_migration()) sys.exit(0) else: print("Migration failed - see errors above") sys.exit(1)