prediqai / scripts /setup_indexes.py
ganesh-vilje's picture
Deploy to Hugging Face Main
f8f02c0
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())