272 lines
8.6 KiB
Python
272 lines
8.6 KiB
Python
|
#!/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)
|