archive/migrate_user_authentication.py

#!/usr/bin/env python3
"""
Database migration script for user authentication feature.

This script safely adds user_id fields to existing collections and creates
the necessary indexes for efficient user-specific queries.

Usage:
    python migrate_user_authentication.py [command]

Commands:
    migrate   - Run full migration (default)
    rollback  - Rollback migration changes
    validate  - Validate migration success
    backup    - Create backup of current data
    status    - Check migration status

The script includes safety checks and can be run multiple times safely.
"""

import asyncio
from datetime import datetime
import json
import logging
import os
import sys
from typing import Any, Dict, List, Optional

from dotenv import load_dotenv

from analytics.create_indexes import create_user_id_indexes, drop_user_id_indexes, list_all_indexes, verify_indexes
from analytics.database import connect_to_database, get_database, test_connection

    create_user_id_indexes,
    verify_indexes,
    drop_user_id_indexes,
    list_all_indexes
)

logger = logging.getLogger(__name__)

class MigrationManager:
    """Manages database migration for user authentication feature"""
    
    def __init__(self):
        self.db = None
        self.backup_dir = "migration_backups"
        self.migration_collection = "migration_history"
        
    async def initialize(self):
        """Initialize database connection"""
        self.db = await get_database()
        if self.db is None:
            raise Exception("Could not connect to database")
        
        # Ensure backup directory exists
        os.makedirs(self.backup_dir, exist_ok=True)
        
    async def check_migration_status(self) -> Dict[str, Any]:
        """Check current migration status"""
        try:
            # Check if migration history collection exists
            collections = await self.db.list_collection_names()
            
            status = {
                "migration_history_exists": self.migration_collection in collections,
                "collections_exist": {
                    "sessions": "sessions" in collections,
                    "messages": "messages" in collections,
                    "search_analytics": "search_analytics" in collections
                },
                "user_id_fields_exist": {},
                "indexes_exist": {},
                "migration_completed": False
            }
            
            # Check if user_id fields exist in collections
            for collection_name in ["sessions", "messages", "search_analytics"]:
                if collection_name in collections:
                    collection = self.db[collection_name]
                    # Check if any document has user_id field
                    sample_doc = await collection.find_one({"user_id": {"$exists": True}})
                    status["user_id_fields_exist"][collection_name] = sample_doc is not None
                else:
                    status["user_id_fields_exist"][collection_name] = False
            
            # Check index status
            status["indexes_exist"] = await self._check_indexes_exist()
            
            # Check migration history
            if status["migration_history_exists"]:
                migration_coll = self.db[self.migration_collection]
                last_migration = await migration_coll.find_one(
                    {"migration_name": "user_authentication"},
                    sort=[("timestamp", -1)]
                )
                if last_migration and last_migration.get("status") == "completed":
                    status["migration_completed"] = True
                    status["last_migration"] = last_migration
            
            return status
            
        except Exception as e:
            logger.error(f"Error checking migration status: {e}")
            return {"error": str(e)}
    
    async def _check_indexes_exist(self) -> Dict[str, bool]:
        """Check if required indexes exist"""
        expected_indexes = {
            "sessions": ["user_id_sparse", "user_id_start_time_compound"],
            "messages": ["user_id_sparse", "user_id_timestamp_compound"],
            "search_analytics": ["user_id_sparse", "user_id_timestamp_compound"]
        }
        
        index_status = {}
        
        for collection_name, expected_names in expected_indexes.items():
            try:
                collection = self.db[collection_name]
                existing_indexes = await collection.list_indexes().to_list(length=None)
                existing_names = [idx["name"] for idx in existing_indexes]
                
                index_status[collection_name] = all(
                    name in existing_names for name in expected_names
                )
            except Exception:
                index_status[collection_name] = False
        
        return index_status
    
    async def create_backup(self) -> str:
        """Create backup of current data"""
        timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
        backup_file = os.path.join(self.backup_dir, f"backup_{timestamp}.json")
        
        logger.info(f"Creating backup: {backup_file}")
        
        try:
            backup_data = {
                "timestamp": timestamp,
                "collections": {}
            }
            
            # Backup each collection
            for collection_name in ["sessions", "messages", "search_analytics"]:
                try:
                    collection = self.db[collection_name]
                    documents = await collection.find({}).to_list(length=None)
                    
                    # Convert ObjectId and datetime to strings for JSON serialization
                    serializable_docs = []
                    for doc in documents:
                        serializable_doc = {}
                        for key, value in doc.items():
                            if isinstance(value, datetime):
                                serializable_doc[key] = value.isoformat()
                            else:
                                serializable_doc[key] = str(value) if hasattr(value, '__str__') else value
                        serializable_docs.append(serializable_doc)
                    
                    backup_data["collections"][collection_name] = {
                        "count": len(documents),
                        "documents": serializable_docs
                    }
                    
                    logger.info(f"Backed up {len(documents)} documents from {collection_name}")
                    
                except Exception as e:
                    logger.warning(f"Could not backup {collection_name}: {e}")
                    backup_data["collections"][collection_name] = {"error": str(e)}
            
            # Write backup file
            with open(backup_file, 'w') as f:
                json.dump(backup_data, f, indent=2)
            
            logger.info(f"Backup completed: {backup_file}")
            return backup_file
            
        except Exception as e:
            logger.error(f"Backup failed: {e}")
            raise
    
    async def run_migration(self) -> bool:
        """Run the complete migration process"""
        try:
            logger.info("Starting user authentication migration...")
            
            # Check current status
            status = await self.check_migration_status()
            if status.get("migration_completed"):
                logger.info("Migration already completed. Use 'validate' to verify.")
                return True
            
            # Create backup
            backup_file = await self.create_backup()
            
            # Record migration start
            await self._record_migration_event("started", {
                "backup_file": backup_file
            })
            
            # Step 1: Add user_id fields to existing documents (set to null)
            logger.info("Step 1: Adding user_id fields to existing documents...")
            field_success = await self._add_user_id_fields()
            
            if not field_success:
                await self._record_migration_event("failed", {"step": "add_fields"})
                return False
            
            # Step 2: Create indexes
            logger.info("Step 2: Creating user_id indexes...")
            index_success = await create_user_id_indexes()
            
            if not index_success:
                await self._record_migration_event("failed", {"step": "create_indexes"})
                return False
            
            # Step 3: Verify migration
            logger.info("Step 3: Verifying migration...")
            validation_success = await self.validate_migration()
            
            if not validation_success:
                await self._record_migration_event("failed", {"step": "validation"})
                return False
            
            # Record successful completion
            await self._record_migration_event("completed", {
                "backup_file": backup_file,
                "validation_passed": True
            })
            
            logger.info("Migration completed successfully!")
            return True
            
        except Exception as e:
            logger.error(f"Migration failed: {e}")
            await self._record_migration_event("failed", {"error": str(e)})
            return False
    
    async def _add_user_id_fields(self) -> bool:
        """Add user_id fields to existing documents"""
        try:
            collections_to_update = ["sessions", "messages", "search_analytics"]
            
            for collection_name in collections_to_update:
                logger.info(f"Adding user_id field to {collection_name}...")
                
                collection = self.db[collection_name]
                
                # Update documents that don't have user_id field
                result = await collection.update_many(
                    {"user_id": {"$exists": False}},
                    {"$set": {"user_id": None}}
                )
                
                logger.info(f"Updated {result.modified_count} documents in {collection_name}")
            
            return True
            
        except Exception as e:
            logger.error(f"Error adding user_id fields: {e}")
            return False
    
    async def _record_migration_event(self, status: str, details: Dict[str, Any]):
        """Record migration event in history"""
        try:
            migration_coll = self.db[self.migration_collection]
            
            event = {
                "migration_name": "user_authentication",
                "status": status,
                "timestamp": datetime.utcnow(),
                "details": details
            }
            
            await migration_coll.insert_one(event)
            
        except Exception as e:
            logger.warning(f"Could not record migration event: {e}")
    
    async def validate_migration(self) -> bool:
        """Validate that migration was successful"""
        try:
            logger.info("Validating migration...")
            
            validation_results = {
                "user_id_fields": True,
                "indexes": True,
                "data_integrity": True
            }
            
            # Check 1: Verify user_id fields exist
            for collection_name in ["sessions", "messages", "search_analytics"]:
                collection = self.db[collection_name]
                
                # Check if all documents have user_id field
                total_docs = await collection.count_documents({})
                docs_with_user_id = await collection.count_documents({"user_id": {"$exists": True}})
                
                if total_docs > 0 and docs_with_user_id != total_docs:
                    logger.error(f"Not all documents in {collection_name} have user_id field")
                    validation_results["user_id_fields"] = False
                else:
                    logger.info(f"✓ All {total_docs} documents in {collection_name} have user_id field")
            
            # Check 2: Verify indexes exist
            index_verification = await verify_indexes()
            validation_results["indexes"] = index_verification
            
            if index_verification:
                logger.info("✓ All required indexes exist")
            else:
                logger.error("✗ Some required indexes are missing")
            
            # Check 3: Data integrity checks
            integrity_check = await self._check_data_integrity()
            validation_results["data_integrity"] = integrity_check
            
            # Overall validation result
            all_passed = all(validation_results.values())
            
            if all_passed:
                logger.info("✅ Migration validation passed!")
            else:
                logger.error("❌ Migration validation failed!")
                logger.error(f"Results: {validation_results}")
            
            return all_passed
            
        except Exception as e:
            logger.error(f"Validation failed: {e}")
            return False
    
    async def _check_data_integrity(self) -> bool:
        """Check data integrity after migration"""
        try:
            # Check that existing data is preserved
            for collection_name in ["sessions", "messages", "search_analytics"]:
                collection = self.db[collection_name]
                
                # Sample a few documents to verify structure
                sample_docs = await collection.find({}).limit(5).to_list(length=5)
                
                for doc in sample_docs:
                    # Verify user_id field exists
                    if "user_id" not in doc:
                        logger.error(f"Document missing user_id field in {collection_name}: {doc.get('_id')}")
                        return False
                    
                    # Verify user_id is None for migrated documents (existing data)
                    # New documents with actual user_ids are fine
                    if doc["user_id"] is not None:
                        # This is fine - could be new data with actual user_ids
                        pass
            
            logger.info("✓ Data integrity checks passed")
            return True
            
        except Exception as e:
            logger.error(f"Data integrity check failed: {e}")
            return False
    
    async def rollback_migration(self) -> bool:
        """Rollback the migration changes"""
        try:
            logger.info("Starting migration rollback...")
            
            # Record rollback start
            await self._record_migration_event("rollback_started", {})
            
            # Step 1: Drop user_id indexes
            logger.info("Step 1: Dropping user_id indexes...")
            index_rollback = await drop_user_id_indexes()
            
            # Step 2: Remove user_id fields from documents
            logger.info("Step 2: Removing user_id fields from documents...")
            field_rollback = await self._remove_user_id_fields()
            
            if index_rollback and field_rollback:
                await self._record_migration_event("rollback_completed", {})
                logger.info("✅ Rollback completed successfully!")
                return True
            else:
                await self._record_migration_event("rollback_failed", {})
                logger.error("❌ Rollback failed!")
                return False
                
        except Exception as e:
            logger.error(f"Rollback failed: {e}")
            await self._record_migration_event("rollback_failed", {"error": str(e)})
            return False
    
    async def _remove_user_id_fields(self) -> bool:
        """Remove user_id fields from all documents"""
        try:
            collections_to_update = ["sessions", "messages", "search_analytics"]
            
            for collection_name in collections_to_update:
                logger.info(f"Removing user_id field from {collection_name}...")
                
                collection = self.db[collection_name]
                
                # Remove user_id field from all documents
                result = await collection.update_many(
                    {},
                    {"$unset": {"user_id": ""}}
                )
                
                logger.info(f"Updated {result.modified_count} documents in {collection_name}")
            
            return True
            
        except Exception as e:
            logger.error(f"Error removing user_id fields: {e}")
            return False

