diff --git a/jarvis/db/types.py b/jarvis/db/types.py index 19ba060..adc764b 100644 --- a/jarvis/db/types.py +++ b/jarvis/db/types.py @@ -4,10 +4,9 @@ from datetime import datetime from typing import Any, Optional from bson import ObjectId -from pymongo import ASCENDING, DESCENDING - from jarvis.config import get_config from jarvis.db import DBManager +from pymongo import ASCENDING, DESCENDING logger = logging.getLogger("mongodb") @@ -128,15 +127,18 @@ class MongoObject: :return: If update was successful """ to_update = self.to_dict() - search = {"_id": to_update["_id"]} - try: - result = self.coll.update_one( - search, {"$set": to_update}, upsert=True - ) - except Exception as e: - logger.error(f"Failed to update {type(self).__name__}", e) + if self._id: + search = {"_id": to_update["_id"]} + try: + result = self.coll.update_one( + search, {"$set": to_update}, upsert=True + ) + except Exception as e: + logger.error(f"Failed to update {type(self).__name__}", e) + else: + return result.modified_count > 0 else: - return result.modified_count > 0 + self.insert() @classmethod def get(cls, **kwargs) -> Optional[object]: