Spaces:
Sleeping
Sleeping
File size: 10,973 Bytes
04aa1ba 838cd23 04aa1ba 838cd23 04aa1ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 | """
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()) |