c595ea47d7
- database.py: add connect_db/disconnect_db aliases, get_db_instance() sync getter, get_db() async dependency, and optional db arg on create_indexes() to match all call sites in main.py, celery_worker.py, and route modules - docker-compose.yml: remove host-port exposure for mongodb/redis to avoid conflicts with existing services on optical-dev (:27017, :6379 already in use) - .env.example: use Docker service-name hostnames (mongodb/redis) instead of localhost so the API container can reach them inside the Compose network Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
96 lines
2.8 KiB
Python
96 lines
2.8 KiB
Python
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
|
|
|
|
from .config import settings
|
|
from .logging import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class MongoDB:
|
|
client: AsyncIOMotorClient = None
|
|
database: AsyncIOMotorDatabase = None
|
|
|
|
|
|
mongodb = MongoDB()
|
|
|
|
|
|
async def connect_to_mongo() -> None:
|
|
logger.info("Connecting to MongoDB...")
|
|
mongodb.client = AsyncIOMotorClient(settings.mongodb_uri)
|
|
mongodb.database = mongodb.client[settings.mongodb_db]
|
|
try:
|
|
await mongodb.client.admin.command("ping")
|
|
logger.info("MongoDB connected")
|
|
except Exception as e:
|
|
logger.error(f"MongoDB connection failed: {e}")
|
|
raise
|
|
|
|
|
|
async def close_mongo_connection() -> None:
|
|
if mongodb.client:
|
|
mongodb.client.close()
|
|
|
|
|
|
async def get_database() -> AsyncIOMotorDatabase:
|
|
return mongodb.database
|
|
|
|
|
|
def get_db_instance() -> AsyncIOMotorDatabase:
|
|
return mongodb.database
|
|
|
|
|
|
async def get_db() -> AsyncIOMotorDatabase:
|
|
return mongodb.database
|
|
|
|
|
|
connect_db = connect_to_mongo
|
|
disconnect_db = close_mongo_connection
|
|
|
|
|
|
async def create_indexes(db: AsyncIOMotorDatabase = None) -> None:
|
|
if db is None:
|
|
db = mongodb.database
|
|
|
|
# workspaces
|
|
await db.workspaces.create_index([("slug", 1)], unique=True)
|
|
|
|
# teams
|
|
await db.teams.create_index([("workspace_id", 1), ("slug", 1)], unique=True)
|
|
|
|
# projects
|
|
await db.projects.create_index([("workspace_id", 1), ("slug", 1)], unique=True)
|
|
await db.projects.create_index([("source_app", 1), ("external_id", 1)], unique=True, sparse=True)
|
|
|
|
# users_mirror
|
|
await db.users_mirror.create_index([("source_app", 1), ("external_id", 1)], unique=True)
|
|
await db.users_mirror.create_index([("workspace_id", 1)])
|
|
|
|
# api_keys
|
|
await db.api_keys.create_index([("key_hash", 1)], unique=True)
|
|
await db.api_keys.create_index([("source_app", 1)])
|
|
|
|
# model_prices
|
|
await db.model_prices.create_index([("provider", 1), ("model", 1), ("effective_from", -1)])
|
|
|
|
# usage_events — primary analytics index
|
|
await db.usage_events.create_index([
|
|
("ts", -1), ("workspace_id", 1), ("project_id", 1),
|
|
("user_external_id", 1), ("model", 1)
|
|
])
|
|
await db.usage_events.create_index([("job_external_id", 1)])
|
|
await db.usage_events.create_index([("source_app", 1), ("ts", -1)])
|
|
|
|
# usage_rollups
|
|
await db.usage_rollups.create_index([
|
|
("date", -1), ("workspace_id", 1), ("project_id", 1)
|
|
])
|
|
await db.usage_rollups.create_index([("date", -1), ("workspace_id", 1), ("user_external_id", 1)])
|
|
|
|
# budgets
|
|
await db.budgets.create_index([("scope_type", 1), ("scope_id", 1)], unique=True)
|
|
|
|
# audit_log
|
|
await db.audit_log.create_index([("ts", -1)])
|
|
await db.audit_log.create_index([("actor_id", 1), ("ts", -1)])
|
|
|
|
logger.info("Indexes created")
|