Allow dequeueing a deleted task

This commit is contained in:
allegroai 2023-07-26 18:32:32 +03:00
parent a2b9fed92d
commit a17485b1bd
4 changed files with 40 additions and 15 deletions

View File

@ -155,7 +155,13 @@ class QueueBLL(object):
task = Task.get_for_writing( task = Task.get_for_writing(
company=company_id, company=company_id,
id=item.task, id=item.task,
_only=["id", "status", "enqueue_status", "project"], _only=[
"id",
"company",
"status",
"enqueue_status",
"project",
],
) )
if not task: if not task:
continue continue
@ -321,16 +327,13 @@ class QueueBLL(object):
return len(entries_to_remove) if res else 0 return len(entries_to_remove) if res else 0
def reposition_task( def reposition_task(
self, self, company_id: str, queue_id: str, task_id: str, move_count: Union[int, str],
company_id: str,
queue_id: str,
task_id: str,
move_count: Union[int, str],
) -> int: ) -> int:
""" """
Moves the task in the queue to the position calculated by pos_func Moves the task in the queue to the position calculated by pos_func
Returns the updated task position in the queue Returns the updated task position in the queue
""" """
def get_queue_and_task_position(): def get_queue_and_task_position():
q = self.get_queue_with_task( q = self.get_queue_with_task(
company_id=company_id, queue_id=queue_id, task_id=task_id company_id=company_id, queue_id=queue_id, task_id=task_id

View File

@ -449,9 +449,9 @@ class TaskBLL:
return ret return ret
@staticmethod @staticmethod
def remove_task_from_all_queues(company_id: str, task: Task) -> int: def remove_task_from_all_queues(company_id: str, task_id: str) -> int:
return Queue.objects(company=company_id, entries__task=task.id).update( return Queue.objects(company=company_id, entries__task=task_id).update(
pull__entries__task=task.id, last_update=datetime.utcnow() pull__entries__task=task_id, last_update=datetime.utcnow()
) )
@classmethod @classmethod
@ -471,7 +471,7 @@ class TaskBLL:
pass pass
if remove_from_all_queues: if remove_from_all_queues:
cls.remove_task_from_all_queues(company_id=company_id, task=task) cls.remove_task_from_all_queues(company_id=company_id, task_id=task.id)
if task.status not in [TaskStatus.queued, TaskStatus.in_progress]: if task.status not in [TaskStatus.queued, TaskStatus.in_progress]:
return {"updated": 0} return {"updated": 0}

View File

@ -46,6 +46,7 @@ def archive_task(
company_id=company_id, company_id=company_id,
only=( only=(
"id", "id",
"company",
"execution", "execution",
"status", "status",
"project", "project",
@ -102,10 +103,23 @@ def dequeue_task(
status_reason: str, status_reason: str,
remove_from_all_queues: bool = False, remove_from_all_queues: bool = False,
) -> Tuple[int, dict]: ) -> Tuple[int, dict]:
query = dict(id=task_id, company=company_id) # get the task without write access to make sure that it actually exists
task = Task.get_for_writing(**query) task = Task.get(
id=task_id,
company=company_id,
_only=(
"id",
"company",
"execution",
"status",
"project",
"enqueue_status",
),
include_public=True,
)
if not task: if not task:
raise errors.bad_request.InvalidTaskId(**query) TaskBLL.remove_task_from_all_queues(company_id, task_id=task_id)
return 1, {"updated": 0}
res = TaskBLL.dequeue_and_change_status( res = TaskBLL.dequeue_and_change_status(
task, task,
@ -301,7 +315,7 @@ def reset_task(
# dequeue may fail if the task was not enqueued # dequeue may fail if the task was not enqueued
pass pass
TaskBLL.remove_task_from_all_queues(company_id=company_id, task=task) TaskBLL.remove_task_from_all_queues(company_id=company_id, task_id=task.id)
cleaned_up = cleanup_task( cleaned_up = cleanup_task(
company=company_id, company=company_id,

View File

@ -1024,7 +1024,15 @@ def archive(call: APICall, company_id, request: ArchiveRequest):
tasks = TaskBLL.assert_exists( tasks = TaskBLL.assert_exists(
company_id, company_id,
task_ids=request.tasks, task_ids=request.tasks,
only=("id", "execution", "status", "project", "system_tags", "enqueue_status"), only=(
"id",
"company",
"execution",
"status",
"project",
"system_tags",
"enqueue_status",
),
) )
archived = 0 archived = 0
for task in tasks: for task in tasks: