New enterprise feature
This commit is contained in:
Urtzi Alfaro
271
scripts/fix_existing_demo_sessions.py
Executable file
271
scripts/fix_existing_demo_sessions.py
Executable file
@@ -0,0 +1,271 @@
|
||||
#!/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)
|
||||
Reference in New Issue
Block a user