Minor updates to use celery states

google · Oct 1, 2024 · 31093f0 · 31093f0
1 parent 4512e15
commit 31093f0
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/turbinia/state_manager.py b/turbinia/state_manager.py
@@ -22,6 +22,7 @@
 import json
 import logging
 import sys
+from celery import states as celery_states
 from datetime import datetime
 from datetime import timedelta
 from typing import Any, List, Dict, Optional
@@ -291,11 +292,11 @@ def update_request_task(self, task) -> None:
           self.redis_client.add_to_list(request_key, 'failed_tasks', task.id)
           statuses_to_remove.remove('failed_tasks')
       task_status = task.celery_state
-      if task_status == 'STARTED':
+      if task_status == celery_states.STARTED:
         self.redis_client.add_to_list(request_key, 'running_tasks', task.id)
         statuses_to_remove.remove('running_tasks')
-      elif (task_status is None or task_status == 'RECEIVED' or
-            task_status == 'PENDING'):
+      elif (task_status is None or task_status == celery_states.RECEIVED or
+            task_status == celery_states.PENDING):
         self.redis_client.add_to_list(request_key, 'queued_tasks', task.id)
         statuses_to_remove.remove('queued_tasks')
       for status_name in statuses_to_remove:

diff --git a/turbinia/workers/__init__.py b/turbinia/workers/__init__.py
@@ -130,7 +130,7 @@ class TurbiniaTaskResult:
   # The list of attributes that we will persist into storage
   STORED_ATTRIBUTES = [
       'report_data', 'report_priority', 'run_time', 'status', 'saved_paths',
-      'successful', 'evidence_size'
+      'successful', 'evidence_size', 'worker_name'
   ]
 
   def __init__(