Handle partial results

CybercentreCanada · Sep 6, 2023 · c755152 · c755152
1 parent 3093e1a
commit c755152
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/assemblyline/common/caching.py b/assemblyline/common/caching.py
@@ -150,7 +150,7 @@ def keys(self):
             return self.cache.keys()
 
 
-def generate_conf_key(service_tool_version: Optional[str] = None, task: Optional[Task] = None) -> str:
+def generate_conf_key(service_tool_version: Optional[str] = None, task: Optional[Task] = None, partial_result: bool = False) -> str:
     ignore_salt = None
     service_config = None
     submission_params_str = None
@@ -165,7 +165,7 @@ def generate_conf_key(service_tool_version: Optional[str] = None, task: Optional
         }
         submission_params_str = json.dumps(sorted(submission_params.items()))
 
-        if task.ignore_cache:
+        if task.ignore_cache or partial_result:
             ignore_salt = get_random_id()
 
     if service_tool_version is None and \

diff --git a/assemblyline/odm/models/result.py b/assemblyline/odm/models/result.py
@@ -115,6 +115,7 @@ class Result(odm.Model):
     type = odm.Optional(odm.Keyword())
     size = odm.Optional(odm.Integer())
     drop_file = odm.Boolean(default=False, description="Use to not pass to other stages after this run")
+    partial = odm.Boolean(default=False, description="Invalidate the current result cache creation")
     from_archive = odm.Boolean(index=False, default=False, description="Was loaded from the archive")
 
     def build_key(self, service_tool_version=None, task=None):
@@ -124,16 +125,17 @@ def build_key(self, service_tool_version=None, task=None):
             self.response.service_version,
             self.is_empty(),
             service_tool_version=service_tool_version,
-            task=task
+            task=task,
+            partial=self.partial
         )
 
     @staticmethod
-    def help_build_key(sha256, service_name, service_version, is_empty, service_tool_version=None, task=None):
+    def help_build_key(sha256, service_name, service_version, is_empty, service_tool_version=None, task=None, partial=False):
         key_list = [
             sha256,
             service_name.replace('.', '_'),
             f"v{service_version.replace('.', '_')}",
-            f"c{generate_conf_key(service_tool_version=service_tool_version, task=task)}",
+            f"c{generate_conf_key(service_tool_version=service_tool_version, task=task, partial_result=partial)}",
         ]
 
         if is_empty: