diff --git a/backend/infra/permission/role_handler.py b/backend/infra/permission/role_handler.py index 08d5d83..2fe4be0 100644 --- a/backend/infra/permission/role_handler.py +++ b/backend/infra/permission/role_handler.py @@ -118,7 +118,7 @@ class RoleHandler: raise RequestValidationError("Role with the provided ID already exists.") new_doc.id = custom_role_id - await new_doc.insert() + await new_doc.create() return new_doc async def query_roles(self, role_key: Optional[str], role_name: Optional[str], skip: int = 0, limit: int = 10) -> \ diff --git a/backend/infra/permission/user_role_handler.py b/backend/infra/permission/user_role_handler.py index 2b02fe8..96bd9fa 100644 --- a/backend/infra/permission/user_role_handler.py +++ b/backend/infra/permission/user_role_handler.py @@ -36,7 +36,7 @@ class UserRoleHandler: user_id=user_id, role_ids=unique_role_ids ) - await user_role_doc.insert() + await user_role_doc.create() return user_role_doc async def get_role_and_permission_by_user_id(self, user_id: str) -> tuple[list[str], list[str]]: diff --git a/backend/models/base_doc.py b/backend/models/base_doc.py index bd6722d..328076b 100644 --- a/backend/models/base_doc.py +++ b/backend/models/base_doc.py @@ -3,6 +3,7 @@ BaseDoc - A custom document class that provides Beanie-like interface using dire """ import asyncio from datetime import datetime, timezone +from bson import ObjectId from typing import Optional, List, Dict, Any, Type, Union from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase from pydantic import BaseModel @@ -282,6 +283,14 @@ class BaseDoc(BaseModel, metaclass=QueryModelMeta): # Convert Decimal objects to float for MongoDB compatibility doc_dict = self._convert_decimals_to_float(doc_dict) + # Respect pre-populated id by mapping to MongoDB _id + id_value = getattr(self, 'id', None) + if id_value and not isinstance(id_value, QueryExpression): + try: + doc_dict['_id'] = ObjectId(id_value) + except Exception: + doc_dict['_id'] = id_value + result = await collection.insert_one(doc_dict) # Set the id field from the inserted document @@ -314,7 +323,18 @@ class BaseDoc(BaseModel, metaclass=QueryModelMeta): elif hasattr(self, 'auth_code'): query['auth_code'] = self.auth_code - if query: + id_value = getattr(self, 'id', None) + if id_value and not isinstance(id_value, QueryExpression): + # Update by primary key when available + try: + object_id = ObjectId(id_value) + except Exception: + object_id = id_value + + result = await collection.update_one({"_id": object_id}, {"$set": doc_dict}, upsert=True) + if result.upserted_id: + self.id = str(result.upserted_id) + elif query: # Update existing document result = await collection.update_one(query, {"$set": doc_dict}, upsert=True) # If it was an insert, set the id field