From 5ea0d60a462f31ac771b0f8e4a51ab8910ce388c Mon Sep 17 00:00:00 2001
From: udgover <frederic.baguelin@gmail.com>
Date: Tue, 10 Sep 2024 15:18:23 +0200
Subject: [PATCH] Exclude related_observables_count from model_dump_json in
 save method (#1134)

---
 core/database_arango.py | 12 ++++++++----
 core/schemas/entity.py  |  1 +
 core/schemas/graph.py   |  2 ++
 core/schemas/model.py   |  1 +
 4 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/core/database_arango.py b/core/database_arango.py
index 95e39b571..db4e52d2f 100644
--- a/core/database_arango.py
+++ b/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"):
diff --git a/core/schemas/entity.py b/core/schemas/entity.py
index 1c147d364..484d006d2 100644
--- a/core/schemas/entity.py
+++ b/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"
diff --git a/core/schemas/graph.py b/core/schemas/graph.py
index 76d41ed5e..06b8852b3 100644
--- a/core/schemas/graph.py
+++ b/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
diff --git a/core/schemas/model.py b/core/schemas/model.py
index e5eaacce4..78e7e07f0 100644
--- a/core/schemas/model.py
+++ b/core/schemas/model.py
@@ -4,6 +4,7 @@
 
 
 class YetiModel(BaseModel):
+    _exclude_overwrite: list[str] = list()
     __id: str | None = None
 
     def __init__(self, **data):