findEthics/Atlas

1"""2Database index creation script for user authentication feature.3 4This script creates the necessary indexes for user_id fields to support5efficient user-specific queries while maintaining performance.6 7Indexes created:81. Sparse indexes on user_id fields (sessions, messages, search_analytics)92. Compound indexes on (user_id, timestamp) for user history queries10 11The script is designed to be safe to run multiple times and on existing data.12"""13 14import asyncio15import logging16from typing import List, Dict, Any17from dotenv import load_dotenv18from analytics.database import get_database, connect_to_database19 20# Load environment variables21load_dotenv()22 23logger = logging.getLogger(__name__)24 25async def create_user_id_indexes() -> bool:26    """27    Create all necessary indexes for user_id fields.28    29    Returns:30        bool: True if all indexes were created successfully, False otherwise31    """32    try:33        # Connect to database34        db = await get_database()35        if db is None:36            logger.error("Could not connect to database")37            return False38        39        logger.info("Starting index creation for user authentication feature...")40        41        # Define indexes to create42        indexes_to_create = [43            # Sessions collection indexes44            {45                "collection": "sessions",46                "indexes": [47                    {48                        "name": "user_id_sparse",49                        "keys": [("user_id", 1)],50                        "options": {"sparse": True, "background": True}51                    },52                    {53                        "name": "user_id_start_time_compound",54                        "keys": [("user_id", 1), ("start_time", -1)],55                        "options": {"sparse": True, "background": True}56                    }57                ]58            },59            # Messages collection indexes60            {61                "collection": "messages",62                "indexes": [63                    {64                        "name": "user_id_sparse",65                        "keys": [("user_id", 1)],66                        "options": {"sparse": True, "background": True}67                    },68                    {69                        "name": "user_id_timestamp_compound",70                        "keys": [("user_id", 1), ("timestamp", -1)],71                        "options": {"sparse": True, "background": True}72                    }73                ]74            },75            # Search analytics collection indexes76            {77                "collection": "search_analytics",78                "indexes": [79                    {80                        "name": "user_id_sparse",81                        "keys": [("user_id", 1)],82                        "options": {"sparse": True, "background": True}83                    },84                    {85                        "name": "user_id_timestamp_compound",86                        "keys": [("user_id", 1), ("timestamp", -1)],87                        "options": {"sparse": True, "background": True}88                    }89                ]90            }91        ]92        93        success_count = 094        total_indexes = sum(len(coll["indexes"]) for coll in indexes_to_create)95        96        # Create indexes for each collection97        for collection_config in indexes_to_create:98            collection_name = collection_config["collection"]99            collection = db[collection_name]100            101            logger.info(f"Creating indexes for {collection_name} collection...")102            103            for index_config in collection_config["indexes"]:104                try:105                    # Check if index already exists106                    existing_indexes = await collection.list_indexes().to_list(length=None)107                    index_names = [idx["name"] for idx in existing_indexes]108                    109                    if index_config["name"] in index_names:110                        logger.info(f"Index {index_config['name']} already exists on {collection_name}, skipping...")111                        success_count += 1112                        continue113                    114                    # Create the index115                    await collection.create_index(116                        index_config["keys"],117                        name=index_config["name"],118                        **index_config["options"]119                    )120                    121                    logger.info(f"Successfully created index {index_config['name']} on {collection_name}")122                    success_count += 1123                    124                except Exception as e:125                    logger.error(f"Failed to create index {index_config['name']} on {collection_name}: {e}")126        127        if success_count == total_indexes:128            logger.info(f"Successfully created all {total_indexes} indexes")129            return True130        else:131            logger.warning(f"Created {success_count}/{total_indexes} indexes")132            return False133            134    except Exception as e:135        logger.error(f"Error during index creation: {e}")136        return False137 138async def verify_indexes() -> bool:139    """140    Verify that all required indexes exist and are properly configured.141    142    Returns:143        bool: True if all indexes exist, False otherwise144    """145    try:146        db = await get_database()147        if db is None:148            logger.error("Could not connect to database for verification")149            return False150        151        logger.info("Verifying index creation...")152        153        # Expected indexes for each collection154        expected_indexes = {155            "sessions": ["user_id_sparse", "user_id_start_time_compound"],156            "messages": ["user_id_sparse", "user_id_timestamp_compound"],157            "search_analytics": ["user_id_sparse", "user_id_timestamp_compound"]158        }159        160        all_verified = True161        162        for collection_name, expected_index_names in expected_indexes.items():163            collection = db[collection_name]164            165            # Get existing indexes166            existing_indexes = await collection.list_indexes().to_list(length=None)167            existing_names = [idx["name"] for idx in existing_indexes]168            169            logger.info(f"Verifying indexes for {collection_name}:")170            171            for expected_name in expected_index_names:172                if expected_name in existing_names:173                    logger.info(f"  ✓ {expected_name} exists")174                else:175                    logger.error(f"  ✗ {expected_name} missing")176                    all_verified = False177        178        if all_verified:179            logger.info("All indexes verified successfully")180        else:181            logger.error("Some indexes are missing")182            183        return all_verified184        185    except Exception as e:186        logger.error(f"Error during index verification: {e}")187        return False188 189async def list_all_indexes() -> Dict[str, List[Dict[str, Any]]]:190    """191    List all indexes for analytics collections.192    193    Returns:194        Dict mapping collection names to their index information195    """196    try:197        db = await get_database()198        if db is None:199            logger.error("Could not connect to database")200            return {}201        202        collections = ["sessions", "messages", "search_analytics"]203        all_indexes = {}204        205        for collection_name in collections:206            collection = db[collection_name]207            indexes = await collection.list_indexes().to_list(length=None)208            all_indexes[collection_name] = indexes209            210            logger.info(f"Indexes for {collection_name}:")211            for idx in indexes:212                logger.info(f"  - {idx['name']}: {idx.get('key', 'N/A')}")213        214        return all_indexes215        216    except Exception as e:217        logger.error(f"Error listing indexes: {e}")218        return {}219 220async def drop_user_id_indexes() -> bool:221    """222    Drop all user_id related indexes (for rollback purposes).223    224    Returns:225        bool: True if all indexes were dropped successfully, False otherwise226    """227    try:228        db = await get_database()229        if db is None:230            logger.error("Could not connect to database")231            return False232        233        logger.info("Dropping user_id indexes for rollback...")234        235        # Indexes to drop236        indexes_to_drop = {237            "sessions": ["user_id_sparse", "user_id_start_time_compound"],238            "messages": ["user_id_sparse", "user_id_timestamp_compound"],239            "search_analytics": ["user_id_sparse", "user_id_timestamp_compound"]240        }241        242        success_count = 0243        total_indexes = sum(len(indexes) for indexes in indexes_to_drop.values())244        245        for collection_name, index_names in indexes_to_drop.items():246            collection = db[collection_name]247            248            for index_name in index_names:249                try:250                    await collection.drop_index(index_name)251                    logger.info(f"Dropped index {index_name} from {collection_name}")252                    success_count += 1253                except Exception as e:254                    # Index might not exist, which is fine for rollback255                    logger.warning(f"Could not drop index {index_name} from {collection_name}: {e}")256                    success_count += 1  # Count as success for rollback257        258        logger.info(f"Rollback completed: {success_count}/{total_indexes} indexes processed")259        return success_count == total_indexes260        261    except Exception as e:262        logger.error(f"Error during index rollback: {e}")263        return False264 265async def main():266    """Main function to create indexes"""267    logging.basicConfig(level=logging.INFO)268    269    try:270        # Connect to database271        await connect_to_database()272        273        # Create indexes274        success = await create_user_id_indexes()275        276        if success:277            # Verify indexes were created278            await verify_indexes()279            280            # List all indexes for confirmation281            await list_all_indexes()282            283            print("\n✅ Index creation completed successfully!")284            print("The following indexes have been created:")285            print("  - Sparse indexes on user_id fields for all collections")286            print("  - Compound indexes on (user_id, timestamp) for efficient user history queries")287            print("  - All indexes are created with background=True for minimal impact")288            289        else:290            print("\n❌ Index creation failed. Check logs for details.")291            292    except Exception as e:293        logger.error(f"Script execution failed: {e}")294        print(f"\n❌ Script failed: {e}")295 296if __name__ == "__main__":297    asyncio.run(main())