From 6c5f966ed43920b0f703bc51b2c5e963905a880c Mon Sep 17 00:00:00 2001 From: allegroai <> Date: Wed, 26 Jul 2023 18:55:05 +0300 Subject: [PATCH] Add new_status field to tasks.dequeue and dequeue_many endpoints --- apiserver/apimodels/tasks.py | 2 ++ apiserver/bll/queue/queue_bll.py | 1 + apiserver/bll/task/task_bll.py | 4 +++- apiserver/bll/task/task_operations.py | 6 ++++++ apiserver/schema/services/tasks.conf | 12 ++++++++++++ apiserver/services/tasks.py | 2 ++ apiserver/tests/automated/test_queues.py | 3 +++ 7 files changed, 29 insertions(+), 1 deletion(-) diff --git a/apiserver/apimodels/tasks.py b/apiserver/apimodels/tasks.py index e16048b..c197fef 100644 --- a/apiserver/apimodels/tasks.py +++ b/apiserver/apimodels/tasks.py @@ -98,6 +98,7 @@ class UpdateRequest(TaskUpdateRequest): class DequeueRequest(UpdateRequest): remove_from_all_queues = BoolField(default=False) + new_status = StringField() class EnqueueRequest(UpdateRequest): @@ -280,6 +281,7 @@ class StopManyRequest(TaskBatchRequest): class DequeueManyRequest(TaskBatchRequest): remove_from_all_queues = BoolField(default=False) + new_status = StringField() class EnqueueManyRequest(TaskBatchRequest): diff --git a/apiserver/bll/queue/queue_bll.py b/apiserver/bll/queue/queue_bll.py index c096b0f..1c33651 100644 --- a/apiserver/bll/queue/queue_bll.py +++ b/apiserver/bll/queue/queue_bll.py @@ -172,6 +172,7 @@ class QueueBLL(object): status_reason="Queue was deleted", status_message="", user_id=user_id, + force=True, ).execute(enqueue_status=None) except Exception as ex: log.exception( diff --git a/apiserver/bll/task/task_bll.py b/apiserver/bll/task/task_bll.py index 689e79a..210adf1 100644 --- a/apiserver/bll/task/task_bll.py +++ b/apiserver/bll/task/task_bll.py @@ -463,6 +463,7 @@ class TaskBLL: status_message: str, status_reason: str, remove_from_all_queues=False, + new_status=None, ): try: cls.dequeue(task, company_id, silent_fail=True) @@ -478,10 +479,11 @@ class TaskBLL: return ChangeStatusRequest( task=task, - new_status=task.enqueue_status or TaskStatus.created, + new_status=new_status or task.enqueue_status or TaskStatus.created, status_reason=status_reason, status_message=status_message, user_id=user_id, + force=True, ).execute(enqueue_status=None) @classmethod diff --git a/apiserver/bll/task/task_operations.py b/apiserver/bll/task/task_operations.py index 04d2c29..8e4598b 100644 --- a/apiserver/bll/task/task_operations.py +++ b/apiserver/bll/task/task_operations.py @@ -23,6 +23,7 @@ from apiserver.database.model.task.task import ( Execution, DEFAULT_LAST_ITERATION, ) +from apiserver.database.utils import get_options from apiserver.utilities.dicts import nested_set log = config.logger(__file__) @@ -102,7 +103,11 @@ def dequeue_task( status_message: str, status_reason: str, remove_from_all_queues: bool = False, + new_status=None, ) -> Tuple[int, dict]: + if new_status and new_status not in get_options(TaskStatus): + raise errors.bad_request.ValidationError(f"Invalid task status: {new_status}") + # get the task without write access to make sure that it actually exists task = Task.get( id=task_id, @@ -128,6 +133,7 @@ def dequeue_task( status_message=status_message, status_reason=status_reason, remove_from_all_queues=remove_from_all_queues, + new_status=new_status, ) return 1, res diff --git a/apiserver/schema/services/tasks.conf b/apiserver/schema/services/tasks.conf index 8e7230d..9b6c6e0 100644 --- a/apiserver/schema/services/tasks.conf +++ b/apiserver/schema/services/tasks.conf @@ -1525,6 +1525,12 @@ dequeue { default: false } } + "2.26": ${dequeue."2.25"} { + request.properties.new_status { + type: string + description: The new status to assign to the task after the dequeue instead of the default one + } + } } dequeue_many { "2.13": ${_definitions.change_many_request} { @@ -1550,6 +1556,12 @@ dequeue_many { default: false } } + "2.26": ${dequeue_many."2.25"} { + request.properties.new_status { + type: string + description: The new status to assign to the task after the dequeue instead of the default one + } + } } set_requirements { "2.1" { diff --git a/apiserver/services/tasks.py b/apiserver/services/tasks.py index fa0399f..fef39b2 100644 --- a/apiserver/services/tasks.py +++ b/apiserver/services/tasks.py @@ -921,6 +921,7 @@ def dequeue(call: APICall, company_id, request: DequeueRequest): status_message=request.status_message, status_reason=request.status_reason, remove_from_all_queues=request.remove_from_all_queues, + new_status=request.new_status, ) call.result.data_model = DequeueResponse(dequeued=dequeued, **res) @@ -937,6 +938,7 @@ def dequeue_many(call: APICall, company_id, request: DequeueManyRequest): status_message=request.status_message, status_reason=request.status_reason, remove_from_all_queues=request.remove_from_all_queues, + new_status=request.new_status, ), ids=request.ids, ) diff --git a/apiserver/tests/automated/test_queues.py b/apiserver/tests/automated/test_queues.py index e7a2e71..1959d91 100644 --- a/apiserver/tests/automated/test_queues.py +++ b/apiserver/tests/automated/test_queues.py @@ -84,6 +84,9 @@ class TestQueues(TestService): res = self.api.queues.get_by_id(queue=queue) self.assertQueueTasks(res.queue, [task]) self.assertTaskTags(task, system_tags=[]) + self.api.tasks.dequeue(task=task, new_status="published") + res = self.api.tasks.get_by_id(task=task) + self.assertEqual(res.task.status, "published") def test_dequeue_not_queued_task(self): # queue = self._temp_queue("TestTempQueue")