Exclude related_observables_count from model_dump_json in save method

core/database_arango.py

@@ -250,13 +250,17 @@ def save(
         Returns:
           The created Yeti object.
         """
-        doc_dict = self.model_dump(exclude_unset=True, exclude=["tags"])
+        exclude = ["tags"] + self._exclude_overwrite
+        doc_dict = self.model_dump(exclude_unset=True, exclude=exclude)
         if doc_dict.get("id") is not None:
-            result = self._update(self.model_dump_json(exclude=["tags"]))
+            exclude = ["tags"] + self._exclude_overwrite
+            result = self._update(self.model_dump_json(exclude=exclude))
         else:
-            result = self._insert(self.model_dump_json(exclude=["tags", "id"]))
+            exclude = ["tags", "id"] + self._exclude_overwrite
+            result = self._insert(self.model_dump_json(exclude=exclude))
             if not result:
-                result = self._update(self.model_dump_json(exclude=exclude_overwrite))
+                exclude = exclude_overwrite + self._exclude_overwrite
+                result = self._update(self.model_dump_json(exclude=exclude))
         yeti_object = self.__class__(**result)
         # TODO: Override this if we decide to implement YetiTagModel
         if hasattr(self, "tags"):

core/schemas/entity.py

@@ -27,6 +27,7 @@ class EntityType(str, Enum):
 
 
 class Entity(YetiTagModel, database_arango.ArangoYetiConnector):
+    _exclude_overwrite: list[str] = ["related_observables_count"]
     _collection_name: ClassVar[str] = "entities"
     _type_filter: ClassVar[str] = ""
     _root_type: Literal["entity"] = "entity"

core/schemas/graph.py

@@ -17,6 +17,7 @@ class GraphFilter(BaseModel):
 # Relationship and TagRelationship do not inherit from YetiModel
 # because they represent and id in the form of collection_name/id
 class Relationship(BaseModel, database_arango.ArangoYetiConnector):
+    _exclude_overwrite: list[str] = list()
     _collection_name: ClassVar[str] = "links"
     _type_filter: ClassVar[str | None] = None
     __id: str | None = None
@@ -44,6 +45,7 @@ def load(cls, object: dict):
 
 
 class TagRelationship(BaseModel, database_arango.ArangoYetiConnector):
+    _exclude_overwrite: list[str] = list()
     _collection_name: ClassVar[str] = "tagged"
     _type_filter: None = None
     __id: str | None = None

core/schemas/model.py

@@ -4,6 +4,7 @@
 
 
 class YetiModel(BaseModel):
+    _exclude_overwrite: list[str] = list()
     __id: str | None = None
 
     def __init__(self, **data):