From 251ee57ffd938535c2995c055a244b0949f5dad2 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Mon, 3 May 2021 18:03:17 +0300
Subject: [PATCH] Fix rapidjson dumps does not support ensure_ascii, only
 Encoder initialization does Add task enqueue status

---
 apiserver/bll/task/task_bll.py        |  4 ++--
 apiserver/bll/task/task_operations.py | 20 ++++++++++++++++----
 apiserver/database/model/task/task.py |  3 +++
 apiserver/service_repo/apicall.py     |  3 ++-
 apiserver/services/tasks.py           |  4 ++--
 apiserver/utilities/json.py           |  1 +
 6 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/apiserver/bll/task/task_bll.py b/apiserver/bll/task/task_bll.py
index cf7286d..33b3bc7 100644
--- a/apiserver/bll/task/task_bll.py
+++ b/apiserver/bll/task/task_bll.py
@@ -668,10 +668,10 @@ class TaskBLL:
 
         return ChangeStatusRequest(
             task=task,
-            new_status=TaskStatus.created,
+            new_status=task.enqueue_status or TaskStatus.created,
             status_reason=status_reason,
             status_message=status_message,
-        ).execute()
+        ).execute(enqueue_status=None)
 
     @classmethod
     def dequeue(cls, task: Task, company_id: str, silent_fail=False):
diff --git a/apiserver/bll/task/task_operations.py b/apiserver/bll/task/task_operations.py
index d5423a9..16a2bb7 100644
--- a/apiserver/bll/task/task_operations.py
+++ b/apiserver/bll/task/task_operations.py
@@ -39,7 +39,14 @@ def archive_task(
         task = TaskBLL.get_task_with_access(
             task,
             company_id=company_id,
-            only=("id", "execution", "status", "project", "system_tags"),
+            only=(
+                "id",
+                "execution",
+                "status",
+                "project",
+                "system_tags",
+                "enqueue_status",
+            ),
             requires_write_access=True,
         )
     try:
@@ -82,7 +89,7 @@ def enqueue_task(
         status_reason=status_reason,
         status_message=status_message,
         allow_same_state_transition=False,
-    ).execute()
+    ).execute(enqueue_status=task.status)
 
     try:
         queue_bll.add_task(company_id=company_id, queue_id=queue_id, task_id=task.id)
@@ -94,7 +101,7 @@ def enqueue_task(
             new_status=task.status,
             force=True,
             status_reason="failed enqueueing",
-        ).execute()
+        ).execute(enqueue_status=None)
         raise
 
     # set the current queue ID in the task
@@ -220,7 +227,12 @@ def reset_task(
         status_reason="reset",
         status_message="reset",
     ).execute(
-        started=None, completed=None, published=None, active_duration=None, **updates,
+        started=None,
+        completed=None,
+        published=None,
+        active_duration=None,
+        enqueue_status=None,
+        **updates,
     )
 
     return dequeued, cleaned_up, res
diff --git a/apiserver/database/model/task/task.py b/apiserver/database/model/task/task.py
index cc6c694..2dd58e4 100644
--- a/apiserver/database/model/task/task.py
+++ b/apiserver/database/model/task/task.py
@@ -261,6 +261,9 @@ class Task(AttributedDocument):
     runtime = SafeDictField(default=dict)
     models: Models = EmbeddedDocumentField(Models, default=Models)
     container = SafeMapField(field=StringField(default=""))
+    enqueue_status = StringField(
+        choices=get_options(TaskStatus), exclude_by_default=True
+    )
 
     def get_index_company(self) -> str:
         """
diff --git a/apiserver/service_repo/apicall.py b/apiserver/service_repo/apicall.py
index 426c0d8..37326fe 100644
--- a/apiserver/service_repo/apicall.py
+++ b/apiserver/service_repo/apicall.py
@@ -634,7 +634,8 @@ class APICall(DataContainer):
             }
             if self.content_type.lower() == JSON_CONTENT_TYPE:
                 try:
-                    res = json.dumps(res, **(self._json_flags or {}))
+                    func = json.dumps if self._json_flags.pop("ensure_ascii", True) else json.dumps_notascii
+                    res = func(res, **(self._json_flags or {}))
                 except Exception as ex:
                     # JSON serialization may fail, probably problem with data or error_data so pop it and try again
                     if not (self.result.data or self.result.error_data):
diff --git a/apiserver/services/tasks.py b/apiserver/services/tasks.py
index c86e2f2..d76d21e 100644
--- a/apiserver/services/tasks.py
+++ b/apiserver/services/tasks.py
@@ -893,7 +893,7 @@ def dequeue(call: APICall, company_id, request: UpdateRequest):
     task = TaskBLL.get_task_with_access(
         request.task,
         company_id=company_id,
-        only=("id", "execution", "status", "project"),
+        only=("id", "execution", "status", "project", "enqueue_status"),
         requires_write_access=True,
     )
     res = DequeueResponse(
@@ -984,7 +984,7 @@ def archive(call: APICall, company_id, request: ArchiveRequest):
     tasks = TaskBLL.assert_exists(
         company_id,
         task_ids=request.tasks,
-        only=("id", "execution", "status", "project", "system_tags"),
+        only=("id", "execution", "status", "project", "system_tags", "enqueue_status"),
     )
     archived = 0
     for task in tasks:
diff --git a/apiserver/utilities/json.py b/apiserver/utilities/json.py
index 6ffbb5c..14cd969 100644
--- a/apiserver/utilities/json.py
+++ b/apiserver/utilities/json.py
@@ -6,4 +6,5 @@ import rapidjson
 DATETIME_MODE = rapidjson.DM_ISO8601 | rapidjson.DM_NAIVE_IS_UTC
 
 dumps = rapidjson.Encoder(datetime_mode=DATETIME_MODE)
+dumps_notascii = rapidjson.Encoder(datetime_mode=DATETIME_MODE, ensure_ascii=False)
 loads = rapidjson.Decoder(datetime_mode=DATETIME_MODE)