#!/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())