Spaces:
Runtime error
Runtime error
File size: 6,291 Bytes
f8f02c0 | 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 | import asyncio
from motor.motor_asyncio import AsyncIOMotorClient
import os
from app.core.config import Config
MONGODB_URI = Config.MONGODB_URI
DB_NAME = Config.MONGODB_DATABASE
async def main():
client = AsyncIOMotorClient(MONGODB_URI, tz_aware=True)
database = client[DB_NAME]
print("[ASYNC_DATABASE] Starting collection and index creation...")
# List existing collections before creation
existing_collections = await database.list_collection_names()
print(f"[ASYNC_DATABASE] Existing collections before creation: {existing_collections}")
# Models collection
print("[ASYNC_DATABASE] Creating models collection...")
models_collection = database["models"]
await models_collection.create_index([("model_id", 1)], unique=True)
await models_collection.create_index([("user_id", 1), ("model_name", 1)])
await models_collection.create_index([("user_id", 1), ("tag", 1)])
await models_collection.create_index([("status", 1)])
await models_collection.create_index([("organization_id", 1)])
await models_collection.create_index([("department_id", 1)])
await models_collection.create_index([("created_at", -1)])
print("[ASYNC_DATABASE] Models collection and indexes created")
# Training runs collection
print("[ASYNC_DATABASE] Creating training_runs collection...")
training_runs_collection = database["training_runs"]
await training_runs_collection.create_index([("training_id", 1)], unique=True)
await training_runs_collection.create_index([("model_id", 1)])
await training_runs_collection.create_index([("user_id", 1)])
await training_runs_collection.create_index([("status", 1)])
await training_runs_collection.create_index([("created_at", -1)])
await training_runs_collection.create_index([("compliance_type", 1)])
print("[ASYNC_DATABASE] Training runs collection and indexes created")
# Validation runs collection
print("[ASYNC_DATABASE] Creating validation_runs collection...")
validation_runs_collection = database["validation_runs"]
await validation_runs_collection.create_index([("validation_id", 1)], unique=True)
await validation_runs_collection.create_index([("training_id", 1)])
await validation_runs_collection.create_index([("model_id", 1)])
await validation_runs_collection.create_index([("model_name", 1)])
await validation_runs_collection.create_index([("user_id", 1)])
await validation_runs_collection.create_index([("created_at", -1)])
await validation_runs_collection.create_index([("compliance_type", 1)])
print("[ASYNC_DATABASE] Validation runs collection and indexes created")
# Prediction runs collection
print("[ASYNC_DATABASE] Creating prediction_runs collection...")
prediction_runs_collection = database["prediction_runs"]
await prediction_runs_collection.create_index([("prediction_id", 1)], unique=True)
await prediction_runs_collection.create_index([("training_id", 1)])
await prediction_runs_collection.create_index([("model_id", 1)])
await prediction_runs_collection.create_index([("model_name", 1)])
await prediction_runs_collection.create_index([("user_id", 1)])
await prediction_runs_collection.create_index([("created_at", -1)])
await prediction_runs_collection.create_index([("compliance_type", 1)])
print("[ASYNC_DATABASE] Prediction runs collection and indexes created")
# System status collection
print("[ASYNC_DATABASE] Creating system_status collection...")
status_collection = database["system_status"]
await status_collection.create_index([("training_id", 1)], unique=True)
await status_collection.create_index([("user_id", 1)])
await status_collection.create_index([("status", 1)])
await status_collection.create_index([("created_at", -1)])
print("[ASYNC_DATABASE] System status collection and indexes created")
# -----------------------------------------------------------------
# Organizations collection
# -----------------------------------------------------------------
print("[ASYNC_DATABASE] Creating organizations collection...")
organizations = database["organizations"]
await organizations.create_index([("name", 1)])
print("[ASYNC_DATABASE] Organizations collection and indexes created")
# -----------------------------------------------------------------
# Departments collection
# -----------------------------------------------------------------
print("[ASYNC_DATABASE] Creating departments collection...")
departments = database["departments"]
await departments.create_index([("name", 1)])
# Foreign key style index to organization
await departments.create_index([("organization_id", 1)])
print("[ASYNC_DATABASE] Departments collection and indexes created")
# -----------------------------------------------------------------
# Users collection
# -----------------------------------------------------------------
print("[ASYNC_DATABASE] Creating users collection...")
users = database["users"]
await users.create_index([("username", 1)])
# Foreign key style references
await users.create_index([("organization_id", 1)])
await users.create_index([("department_id", 1)])
# Admin type filter
await users.create_index([("admin_type", 1)])
print("[ASYNC_DATABASE] Users collection and indexes created")
# -----------------------------------------------------------------
# Metadata collection
# -----------------------------------------------------------------
print("[ASYNC_DATABASE] Creating Metadata collection...")
meta_data = database["metadata"]
# await meta_data.create_index([("metadata_id", 1)])
# await meta_data.create_index([("feature_data", 1)])
print("[ASYNC_DATABASE] Metadata collection and indexes created")
# Verify collections were created
final_collections = await database.list_collection_names()
print(f"[ASYNC_DATABASE] Final collections after creation: {final_collections}")
print("[ASYNC_DATABASE] All collections and indexes created successfully")
print("[INDEX] All indexes created successfully.")
client.close()
if __name__ == "__main__":
asyncio.run(main())
|