def setup_logging():
    """Setup logging configuration"""
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )

async def main():
    """Main CLI function"""
    setup_logging()
    
    # Get command from arguments
    command = sys.argv[1] if len(sys.argv) > 1 else "migrate"
    
    print(f"🚀 User Authentication Migration Tool")
    print(f"Command: {command}")
    print("-" * 50)
    
    try:
        # Test database connection first
        if not await test_connection():
            print("❌ Database connection failed. Check your MONGODB_URL.")
            sys.exit(1)
        
        # Initialize migration manager
        manager = MigrationManager()
        await manager.initialize()
        
        if command == "migrate":
            print("Running full migration...")
            success = await manager.run_migration()
            if not success:
                print("❌ Migration failed. Check logs for details.")
                sys.exit(1)
                
        elif command == "rollback":
            print("Rolling back migration...")
            success = await manager.rollback_migration()
            if not success:
                print("❌ Rollback failed. Check logs for details.")
                sys.exit(1)
                
        elif command == "validate":
            print("Validating migration...")
            success = await manager.validate_migration()
            if not success:
                print("❌ Validation failed. Check logs for details.")
                sys.exit(1)
                
        elif command == "backup":
            print("Creating backup...")
            backup_file = await manager.create_backup()
            print(f"✅ Backup created: {backup_file}")
            
        elif command == "status":
            print("Checking migration status...")
            status = await manager.check_migration_status()
            
            print("\n📊 Migration Status:")
            print(f"Migration completed: {status.get('migration_completed', False)}")
            print(f"Collections exist: {status.get('collections_exist', {})}")
            print(f"User ID fields exist: {status.get('user_id_fields_exist', {})}")
            print(f"Indexes exist: {status.get('indexes_exist', {})}")
            
            if status.get('last_migration'):
                last = status['last_migration']
                print(f"Last migration: {last.get('timestamp')} ({last.get('status')})")
                
        else:
            print(f"❌ Unknown command: {command}")
            print("Available commands: migrate, rollback, validate, backup, status")
            sys.exit(1)
            
        print("✅ Operation completed successfully!")
        
    except Exception as e:
        print(f"❌ Operation failed: {e}")
        sys.exit(1)

if __name__ == "__main__":
    asyncio.run(main())