""" Database index creation script for user authentication feature. This script creates the necessary indexes for user_id fields to support efficient user-specific queries while maintaining performance. Indexes created: 1. Sparse indexes on user_id fields (sessions, messages, search_analytics) 2. Compound indexes on (user_id, timestamp) for user history queries The script is designed to be safe to run multiple times and on existing data. """ import asyncio import logging from typing import List, Dict, Any from dotenv import load_dotenv from analytics.database import get_database, connect_to_database # Load environment variables load_dotenv() logger = logging.getLogger(__name__) async def create_user_id_indexes() -> bool: """ Create all necessary indexes for user_id fields. Returns: bool: True if all indexes were created successfully, False otherwise """ try: # Connect to database db = await get_database() if db is None: logger.error("Could not connect to database") return False logger.info("Starting index creation for user authentication feature...") # Define indexes to create indexes_to_create = [ # Sessions collection indexes { "collection": "sessions", "indexes": [ { "name": "user_id_sparse", "keys": [("user_id", 1)], "options": {"sparse": True, "background": True} }, { "name": "user_id_start_time_compound", "keys": [("user_id", 1), ("start_time", -1)], "options": {"sparse": True, "background": True} } ] }, # Messages collection indexes { "collection": "messages", "indexes": [ { "name": "user_id_sparse", "keys": [("user_id", 1)], "options": {"sparse": True, "background": True} }, { "name": "user_id_timestamp_compound", "keys": [("user_id", 1), ("timestamp", -1)], "options": {"sparse": True, "background": True} } ] }, # Search analytics collection indexes { "collection": "search_analytics", "indexes": [ { "name": "user_id_sparse", "keys": [("user_id", 1)], "options": {"sparse": True, "background": True} }, { "name": "user_id_timestamp_compound", "keys": [("user_id", 1), ("timestamp", -1)], "options": {"sparse": True, "background": True} } ] } ] success_count = 0 total_indexes = sum(len(coll["indexes"]) for coll in indexes_to_create) # Create indexes for each collection for collection_config in indexes_to_create: collection_name = collection_config["collection"] collection = db[collection_name] logger.info(f"Creating indexes for {collection_name} collection...") for index_config in collection_config["indexes"]: try: # Check if index already exists existing_indexes = await collection.list_indexes().to_list(length=None) index_names = [idx["name"] for idx in existing_indexes] if index_config["name"] in index_names: logger.info(f"Index {index_config['name']} already exists on {collection_name}, skipping...") success_count += 1 continue # Create the index await collection.create_index( index_config["keys"], name=index_config["name"], **index_config["options"] ) logger.info(f"Successfully created index {index_config['name']} on {collection_name}") success_count += 1 except Exception as e: logger.error(f"Failed to create index {index_config['name']} on {collection_name}: {e}") if success_count == total_indexes: logger.info(f"Successfully created all {total_indexes} indexes") return True else: logger.warning(f"Created {success_count}/{total_indexes} indexes") return False except Exception as e: logger.error(f"Error during index creation: {e}") return False async def verify_indexes() -> bool: """ Verify that all required indexes exist and are properly configured. Returns: bool: True if all indexes exist, False otherwise """ try: db = await get_database() if db is None: logger.error("Could not connect to database for verification") return False logger.info("Verifying index creation...") # Expected indexes for each collection 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"] } all_verified = True for collection_name, expected_index_names in expected_indexes.items(): collection = db[collection_name] # Get existing indexes existing_indexes = await collection.list_indexes().to_list(length=None) existing_names = [idx["name"] for idx in existing_indexes] logger.info(f"Verifying indexes for {collection_name}:") for expected_name in expected_index_names: if expected_name in existing_names: logger.info(f" ✓ {expected_name} exists") else: logger.error(f" ✗ {expected_name} missing") all_verified = False if all_verified: logger.info("All indexes verified successfully") else: logger.error("Some indexes are missing") return all_verified except Exception as e: logger.error(f"Error during index verification: {e}") return False async def list_all_indexes() -> Dict[str, List[Dict[str, Any]]]: """ List all indexes for analytics collections. Returns: Dict mapping collection names to their index information """ try: db = await get_database() if db is None: logger.error("Could not connect to database") return {} collections = ["sessions", "messages", "search_analytics"] all_indexes = {} for collection_name in collections: collection = db[collection_name] indexes = await collection.list_indexes().to_list(length=None) all_indexes[collection_name] = indexes logger.info(f"Indexes for {collection_name}:") for idx in indexes: logger.info(f" - {idx['name']}: {idx.get('key', 'N/A')}") return all_indexes except Exception as e: logger.error(f"Error listing indexes: {e}") return {} async def drop_user_id_indexes() -> bool: """ Drop all user_id related indexes (for rollback purposes). Returns: bool: True if all indexes were dropped successfully, False otherwise """ try: db = await get_database() if db is None: logger.error("Could not connect to database") return False logger.info("Dropping user_id indexes for rollback...") # Indexes to drop indexes_to_drop = { "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"] } success_count = 0 total_indexes = sum(len(indexes) for indexes in indexes_to_drop.values()) for collection_name, index_names in indexes_to_drop.items(): collection = db[collection_name] for index_name in index_names: try: await collection.drop_index(index_name) logger.info(f"Dropped index {index_name} from {collection_name}") success_count += 1 except Exception as e: # Index might not exist, which is fine for rollback logger.warning(f"Could not drop index {index_name} from {collection_name}: {e}") success_count += 1 # Count as success for rollback logger.info(f"Rollback completed: {success_count}/{total_indexes} indexes processed") return success_count == total_indexes except Exception as e: logger.error(f"Error during index rollback: {e}") return False async def main(): """Main function to create indexes""" logging.basicConfig(level=logging.INFO) try: # Connect to database await connect_to_database() # Create indexes success = await create_user_id_indexes() if success: # Verify indexes were created await verify_indexes() # List all indexes for confirmation await list_all_indexes() print("\n✅ Index creation completed successfully!") print("The following indexes have been created:") print(" - Sparse indexes on user_id fields for all collections") print(" - Compound indexes on (user_id, timestamp) for efficient user history queries") print(" - All indexes are created with background=True for minimal impact") else: print("\n❌ Index creation failed. Check logs for details.") except Exception as e: logger.error(f"Script execution failed: {e}") print(f"\n❌ Script failed: {e}") if __name__ == "__main__": asyncio.run(main())