archive/validate_user_migration.py

#!/usr/bin/env python3
"""
Data validation script for user authentication migration.

This script performs comprehensive validation to ensure the migration
was successful and data integrity is maintained.

Usage:
    python validate_user_migration.py [--detailed] [--fix-issues]

Options:
    --detailed    Show detailed validation results
    --fix-issues  Attempt to fix minor issues found during validation
"""

import argparse
import asyncio
from datetime import datetime
import logging
import sys
from typing import Any, Dict, List, Tuple

from analytics.create_indexes import verify_indexes
from analytics.database import get_database, test_connection

logger = logging.getLogger(__name__)

class MigrationValidator:
    """Comprehensive validation for user authentication migration"""
    
    def __init__(self, detailed: bool = False, fix_issues: bool = False):
        self.db = None
        self.detailed = detailed
        self.fix_issues = fix_issues
        self.validation_results = {}
        
    async def initialize(self):
        """Initialize database connection"""
        self.db = await get_database()
        if self.db is None:
            raise Exception("Could not connect to database")
    
    async def run_full_validation(self) -> Dict[str, Any]:
        """Run comprehensive validation suite"""
        logger.info("Starting comprehensive migration validation...")
        
        validation_suite = [
            ("Schema Validation", self._validate_schema),
            ("Index Validation", self._validate_indexes),
            ("Data Integrity", self._validate_data_integrity),
            ("Performance Check", self._validate_performance),
            ("Backward Compatibility", self._validate_backward_compatibility),
            ("Migration History", self._validate_migration_history)
        ]
        
        overall_results = {
            "validation_timestamp": datetime.utcnow().isoformat(),
            "tests_passed": 0,
            "tests_failed": 0,
            "tests_total": len(validation_suite),
            "details": {},
            "issues_found": [],
            "recommendations": []
        }
        
        for test_name, test_func in validation_suite:
            logger.info(f"Running {test_name}...")
            
            try:
                result = await test_func()
                overall_results["details"][test_name] = result
                
                if result.get("passed", False):
                    overall_results["tests_passed"] += 1
                    if self.detailed:
                        print(f"✅ {test_name}: PASSED")
                else:
                    overall_results["tests_failed"] += 1
                    print(f"❌ {test_name}: FAILED")
                    
                    # Collect issues
                    if "issues" in result:
                        overall_results["issues_found"].extend(result["issues"])
                    
                    # Collect recommendations
                    if "recommendations" in result:
                        overall_results["recommendations"].extend(result["recommendations"])
                
                if self.detailed and "details" in result:
                    for detail in result["details"]:
                        print(f"   {detail}")
                        
            except Exception as e:
                logger.error(f"{test_name} failed with exception: {e}")
                overall_results["tests_failed"] += 1
                overall_results["details"][test_name] = {
                    "passed": False,
                    "error": str(e)
                }
        
        # Calculate success rate
        success_rate = (overall_results["tests_passed"] / overall_results["tests_total"]) * 100
        overall_results["success_rate"] = success_rate
        
        return overall_results
    
    async def _validate_schema(self) -> Dict[str, Any]:
        """Validate database schema changes"""
        result = {
            "passed": True,
            "details": [],
            "issues": [],
            "recommendations": []
        }
        
        collections_to_check = ["sessions", "messages", "search_analytics"]
        
        for collection_name in collections_to_check:
            try:
                collection = self.db[collection_name]
                
                # Check if collection exists
                collections = await self.db.list_collection_names()
                if collection_name not in collections:
                    result["issues"].append(f"Collection {collection_name} does not exist")
                    result["passed"] = False
                    continue
                
                # Check user_id field presence
                total_docs = await collection.count_documents({})
                docs_with_user_id = await collection.count_documents({"user_id": {"$exists": True}})
                
                if total_docs > 0:
                    if docs_with_user_id == 0:
                        result["issues"].append(f"No documents in {collection_name} have user_id field")
                        result["passed"] = False
                    elif docs_with_user_id < total_docs:
                        missing_count = total_docs - docs_with_user_id
                        result["issues"].append(f"{missing_count} documents in {collection_name} missing user_id field")
                        result["passed"] = False
                        
                        if self.fix_issues:
                            # Fix missing user_id fields
                            fix_result = await collection.update_many(
                                {"user_id": {"$exists": False}},
                                {"$set": {"user_id": None}}
                            )
                            result["details"].append(f"Fixed {fix_result.modified_count} documents in {collection_name}")
                    else:
                        result["details"].append(f"✓ All {total_docs} documents in {collection_name} have user_id field")
                else:
                    result["details"].append(f"✓ {collection_name} is empty (no validation needed)")
                
                # Check user_id field types
                sample_docs = await collection.find({"user_id": {"$ne": None}}).limit(10).to_list(length=10)
                for doc in sample_docs:
                    user_id = doc.get("user_id")
                    if user_id is not None and not isinstance(user_id, str):
                        result["issues"].append(f"Invalid user_id type in {collection_name}: {type(user_id)}")
                        result["passed"] = False
                
            except Exception as e:
                result["issues"].append(f"Error validating {collection_name}: {e}")
                result["passed"] = False
        
        return result
    
    async def _validate_indexes(self) -> Dict[str, Any]:
        """Validate index creation"""
        result = {
            "passed": True,
            "details": [],
            "issues": [],
            "recommendations": []
        }
        
        # Use existing index verification
        indexes_valid = await verify_indexes()
        
        if indexes_valid:
            result["details"].append("✓ All required indexes exist")
        else:
            result["passed"] = False
            result["issues"].append("Some required indexes are missing")
            result["recommendations"].append("Run: python scripts/deployment/create_user_indexes.py create")
        
        # Additional index performance checks
        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"]
        }
        
        for collection_name, expected_names in expected_indexes.items():
            try:
                collection = self.db[collection_name]
                indexes = await collection.list_indexes().to_list(length=None)
                
                for index in indexes:
                    if index["name"] in expected_names:
                        # Check if index is sparse (important for null user_id values)
                        if index["name"].endswith("_sparse") and not index.get("sparse", False):
                            result["issues"].append(f"Index {index['name']} should be sparse")
                            result["passed"] = False
                        
                        result["details"].append(f"✓ Index {index['name']} configured correctly")
                        
            except Exception as e:
                result["issues"].append(f"Error checking indexes for {collection_name}: {e}")
                result["passed"] = False
        
        return result
    
    async def _validate_data_integrity(self) -> Dict[str, Any]:
        """Validate data integrity after migration"""
        result = {
            "passed": True,
            "details": [],
            "issues": [],
            "recommendations": []
        }
        
        try:
            # Check session-message consistency
            sessions_coll = self.db["sessions"]
            messages_coll = self.db["messages"]
            
            # Sample sessions with user_id
            sessions_with_user_id = await sessions_coll.find({"user_id": {"$ne": None}}).limit(10).to_list(length=10)
            
            for session in sessions_with_user_id:
                session_id = session["session_id"]
                session_user_id = session["user_id"]
                
                # Check messages for this session
                messages = await messages_coll.find({"session_id": session_id}).to_list(length=None)
                
                for message in messages:
                    message_user_id = message.get("user_id")
                    
                    # User IDs should match (or message user_id can be None for backward compatibility)
                    if message_user_id is not None and message_user_id != session_user_id:
                        result["issues"].append(
                            f"User ID mismatch: session {session_id} has user_id '{session_user_id}' "
                            f"but message {message['message_id']} has user_id '{message_user_id}'"
                        )
                        result["passed"] = False
            
            # Check for orphaned messages (messages without corresponding sessions)
            pipeline = [
                {
                    "$lookup": {
                        "from": "sessions",
                        "localField": "session_id",
                        "foreignField": "session_id",
                        "as": "session"
                    }
                },
                {
                    "$match": {
                        "session": {"$size": 0}
                    }
                },
                {
                    "$count": "orphaned_messages"
                }
            ]
            
            orphaned_result = await messages_coll.aggregate(pipeline).to_list(length=1)
            if orphaned_result and orphaned_result[0]["orphaned_messages"] > 0:
                orphaned_count = orphaned_result[0]["orphaned_messages"]
                result["issues"].append(f"Found {orphaned_count} orphaned messages without corresponding sessions")
                result["recommendations"].append("Review and clean up orphaned messages")
            else:
                result["details"].append("✓ No orphaned messages found")
            
            # Check user_id format validation
            for collection_name in ["sessions", "messages", "search_analytics"]:
                collection = self.db[collection_name]
                
                # Check for invalid user_id formats
                invalid_user_ids = await collection.find({
                    "user_id": {
                        "$ne": None,
                        "$not": {"$regex": "^[a-zA-Z0-9_-]+$"}
                    }
                }).to_list(length=10)
                
                if invalid_user_ids:
                    result["issues"].append(f"Found {len(invalid_user_ids)} documents with invalid user_id format in {collection_name}")
                    result["passed"] = False
                else:
                    result["details"].append(f"✓ All user_id values in {collection_name} have valid format")
            
        except Exception as e:
            result["issues"].append(f"Data integrity check failed: {e}")
            result["passed"] = False
        
        return result
    
    async def _validate_performance(self) -> Dict[str, Any]:
        """Validate query performance with new indexes"""
        result = {
            "passed": True,
            "details": [],
            "issues": [],
            "recommendations": []
        }
        
        try:
            # Test user_id query performance
            collections_to_test = ["sessions", "messages", "search_analytics"]
            
            for collection_name in collections_to_test:
                collection = self.db[collection_name]
                
                # Find a user_id to test with
                sample_doc = await collection.find_one({"user_id": {"$ne": None}})
                
                if sample_doc and sample_doc.get("user_id"):
                    user_id = sample_doc["user_id"]
                    
                    # Test query with explain
                    explain_result = await collection.find({"user_id": user_id}).explain()
                    
                    # Check if index was used
                    execution_stats = explain_result.get("executionStats", {})
                    if execution_stats.get("totalDocsExamined", 0) > execution_stats.get("totalDocsReturned", 0) * 2:
                        result["issues"].append(f"Query performance concern in {collection_name}: examining too many documents")
                        result["recommendations"].append(f"Review index usage for {collection_name}")
                    else:
                        result["details"].append(f"✓ Query performance good for {collection_name}")
                else:
                    result["details"].append(f"✓ No user_id data to test in {collection_name}")
            
        except Exception as e:
            result["issues"].append(f"Performance validation failed: {e}")
            result["passed"] = False
        
        return result
    
    async def _validate_backward_compatibility(self) -> Dict[str, Any]:
        """Validate backward compatibility for anonymous users"""
        result = {
            "passed": True,
            "details": [],
            "issues": [],
            "recommendations": []
        }
        
        try:
            # Check that anonymous sessions (user_id = null) still work
            sessions_coll = self.db["sessions"]
            messages_coll = self.db["messages"]
            
            # Count anonymous sessions
            anonymous_sessions = await sessions_coll.count_documents({"user_id": None})
            total_sessions = await sessions_coll.count_documents({})
            
            if total_sessions > 0:
                anonymous_percentage = (anonymous_sessions / total_sessions) * 100
                result["details"].append(f"✓ {anonymous_sessions}/{total_sessions} sessions are anonymous ({anonymous_percentage:.1f}%)")
                
                # Verify anonymous sessions have all required fields
                sample_anonymous = await sessions_coll.find_one({"user_id": None})
                if sample_anonymous:
                    required_fields = ["session_id", "start_time", "message_count", "search_used", "status"]
                    missing_fields = [field for field in required_fields if field not in sample_anonymous]
                    
                    if missing_fields:
                        result["issues"].append(f"Anonymous sessions missing fields: {missing_fields}")
                        result["passed"] = False
                    else:
                        result["details"].append("✓ Anonymous sessions have all required fields")
            else:
                result["details"].append("✓ No sessions to validate (empty database)")
            
            # Check anonymous messages
            anonymous_messages = await messages_coll.count_documents({"user_id": None})
            total_messages = await messages_coll.count_documents({})
            
            if total_messages > 0:
                anonymous_msg_percentage = (anonymous_messages / total_messages) * 100
                result["details"].append(f"✓ {anonymous_messages}/{total_messages} messages are anonymous ({anonymous_msg_percentage:.1f}%)")
            
        except Exception as e:
            result["issues"].append(f"Backward compatibility check failed: {e}")
            result["passed"] = False
        
        return result
    
    async def _validate_migration_history(self) -> Dict[str, Any]:
        """Validate migration history and tracking"""
        result = {
            "passed": True,
            "details": [],
            "issues": [],
            "recommendations": []
        }
        
        try:
            # Check if migration history collection exists
            collections = await self.db.list_collection_names()
            
            if "migration_history" not in collections:
                result["issues"].append("Migration history collection not found")
                result["recommendations"].append("Migration tracking is not available")
                # This is not a critical failure
            else:
                migration_coll = self.db["migration_history"]
                
                # Check for user authentication migration records
                user_auth_migrations = await migration_coll.find(
                    {"migration_name": "user_authentication"}
                ).sort("timestamp", -1).to_list(length=10)
                
                if not user_auth_migrations:
                    result["issues"].append("No user authentication migration records found")
                    result["recommendations"].append("Migration may not have been run through the migration script")
                else:
                    latest_migration = user_auth_migrations[0]
                    result["details"].append(f"✓ Latest migration: {latest_migration['status']} at {latest_migration['timestamp']}")
                    
                    if latest_migration["status"] != "completed":
                        result["issues"].append(f"Latest migration status is '{latest_migration['status']}', not 'completed'")
                        result["passed"] = False
        
        except Exception as e:
            result["issues"].append(f"Migration history validation failed: {e}")
            result["passed"] = False
        
        return result

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"""
    parser = argparse.ArgumentParser(description="Validate user authentication migration")
    parser.add_argument("--detailed", action="store_true", help="Show detailed validation results")
    parser.add_argument("--fix-issues", action="store_true", help="Attempt to fix minor issues")
    
    args = parser.parse_args()
    
    setup_logging()
    
    print(f"🔍 User Authentication Migration Validator")
    print(f"Detailed mode: {args.detailed}")
    print(f"Fix issues: {args.fix_issues}")
    print("-" * 50)
    
    try:
        # Test database connection
        if not await test_connection():
            print("❌ Database connection failed. Check your MONGODB_URL.")
            sys.exit(1)
        
        # Initialize validator
        validator = MigrationValidator(detailed=args.detailed, fix_issues=args.fix_issues)
        await validator.initialize()
        
        # Run validation
        results = await validator.run_full_validation()
        
        # Print summary
        print(f"\n📊 Validation Summary:")
        print(f"Tests passed: {results['tests_passed']}/{results['tests_total']}")
        print(f"Success rate: {results['success_rate']:.1f}%")
        
        if results["issues_found"]:
            print(f"\n⚠️  Issues found ({len(results['issues_found'])}):")
            for issue in results["issues_found"]:
                print(f"  - {issue}")
        
        if results["recommendations"]:
            print(f"\n💡 Recommendations ({len(results['recommendations'])}):")
            for rec in results["recommendations"]:
                print(f"  - {rec}")
        
        # Exit with appropriate code
        if results["tests_failed"] == 0:
            print("\n✅ All validation tests passed!")
            sys.exit(0)
        else:
            print(f"\n❌ {results['tests_failed']} validation tests failed!")
            sys.exit(1)
            
    except Exception as e:
        print(f"❌ Validation failed: {e}")
        sys.exit(1)

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