g_solution / app /core /db.py
minhvtt's picture
Upload 20 files
94b52f1 verified
Raw
History Blame Contribute Delete
2.09 kB
from datetime import timezone
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
from pymongo.errors import OperationFailure
from app.core.config import settings
_client: AsyncIOMotorClient | None = None
def get_db() -> AsyncIOMotorDatabase:
global _client
if _client is None:
_client = AsyncIOMotorClient(
settings.mongo_uri,
tz_aware=True,
tzinfo=timezone.utc,
)
return _client[settings.mongo_db]
async def ensure_indexes() -> None:
db = get_db()
async def safe_create_index(collection, keys, **options) -> None:
"""Create index and ignore already-existing/conflicting definitions."""
try:
await collection.create_index(keys, **options)
except OperationFailure as exc:
# Atlas/PyMongo can raise conflicts when an index with same name/key
# already exists but with different options (e.g., unique vs non-unique).
if exc.code in {85, 86}:
return
raise
await safe_create_index(db.devices, "device_id", unique=True)
await safe_create_index(db.policies, "device_id", unique=True)
await safe_create_index(db.key_states, "device_id", unique=True)
await safe_create_index(db.key_states, "expires_at")
await safe_create_index(db.key_states, "lock_until")
await safe_create_index(db.violations, "device_id")
await safe_create_index(db.violations, "timestamp")
await safe_create_index(db.violations, "expires_at")
await safe_create_index(db.screenshot_audit, "created_at")
await safe_create_index(db.alert_audit, "created_at")
await safe_create_index(db.alert_audit, "type")
await safe_create_index(db.url_verdicts, "url", unique=True)
await safe_create_index(db.url_verdicts, "verdict")
await safe_create_index(db.game_lock_reviews, "review_id", unique=True)
await safe_create_index(db.game_lock_reviews, "status")
await safe_create_index(db.game_lock_reviews, "created_at")