From b90165b4e4235fb20aac3b380868d2a715ebe5b9 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Tue, 17 May 2022 16:04:34 +0300
Subject: [PATCH] Support queue_name in tasks enqueue

---
 apiserver/apimodels/tasks.py                     |  2 ++
 apiserver/bll/queue/queue_bll.py                 | 12 ++++++++++++
 apiserver/bll/task/task_operations.py            | 14 ++++++++++++++
 apiserver/schema/services/tasks.conf             | 16 ++++++++++++++--
 apiserver/services/tasks.py                      |  2 ++
 .../tests/automated/test_batch_operations.py     |  2 +-
 6 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/apiserver/apimodels/tasks.py b/apiserver/apimodels/tasks.py
index 7acd531..1b70e42 100644
--- a/apiserver/apimodels/tasks.py
+++ b/apiserver/apimodels/tasks.py
@@ -96,6 +96,7 @@ class UpdateRequest(TaskUpdateRequest):
 
 class EnqueueRequest(UpdateRequest):
     queue = StringField()
+    queue_name = StringField()
 
 
 class DeleteRequest(UpdateRequest):
@@ -262,6 +263,7 @@ class StopManyRequest(TaskBatchRequest):
 
 class EnqueueManyRequest(TaskBatchRequest):
     queue = StringField()
+    queue_name = StringField()
     validate_tasks = BoolField(default=False)
 
 
diff --git a/apiserver/bll/queue/queue_bll.py b/apiserver/bll/queue/queue_bll.py
index a4b0c8b..d8b698f 100644
--- a/apiserver/bll/queue/queue_bll.py
+++ b/apiserver/bll/queue/queue_bll.py
@@ -50,6 +50,18 @@ class QueueBLL(object):
             queue.save()
             return queue
 
+    def get_by_name(
+        self,
+        company_id: str,
+        queue_name: str,
+        only: Optional[Sequence[str]] = None,
+    ) -> Queue:
+        qs = Queue.objects(name=queue_name, company=company_id)
+        if only:
+            qs = qs.only(*only)
+
+        return qs.first()
+
     def get_by_id(
         self, company_id: str, queue_id: str, only: Optional[Sequence[str]] = None
     ) -> Queue:
diff --git a/apiserver/bll/task/task_operations.py b/apiserver/bll/task/task_operations.py
index 1e34d4d..30b53bb 100644
--- a/apiserver/bll/task/task_operations.py
+++ b/apiserver/bll/task/task_operations.py
@@ -108,9 +108,23 @@ def enqueue_task(
     queue_id: str,
     status_message: str,
     status_reason: str,
+    queue_name: str = None,
     validate: bool = False,
     force: bool = False,
 ) -> Tuple[int, dict]:
+    if queue_id and queue_name:
+        raise errors.bad_request.ValidationError(
+            "Either queue id or queue name should be provided"
+        )
+
+    if queue_name:
+        queue = queue_bll.get_by_name(
+            company_id=company_id, queue_name=queue_name, only=("id",)
+        )
+        if not queue:
+            queue = queue_bll.create(company_id=company_id, name=queue_name)
+        queue_id = queue.id
+
     if not queue_id:
         # try to get default queue
         queue_id = queue_bll.get_default(company_id).id
diff --git a/apiserver/schema/services/tasks.conf b/apiserver/schema/services/tasks.conf
index 8721707..51e36b1 100644
--- a/apiserver/schema/services/tasks.conf
+++ b/apiserver/schema/services/tasks.conf
@@ -1898,7 +1898,7 @@ Fails if the following parameters in the task were not filled:
             ]
             properties {
                 queue {
-                    description: "Queue id. If not provided, task is added to the default queue."
+                    description: "Queue id. If not provided and no queue name is passed then task is added to the default queue."
                     type: string
                 }
             }
@@ -1914,6 +1914,12 @@ Fails if the following parameters in the task were not filled:
             }
         }
     }
+    "2.19": ${enqueue."1.5"} {
+        request.properties.queue_name {
+            description: The name of the queue. If the queue does not exist then it is auto-created. Cannot be used together with the queue id
+            type: string
+        }
+    }
 }
 enqueue_many {
     "2.13": ${_definitions.change_many_request} {
@@ -1922,7 +1928,7 @@ enqueue_many {
             properties {
                 ids.description: "IDs of the tasks to enqueue"
                 queue {
-                    description: "Queue id. If not provided, tasks are added to the default queue."
+                    description: "Queue id. If not provided and no queue name is passed then tasks are added to the default queue."
                     type: string
                 }
                 validate_tasks {
@@ -1941,6 +1947,12 @@ enqueue_many {
             }
         }
     }
+    "2.19": ${enqueue_many."2.13"} {
+        request.properties.queue_name {
+            description: The name of the queue. If the queue does not exist then it is auto-created. Cannot be used together with the queue id
+            type: string
+        }
+    }
 }
 dequeue {
     "1.5" {
diff --git a/apiserver/services/tasks.py b/apiserver/services/tasks.py
index 5f5dbd2..700497a 100644
--- a/apiserver/services/tasks.py
+++ b/apiserver/services/tasks.py
@@ -861,6 +861,7 @@ def enqueue(call: APICall, company_id, request: EnqueueRequest):
         queue_id=request.queue,
         status_message=request.status_message,
         status_reason=request.status_reason,
+        queue_name=request.queue_name,
         force=request.force,
     )
     call.result.data_model = EnqueueResponse(queued=queued, **res)
@@ -879,6 +880,7 @@ def enqueue_many(call: APICall, company_id, request: EnqueueManyRequest):
             queue_id=request.queue,
             status_message=request.status_message,
             status_reason=request.status_reason,
+            queue_name=request.queue_name,
             validate=request.validate_tasks,
         ),
         ids=request.ids,
diff --git a/apiserver/tests/automated/test_batch_operations.py b/apiserver/tests/automated/test_batch_operations.py
index 48ca572..98b5356 100644
--- a/apiserver/tests/automated/test_batch_operations.py
+++ b/apiserver/tests/automated/test_batch_operations.py
@@ -20,7 +20,7 @@ class TestBatchOperations(TestService):
         ids = [*tasks, missing_id]
 
         # enqueue
-        res = self.api.tasks.enqueue_many(ids=ids)
+        res = self.api.tasks.enqueue_many(ids=ids, queue_name="test")
         self._assert_succeeded(res, tasks)
         self._assert_failed(res, [missing_id])
         data = self.api.tasks.get_all_ex(id=ids).tasks