fix: change some un-exist insert methods to create

This commit is contained in:
haolou 2025-11-07 17:29:44 +08:00
parent 2409fdbde0
commit 9f8892a8a9
3 changed files with 21 additions and 3 deletions

View File

@ -118,7 +118,7 @@ class RoleHandler:
raise RequestValidationError("Role with the provided ID already exists.") raise RequestValidationError("Role with the provided ID already exists.")
new_doc.id = custom_role_id new_doc.id = custom_role_id
await new_doc.insert() await new_doc.create()
return new_doc return new_doc
async def query_roles(self, role_key: Optional[str], role_name: Optional[str], skip: int = 0, limit: int = 10) -> \ async def query_roles(self, role_key: Optional[str], role_name: Optional[str], skip: int = 0, limit: int = 10) -> \

View File

@ -36,7 +36,7 @@ class UserRoleHandler:
user_id=user_id, user_id=user_id,
role_ids=unique_role_ids role_ids=unique_role_ids
) )
await user_role_doc.insert() await user_role_doc.create()
return user_role_doc return user_role_doc
async def get_role_and_permission_by_user_id(self, user_id: str) -> tuple[list[str], list[str]]: async def get_role_and_permission_by_user_id(self, user_id: str) -> tuple[list[str], list[str]]:

View File

@ -3,6 +3,7 @@ BaseDoc - A custom document class that provides Beanie-like interface using dire
""" """
import asyncio import asyncio
from datetime import datetime, timezone from datetime import datetime, timezone
from bson import ObjectId
from typing import Optional, List, Dict, Any, Type, Union from typing import Optional, List, Dict, Any, Type, Union
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
from pydantic import BaseModel from pydantic import BaseModel
@ -282,6 +283,13 @@ class BaseDoc(BaseModel, metaclass=QueryModelMeta):
# Convert Decimal objects to float for MongoDB compatibility # Convert Decimal objects to float for MongoDB compatibility
doc_dict = self._convert_decimals_to_float(doc_dict) doc_dict = self._convert_decimals_to_float(doc_dict)
# Respect pre-populated id by mapping to MongoDB _id
if getattr(self, 'id', None):
try:
doc_dict['_id'] = ObjectId(self.id)
except Exception:
doc_dict['_id'] = self.id
result = await collection.insert_one(doc_dict) result = await collection.insert_one(doc_dict)
# Set the id field from the inserted document # Set the id field from the inserted document
@ -314,7 +322,17 @@ class BaseDoc(BaseModel, metaclass=QueryModelMeta):
elif hasattr(self, 'auth_code'): elif hasattr(self, 'auth_code'):
query['auth_code'] = self.auth_code query['auth_code'] = self.auth_code
if query: if getattr(self, 'id', None):
# Update by primary key when available
try:
object_id = ObjectId(self.id)
except Exception:
object_id = self.id
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 # Update existing document
result = await collection.update_one(query, {"$set": doc_dict}, upsert=True) result = await collection.update_one(query, {"$set": doc_dict}, upsert=True)
# If it was an insert, set the id field # If it was an insert, set the id field