trauma/api/data/db_requests.py

import asyncio
import json
import re

from fastapi import HTTPException

from trauma.api.data.model import EntityModel
from trauma.api.data.schemas import SearchRequest
from trauma.core.config import settings


async def get_facility_by_id(facility_id: str) -> EntityModel:
    facility = await settings.DB_CLIENT.entities.find_one({"id": facility_id}, {"embedding": 0})
    if not facility:
        raise HTTPException(status_code=404, detail="Country with specified id doesn't exists.")
    return EntityModel.from_mongo(facility)


async def get_all_model_obj() -> list[EntityModel]:
    sort_v = -1
    objects = await settings.DB_CLIENT.entities.find({}, {"embedding": 0}).sort("_id", sort_v).to_list(length=None)
    return objects


async def search_facilities_obj(data: SearchRequest) -> tuple[list[EntityModel], int]:
    skip = data.pageSize * data.pageIndex
    sort_v = -1
    regex_filter = {"name": {"$regex": f"^{re.escape(data.name)}", "$options": "i"}}
    objects, total_count = await asyncio.gather(
        settings.DB_CLIENT.entities
        .find(regex_filter, {"embeddings": 0})
        .sort("_id", sort_v)
        .skip(skip)
        .limit(data.pageSize)
        .to_list(length=data.pageSize),
        settings.DB_CLIENT.entities.count_documents(regex_filter)
    )
    return [EntityModel.from_mongo(ent) for ent in objects], total_count


async def check_instructions():
    pipeline = [
        {"$unwind": "$contactDetails.postalCode"},
        {"$group": {"_id": None, "uniqueTreatmentMethods": {"$addToSet": "$contactDetails.postalCode"}}}
    ]
    result = await settings.DB_CLIENT.entities.aggregate(pipeline).to_list(length=1)
    if result:
        t = result[0]["uniqueTreatmentMethods"]
        with open('test.json', 'w') as f:
            f.write(json.dumps({"t": t}, indent=2))
    return []

if __name__ == "__main__":
    asyncio.run(check_instructions())