From cb9812caeeac463c35f1a16f912ee3fd6d61f8e0 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Thu, 20 Jun 2024 17:44:17 +0300
Subject: [PATCH] Do not return any mongodb instructions as a result of task
 update operations

---
 apiserver/bll/task/utils.py | 13 +++++++++++--
 apiserver/services/tasks.py | 29 +++++++++--------------------
 2 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/apiserver/bll/task/utils.py b/apiserver/bll/task/utils.py
index e6d6560..5cb9627 100644
--- a/apiserver/bll/task/utils.py
+++ b/apiserver/bll/task/utils.py
@@ -4,6 +4,7 @@ from typing import Sequence
 import attr
 import six
 from mongoengine import Q
+from mongoengine.base import UPDATE_OPERATORS
 
 from apiserver.apierrors import errors
 from apiserver.bll.util import update_project_time
@@ -78,8 +79,16 @@ class ChangeStatusRequest(object):
 
         update_project_time(project_id)
 
-        # make sure that _raw_ queries are not returned back to the client
-        fields.pop("__raw__", None)
+        def is_mongo_operator(field: str) -> bool:
+            head, _, tail = field.partition("__")
+            return tail and (head in UPDATE_OPERATORS)
+
+        # make sure to not return _raw_ queries or any of the update operators
+        fields = {
+            key: value
+            for key, value in fields.items()
+            if not (key == "__raw__" or is_mongo_operator(key))
+        }
 
         return dict(updated=updated, fields=fields)
 
diff --git a/apiserver/services/tasks.py b/apiserver/services/tasks.py
index be6a15b..540c772 100644
--- a/apiserver/services/tasks.py
+++ b/apiserver/services/tasks.py
@@ -1027,14 +1027,7 @@ def reset(call: APICall, company_id, request: ResetRequest):
         clear_all=request.clear_all,
         delete_external_artifacts=request.delete_external_artifacts,
     )
-
-    res = ResetResponse(**updates, dequeued=dequeued)
-    # do not return artifacts since they are not serializable
-    res.fields.pop("execution.artifacts", None)
-
-    for key, value in attr.asdict(cleanup_res).items():
-        setattr(res, key, value)
-
+    res = ResetResponse(**updates, **attr.asdict(cleanup_res), dequeued=dequeued)
     call.result.data_model = res
 
 
@@ -1058,23 +1051,19 @@ def reset_many(call: APICall, company_id, request: ResetManyRequest):
         ids=request.ids,
     )
 
-    def clean_res(res: dict) -> dict:
-        # do not return artifacts since they are not serializable
-        fields = res.get("fields")
-        if fields:
-            fields.pop("execution.artifacts", None)
-        return res
-
-    call.result.data_model = ResetManyResponse(
-        succeeded=[
+    succeeded = []
+    for _id, (dequeued, cleanup, res) in results:
+        succeeded.append(
             ResetBatchItem(
                 id=_id,
                 dequeued=bool(dequeued.get("removed")) if dequeued else False,
                 **attr.asdict(cleanup),
-                **clean_res(res),
+                **res,
             )
-            for _id, (dequeued, cleanup, res) in results
-        ],
+        )
+
+    call.result.data_model = ResetManyResponse(
+        succeeded=succeeded,
         failed=failures,
     )