Made get_request_data more efficient

turbinia/api/models/request_status.py

@@ -54,11 +54,74 @@ def get_request_data(
     Returns:
       bool: True if the request has at least one task associated with it.
     """
-    if not tasks:
-      _state_manager = state_manager.get_state_manager()
-      tasks = _state_manager.get_task_data(
-          instance=turbinia_config.INSTANCE_ID, request_id=request_id)
+    _state_manager = state_manager.get_state_manager()
+
+    self.request_id = request_id
+
+    # Gets the information from the request if it is stored in Redis
+    if _state_manager.key_exists(f'TurbiniaRequest:{request_id}'):
+      saved_request = _state_manager.get_request_data(request_id)
+      self.evidence_name = saved_request.get('original_evidence').get('name')
+      self.evidence_id = saved_request.get('original_evidence').get('id')
+      self.requester = saved_request.get('requester')
+      self.reason = saved_request.get('reason')
+      self.task_status = saved_request.get('status')
+      self.task_last_update = saved_request.get('last_update')
+      self.sucessful_tasks = len(saved_request.get('succesful_tasks'))
+      self.failed_tasks = len(saved_request.get('failed_tasks'))
+      self.queued_tasks = len(saved_request.get('queued_tasks'))
+      self.running_tasks = len(saved_request.get('running_tasks'))
+    # If the request is not stored in redis, uses legacy get_request_data
+    else:
+      if not tasks:
+        tasks = _state_manager.get_task_data(
+            instance=turbinia_config.INSTANCE_ID, request_id=request_id)
+      self.get_request_data_legacy(request_id, tasks, summary)
+
+    if self.last_task_update_time:
+      if isinstance(self.last_task_update_time, float):
+        self.last_task_update_time = datetime.datetime.fromtimestamp(
+            self.last_task_update_time).strftime(
+                turbinia_config.DATETIME_FORMAT)
+
+    completed_tasks = self.successful_tasks + self.failed_tasks
+
+    if completed_tasks == self.task_count and self.failed_tasks > 0:
+      self.status = 'completed_with_errors'
+    elif self.failed_tasks == self.task_count:
+      self.status = 'failed'
+    elif self.successful_tasks == self.task_count:
+      self.status = 'successful'
+    else:
+      # TODO(leaniz): Add a 'pending' state to tasks for cases 2 and 3.
+      # ref: https://github.com/google/turbinia/issues/1239
+      #
+      # A 'running' status for a request covers multiple cases:
+      #  1) One or more tasks are still in a running status.
+      #  2) Zero tasks are running, zero or more tasks are queued
+      #    and none have failed/succeeded.
+      #  (e.g. all tasks scheduled on the Turbinia server and none picked
+      #    up by any worker yet.)
+      #  3) Zero tasks are running, one or more tasks are queued
+      #    and some have failed/succeeded.
+      #  (e.g. some tasks have completed, others are scheduled on the
+      #    Turbinia server but not picked up by a worker yet.)
+      #
+      # Note that this method is concerned with a Turbiania request's status
+      # which is different than the status of an individual task.
+      self.status = 'running'
+
+    return bool(self.tasks)
+
+  def get_request_data_legacy(self, request_id: str, tasks: Optional[List[Dict]] = None,
+      summary: bool = False):
+    """Gets information about the original evidence for a specific Turbinia
+    request.
 
+    Args:
+      request_id (str): A Turbinia request identifier.
+      tasks (list): List of tasks.
+    """      
     if not summary:
       for task in tasks:
         current_request_id = task.get('request_id')
@@ -73,7 +136,6 @@ def get_request_data(
     # evidence name. There is a small chance of the first task having a
     # different evidence_name, so getting it from arguments is prefered when
     # they exist.
-    # todo(igormr): Save request information in redis to get the evidence_name
     name_from_args = False
     if tasks:
       if tasks[0].get('all_args'):
@@ -86,11 +148,12 @@ def get_request_data(
 
     initial_start_time = datetime.datetime.now().strftime(
         turbinia_config.DATETIME_FORMAT)
+
+    self.taks_count = len(tasks)
+
     for task in tasks:
-      self.request_id = task.get('request_id')
       self.requester = task.get('requester')
       self.reason = task.get('reason')
-      self.task_count = len(tasks)
       task_status = task.get('status')
       # Gets the evidence_name from the first started task.
       if name_from_args and task.get('evidence_name') == self.evidence_name:
@@ -127,42 +190,6 @@ def get_request_data(
         task['last_update'] = task['last_update'].strftime(
             turbinia_config.DATETIME_FORMAT)
 
-    if self.last_task_update_time:
-      if isinstance(self.last_task_update_time, float):
-        self.last_task_update_time = datetime.datetime.fromtimestamp(
-            self.last_task_update_time).strftime(
-                turbinia_config.DATETIME_FORMAT)
-
-    completed_tasks = self.successful_tasks + self.failed_tasks
-
-    if completed_tasks == self.task_count and self.failed_tasks > 0:
-      self.status = 'completed_with_errors'
-    elif self.failed_tasks == self.task_count:
-      self.status = 'failed'
-    elif self.successful_tasks == self.task_count:
-      self.status = 'successful'
-    else:
-      # TODO(leaniz): Add a 'pending' state to tasks for cases 2 and 3.
-      # ref: https://github.com/google/turbinia/issues/1239
-      #
-      # A 'running' status for a request covers multiple cases:
-      #  1) One or more tasks are still in a running status.
-      #  2) Zero tasks are running, zero or more tasks are queued
-      #    and none have failed/succeeded.
-      #  (e.g. all tasks scheduled on the Turbinia server and none picked
-      #    up by any worker yet.)
-      #  3) Zero tasks are running, one or more tasks are queued
-      #    and some have failed/succeeded.
-      #  (e.g. some tasks have completed, others are scheduled on the
-      #    Turbinia server but not picked up by a worker yet.)
-      #
-      # Note that this method is concerned with a Turbiania request's status
-      # which is different than the status of an individual task.
-      self.status = 'running'
-
-    return bool(self.tasks)
-
-
 class RequestsSummary(BaseModel):
   """Represents a summary view of multiple Turbinia requests."""
   requests_status: List[RequestStatus] = []

turbinia/message.py

@@ -28,7 +28,6 @@
 
 log = logging.getLogger('turbinia')
 
-
 class TurbiniaRequest:
   """An object to request evidence to be processed.
 
@@ -54,6 +53,9 @@ def __init__(
     self.recipe = recipe if recipe else {'globals': {}}
     self.context = context if context else {}
     self.evidence = evidence if evidence else []
+    self.original_evidence = {}
+    if evidence and len(evidence) > 0:
+      self.original_evidence = {'id': evidence[0].id, 'name': evidence[0].name}
     self.group_name = group_name if group_name else ''
     self.reason = reason if reason else ''
     self.all_args = all_args if all_args else ''
@@ -70,13 +72,11 @@ def to_json(self, json_values=False):
       A JSON serialized object.
     """
     serializable = copy.deepcopy(self.__dict__)
-
     if json_values:
       if evidence_list := serializable.pop('evidence'):
-        serializable['original_evidence'] = {
-            'name': evidence_list[0].name,
-            'id': evidence_list[0].id
-        }
+        if not serializable.get('original_evidence') and len(evidence_list) > 0:
+          serializable['original_evidence'] = {'name': evidence_list[0].name,
+                                              'id': evidence_list[0].id}
         serializable['evidence_ids'] = [
             evidence.id for evidence in evidence_list
         ]