data_studio/db.py

from __future__ import annotations

from pymongo import ASCENDING, DESCENDING, MongoClient
from pymongo.database import Database

from data_studio.config import StudioSettings


def create_client(settings: StudioSettings) -> MongoClient:
    return MongoClient(settings.mongodb_uri)


def get_database(client: MongoClient, settings: StudioSettings) -> Database:
    return client[settings.mongodb_db_name]


def ensure_indexes(db: Database) -> None:
    db["source_configs"].create_index([("dataset_key", ASCENDING)], unique=True)
    db["users"].create_index([("username", ASCENDING)], unique=True)
    db["samples"].create_index([("sample_id", ASCENDING)], unique=True)
    db["samples"].create_index(
        [("source.dataset_key", ASCENDING), ("source.source_id", ASCENDING)],
        unique=True,
    )
    db["samples"].create_index([("queue_name", ASCENDING), ("status", ASCENDING)])
    db["samples"].create_index([("queue_name", ASCENDING), ("lease.expires_at", ASCENDING)])
    db["samples"].create_index([("updated_at", DESCENDING)])
    db["review_events"].create_index([("sample_id", ASCENDING), ("created_at", DESCENDING)])
    db["sessions"].create_index([("token_hash", ASCENDING)], unique=True)
    db["sessions"].create_index([("expires_at", ASCENDING)], expireAfterSeconds=0)
    db["import_jobs"].create_index([("dataset_key", ASCENDING), ("created_at", DESCENDING)])
    db["import_checkpoints"].create_index([("dataset_key", ASCENDING)], unique=True)
    db["exports"].create_index([("created_at", DESCENDING)])
    db["dataset_versions"].create_index([("version", ASCENDING)], unique=True)
    db["source_registry"].create_index([("source_id", ASCENDING), ("hf_revision", ASCENDING)], unique=True)
    db["raw_records"].create_index([("record_id", ASCENDING)], unique=True)
    db["raw_records"].create_index([("source_id", ASCENDING), ("source_revision", ASCENDING)])
    db["governance_events"].create_index([("entity_type", ASCENDING), ("entity_id", ASCENDING), ("occurred_at", DESCENDING)])
    db["protected_benchmarks"].create_index([("benchmark_id", ASCENDING), ("exact_revision", ASCENDING)], unique=True)
    db["rubrics"].create_index([("rubric_id", ASCENDING), ("rubric_version", ASCENDING)], unique=True)
    db["sampling_manifests"].create_index([("sample_id", ASCENDING)], unique=True)
    db["cleaner_versions"].create_index([("cleaner_version", ASCENDING)], unique=True)
    db["jobs"].create_index([("idempotency_key", ASCENDING)], unique=True)
    db["bulk_batches"].create_index([("batch_id", ASCENDING)], unique=True)
    db["audit_samples"].create_index([("audit_sample_id", ASCENDING)], unique=True)
    db["audit_results"].create_index([("audit_sample_id", ASCENDING), ("audited_by", ASCENDING)], unique=True)
    db["promotion_decisions"].create_index([("dataset_version", ASCENDING), ("occurred_at", DESCENDING)])
    db["tokenizer_versions"].create_index([("tokenizer_version", ASCENDING)], unique=True)
    db["training_exports"].create_index([("export_version", ASCENDING)], unique=True)