Spaces:
Sleeping
Sleeping
| """ | |
| 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()) |