from datetime import datetime, timezone from bson import ObjectId from fastapi import HTTPException, status from motor.motor_asyncio import AsyncIOMotorDatabase from pymongo import ReturnDocument from pymongo.errors import DuplicateKeyError from app.schemas.records import ( ChildCreateRequest, ChildRecordResponse, ChildUpdateRequest, ParentCreateRequest, ParentRecordResponse, ParentUpdateRequest, ) PARENTS_COLLECTION = "parents" CHILDREN_COLLECTION = "children" class RecordNotFoundError(HTTPException): def __init__(self, detail: str = "Record not found.") -> None: super().__init__(status_code=status.HTTP_404_NOT_FOUND, detail=detail) class BadRequestError(HTTPException): def __init__(self, detail: str = "Invalid request.") -> None: super().__init__(status_code=status.HTTP_400_BAD_REQUEST, detail=detail) def _now() -> datetime: return datetime.now(timezone.utc) def _to_object_id(value: str, field_name: str) -> ObjectId: if not ObjectId.is_valid(value): raise BadRequestError(f"Invalid {field_name}.") return ObjectId(value) def _serialize_parent(document: dict) -> ParentRecordResponse: return ParentRecordResponse( id=str(document["_id"]), ownerId=document["ownerId"], iv=document["iv"], ciphertext=document["ciphertext"], createdAt=document["createdAt"], updatedAt=document["updatedAt"], ) def _serialize_child(document: dict) -> ChildRecordResponse: return ChildRecordResponse( id=str(document["_id"]), parentId=str(document["parentId"]), ownerId=document["ownerId"], iv=document["iv"], ciphertext=document["ciphertext"], createdAt=document["createdAt"], updatedAt=document["updatedAt"], ) async def ensure_record_indexes(database: AsyncIOMotorDatabase) -> None: await database[PARENTS_COLLECTION].create_index([("ownerId", 1)], name="ix_parents_ownerId") await database[CHILDREN_COLLECTION].create_index([("ownerId", 1)], name="ix_children_ownerId") await database[CHILDREN_COLLECTION].create_index([("parentId", 1)], name="ix_children_parentId") await database[CHILDREN_COLLECTION].create_index( [("ownerId", 1), ("parentId", 1)], name="ix_children_ownerId_parentId", ) async def create_parent(database: AsyncIOMotorDatabase, owner_id: str, payload: ParentCreateRequest) -> ParentRecordResponse: now = _now() document = { "ownerId": owner_id, "iv": payload.iv, "ciphertext": payload.ciphertext, "createdAt": now, "updatedAt": now, } result = await database[PARENTS_COLLECTION].insert_one(document) document["_id"] = result.inserted_id return _serialize_parent(document) async def list_parents(database: AsyncIOMotorDatabase, owner_id: str) -> list[ParentRecordResponse]: parents: list[ParentRecordResponse] = [] cursor = database[PARENTS_COLLECTION].find({"ownerId": owner_id}).sort("createdAt", -1) async for document in cursor: parents.append(_serialize_parent(document)) return parents async def get_parent(database: AsyncIOMotorDatabase, owner_id: str, parent_id: str) -> ParentRecordResponse: document = await database[PARENTS_COLLECTION].find_one( {"_id": _to_object_id(parent_id, "parentId"), "ownerId": owner_id} ) if document is None: raise RecordNotFoundError("Parent record not found.") return _serialize_parent(document) async def update_parent( database: AsyncIOMotorDatabase, owner_id: str, parent_id: str, payload: ParentUpdateRequest, ) -> ParentRecordResponse: document = await database[PARENTS_COLLECTION].find_one_and_update( {"_id": _to_object_id(parent_id, "parentId"), "ownerId": owner_id}, { "$set": { "iv": payload.iv, "ciphertext": payload.ciphertext, "updatedAt": _now(), } }, return_document=ReturnDocument.AFTER, ) if document is None: raise RecordNotFoundError("Parent record not found.") return _serialize_parent(document) async def delete_parent(database: AsyncIOMotorDatabase, owner_id: str, parent_id: str) -> None: client = database.client if client is None: raise RuntimeError("MongoDB client has not been initialized") parent_object_id = _to_object_id(parent_id, "parentId") async with await client.start_session() as session: async with session.start_transaction(): parent = await database[PARENTS_COLLECTION].find_one( {"_id": parent_object_id, "ownerId": owner_id}, session=session, ) if parent is None: raise RecordNotFoundError("Parent record not found.") await database[CHILDREN_COLLECTION].delete_many( {"parentId": parent_object_id, "ownerId": owner_id}, session=session, ) await database[PARENTS_COLLECTION].delete_one( {"_id": parent_object_id, "ownerId": owner_id}, session=session, ) async def create_child( database: AsyncIOMotorDatabase, owner_id: str, payload: ChildCreateRequest, ) -> ChildRecordResponse: parent_object_id = _to_object_id(payload.parentId, "parentId") parent = await database[PARENTS_COLLECTION].find_one({"_id": parent_object_id, "ownerId": owner_id}) if parent is None: raise RecordNotFoundError("Parent record not found.") now = _now() document = { "ownerId": owner_id, "parentId": parent_object_id, "iv": payload.iv, "ciphertext": payload.ciphertext, "createdAt": now, "updatedAt": now, } result = await database[CHILDREN_COLLECTION].insert_one(document) document["_id"] = result.inserted_id return _serialize_child(document) async def list_children(database: AsyncIOMotorDatabase, owner_id: str, parent_id: str) -> list[ChildRecordResponse]: parent_object_id = _to_object_id(parent_id, "parentId") children: list[ChildRecordResponse] = [] cursor = database[CHILDREN_COLLECTION].find({"ownerId": owner_id, "parentId": parent_object_id}).sort("createdAt", -1) async for document in cursor: children.append(_serialize_child(document)) return children async def get_child(database: AsyncIOMotorDatabase, owner_id: str, parent_id: str, child_id: str) -> ChildRecordResponse: document = await database[CHILDREN_COLLECTION].find_one( { "_id": _to_object_id(child_id, "childId"), "parentId": _to_object_id(parent_id, "parentId"), "ownerId": owner_id, } ) if document is None: raise RecordNotFoundError("Child record not found.") return _serialize_child(document) async def update_child( database: AsyncIOMotorDatabase, owner_id: str, parent_id: str, child_id: str, payload: ChildUpdateRequest, ) -> ChildRecordResponse: document = await database[CHILDREN_COLLECTION].find_one_and_update( { "_id": _to_object_id(child_id, "childId"), "parentId": _to_object_id(parent_id, "parentId"), "ownerId": owner_id, }, { "$set": { "iv": payload.iv, "ciphertext": payload.ciphertext, "updatedAt": _now(), } }, return_document=ReturnDocument.AFTER, ) if document is None: raise RecordNotFoundError("Child record not found.") return _serialize_child(document) async def delete_child(database: AsyncIOMotorDatabase, owner_id: str, parent_id: str, child_id: str) -> None: result = await database[CHILDREN_COLLECTION].delete_one( { "_id": _to_object_id(child_id, "childId"), "parentId": _to_object_id(parent_id, "parentId"), "ownerId": owner_id, } ) if result.deleted_count == 0: raise RecordNotFoundError("Child record not found.")