168 lines
5.1 KiB
Python
168 lines
5.1 KiB
Python
"""MongoDB client + collection provider.
|
|
|
|
Every query is (a) tenant-scoped and (b) backed by explicit indexes.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import asynccontextmanager
|
|
from typing import AsyncIterator
|
|
|
|
from motor.motor_asyncio import (
|
|
AsyncIOMotorClient,
|
|
AsyncIOMotorCollection,
|
|
AsyncIOMotorDatabase,
|
|
)
|
|
from pymongo import IndexModel
|
|
import structlog
|
|
|
|
from ..config import get_settings
|
|
|
|
log = structlog.get_logger(__name__)
|
|
|
|
|
|
_client: AsyncIOMotorClient | None = None
|
|
_db: AsyncIOMotorDatabase | None = None
|
|
|
|
|
|
class MongoCollections(str):
|
|
ORGANIZATIONS = "organizations"
|
|
USERS = "users"
|
|
TEAMS = "teams"
|
|
PROJECTS = "projects"
|
|
ARTIFACTS = "artifacts"
|
|
ARTIFACT_RELATIONSHIPS = "artifact_relationships"
|
|
INTEGRATIONS = "integrations"
|
|
JOBS = "jobs"
|
|
AI_EXECUTIONS = "ai_executions"
|
|
AUDIT_EVENTS = "audit_events"
|
|
|
|
|
|
def get_mongo_client() -> AsyncIOMotorClient:
|
|
global _client
|
|
if _client is None:
|
|
s = get_settings()
|
|
_client = AsyncIOMotorClient(
|
|
s.mongo_uri,
|
|
minPoolSize=s.mongo_min_pool_size,
|
|
maxPoolSize=s.mongo_max_pool_size,
|
|
uuidRepresentation="standard",
|
|
tz_aware=True,
|
|
)
|
|
return _client
|
|
|
|
|
|
def get_mongo_db() -> AsyncIOMotorDatabase:
|
|
global _db
|
|
if _db is None:
|
|
s = get_settings()
|
|
_db = get_mongo_client()[s.mongo_db]
|
|
return _db
|
|
|
|
|
|
def get_collection(name: MongoCollections | str) -> AsyncIOMotorCollection:
|
|
return get_mongo_db()[name]
|
|
|
|
|
|
async def create_indexes() -> None:
|
|
"""Create (idempotent) indexes. Called once on service startup.
|
|
|
|
Rule of thumb for tenanted data: compound index with ``tenant_id`` first,
|
|
then the most common filter / sort key.
|
|
"""
|
|
|
|
db = get_mongo_db()
|
|
|
|
# --- organizations ---
|
|
await db.organizations.create_indexes(
|
|
[
|
|
IndexModel("slug", unique=True, name="orgs_slug_unique"),
|
|
]
|
|
)
|
|
|
|
# --- users ---
|
|
await db.users.create_indexes(
|
|
[
|
|
IndexModel([("tenant_id", 1), ("email", 1)], unique=True, name="users_tenant_email_unique"),
|
|
]
|
|
)
|
|
|
|
# --- projects ---
|
|
await db.projects.create_indexes(
|
|
[
|
|
IndexModel([("tenant_id", 1), ("key", 1)], unique=True, name="projects_tenant_key_unique"),
|
|
IndexModel([("tenant_id", 1), ("is_active", 1)], name="projects_tenant_active"),
|
|
]
|
|
)
|
|
|
|
# --- artifacts ---
|
|
await db.artifacts.create_indexes(
|
|
[
|
|
IndexModel([("tenant_id", 1), ("project_id", 1)], name="arts_tenant_project"),
|
|
IndexModel([("tenant_id", 1), ("project_id", 1), ("type", 1)], name="arts_tenant_project_type"),
|
|
IndexModel([("tenant_id", 1), ("external_id", 1)], name="arts_tenant_external_id", sparse=True),
|
|
IndexModel([("tenant_id", 1), ("created_at", -1)], name="arts_tenant_created_desc"),
|
|
]
|
|
)
|
|
|
|
# --- artifact_relationships ---
|
|
# Two indexes: one for "where does this artifact point TO", one for "what points TO this artifact"
|
|
await db.artifact_relationships.create_indexes(
|
|
[
|
|
IndexModel([("tenant_id", 1), ("source_artifact_id", 1), ("relationship_type", 1)],
|
|
name="rels_tenant_source_type"),
|
|
IndexModel([("tenant_id", 1), ("target_artifact_id", 1), ("relationship_type", 1)],
|
|
name="rels_tenant_target_type"),
|
|
IndexModel([("source_artifact_id", 1), ("target_artifact_id", 1), ("relationship_type", 1)],
|
|
unique=True, name="rels_unique_edge"),
|
|
]
|
|
)
|
|
|
|
# --- jobs ---
|
|
await db.jobs.create_indexes(
|
|
[
|
|
IndexModel([("tenant_id", 1), ("status", 1), ("priority", -1), ("created_at", 1)],
|
|
name="jobs_tenant_status_priority"),
|
|
IndexModel([("tenant_id", 1), ("project_id", 1), ("created_at", -1)],
|
|
name="jobs_tenant_project_created_desc"),
|
|
IndexModel([("correlation_id", 1)], name="jobs_correlation_id", sparse=True),
|
|
]
|
|
)
|
|
|
|
# --- audit_events (append-only, oldest-to-newest scans normal) ---
|
|
await db.audit_events.create_indexes(
|
|
[
|
|
IndexModel([("tenant_id", 1), ("created_at", -1)], name="audit_tenant_created_desc"),
|
|
IndexModel([("tenant_id", 1), ("actor_id", 1), ("created_at", -1)], name="audit_tenant_actor_desc"),
|
|
IndexModel([("correlation_id", 1)], name="audit_correlation_id", sparse=True),
|
|
]
|
|
)
|
|
|
|
# --- ai_executions ---
|
|
await db.ai_executions.create_indexes(
|
|
[
|
|
IndexModel([("tenant_id", 1), ("created_at", -1)], name="aie_tenant_created_desc"),
|
|
IndexModel([("tenant_id", 1), ("job_id", 1)], name="aie_tenant_job", sparse=True),
|
|
]
|
|
)
|
|
|
|
await log.adebug("mongo_indexes_ready")
|
|
|
|
|
|
async def close_mongo() -> None:
|
|
global _client, _db
|
|
if _client is not None:
|
|
_client.close()
|
|
_client = None
|
|
_db = None
|
|
|
|
|
|
@asynccontextmanager
|
|
async def mongo_lifespan() -> AsyncIterator[None]:
|
|
"""Async lifespan helper for FastAPI startup / shutdown."""
|
|
await create_indexes()
|
|
try:
|
|
yield
|
|
finally:
|
|
await close_mongo()
|