scripts/migrate_participants_key.py

#!/usr/bin/env python3
"""
Migration script: Add participants_key to all existing conversations
and fix the unique index
"""

import asyncio
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from app.database import connect_db, disconnect_db, get_db
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


async def migrate_conversations():
    """Add participants_key to all existing conversations"""
    
    try:
        await connect_db()
        db = await get_db()
        
        # Step 1: Drop the problematic unique index on participants array
        logger.info("Step 1: Dropping old participants index...")
        try:
            await db.conversations.drop_index("participants_1")
            logger.info("  ✅ Dropped participants_1 index")
        except Exception as e:
            logger.info(f"  Index doesn't exist: {e}")
        
        # Step 2: Add participants_key to all existing conversations
        logger.info("\nStep 2: Adding participants_key to existing conversations...")
        conversations = await db.conversations.find({}).to_list(length=None)
        logger.info(f"Found {len(conversations)} conversations to migrate")
        
        migrated = 0
        for conv in conversations:
            participants = conv.get("participants", [])
            
            # Check if participants is valid
            if not isinstance(participants, list) or len(participants) != 2:
                logger.warning(f"  ⚠️ Skipping invalid conversation {conv['_id']}: {participants}")
                continue
            
            # Create participants_key
            sorted_participants = sorted(participants)
            participants_key = "::".join(sorted_participants)
            
            # Update the document
            await db.conversations.update_one(
                {"_id": conv["_id"]},
                {"$set": {
                    "participants": sorted_participants,  # Ensure sorted
                    "participants_key": participants_key
                }}
            )
            migrated += 1
            logger.info(f"  ✅ Migrated: {conv['_id']} -> {participants_key}")
        
        # Step 3: Create new unique index on participants_key
        logger.info("\nStep 3: Creating new unique index on participants_key...")
        try:
            await db.conversations.drop_index("participants_key_1")
        except:
            pass
        
        await db.conversations.create_index("participants_key", unique=True)
        logger.info("  ✅ Created unique index on participants_key")
        
        # Step 4: Create non-unique index on participants for queries
        logger.info("\nStep 4: Creating non-unique index on participants for queries...")
        await db.conversations.create_index("participants")
        logger.info("  ✅ Created non-unique index on participants")
        
        logger.info(f"\n{'='*60}")
        logger.info(f"✅ MIGRATION COMPLETE!")
        logger.info(f"Migrated {migrated} conversations")
        logger.info(f"New unique index on participants_key")
        logger.info(f"{'='*60}\n")
        
    except Exception as e:
        logger.error(f"Migration failed: {e}")
        raise
    finally:
        await disconnect_db()


if __name__ == "__main__":
    logger.info("🔄 MIGRATING CONVERSATIONS TO USE participants_key...")
    asyncio.run(migrate_conversations())