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

View File

@ -449,9 +449,9 @@ class TaskBLL:
return ret
@staticmethod
def remove_task_from_all_queues(company_id: str, task: Task) -> int:
return Queue.objects(company=company_id, entries__task=task.id).update(
pull__entries__task=task.id, last_update=datetime.utcnow()
def remove_task_from_all_queues(company_id: str, task_id: str) -> int:
return Queue.objects(company=company_id, entries__task=task_id).update(
pull__entries__task=task_id, last_update=datetime.utcnow()
)
@classmethod
@ -471,7 +471,7 @@ class TaskBLL:
pass
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]:
return {"updated": 0}

View File

@ -46,6 +46,7 @@ def archive_task(
company_id=company_id,
only=(
"id",
"company",
"execution",
"status",
"project",
@ -102,10 +103,23 @@ def dequeue_task(
status_reason: str,
remove_from_all_queues: bool = False,
) -> Tuple[int, dict]:
query = dict(id=task_id, company=company_id)
task = Task.get_for_writing(**query)
# get the task without write access to make sure that it actually exists
task = Task.get(
id=task_id,
company=company_id,
_only=(
"id",
"company",
"execution",
"status",
"project",
"enqueue_status",
),
include_public=True,
)
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(
task,
@ -301,7 +315,7 @@ def reset_task(
# dequeue may fail if the task was not enqueued
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(
company=company_id,

View File

@ -1024,7 +1024,15 @@ 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", "enqueue_status"),
only=(
"id",
"company",
"execution",
"status",
"project",
"system_tags",
"enqueue_status",
),
)
archived = 0
for task in tasks: