From a566e8763fdd4e39e2a71a45ea7af70217252e28 Mon Sep 17 00:00:00 2001
From: Keshav Priyadarshi <git@keshav.space>
Date: Fri, 8 Nov 2024 22:04:27 +0530
Subject: [PATCH] Add on_failure to handle cleanup during pipeline failure

- Resolves https://github.com/aboutcode-org/vulnerablecode/issues/1639

Signed-off-by: Keshav Priyadarshi <git@keshav.space>
---
 vulnerabilities/pipelines/__init__.py        | 53 ++++++++++++++++++++
 vulnerabilities/pipelines/gitlab_importer.py |  3 ++
 vulnerabilities/pipelines/npm_importer.py    |  3 ++
 vulnerabilities/pipelines/pypa_importer.py   |  3 ++
 4 files changed, 62 insertions(+)

diff --git a/vulnerabilities/pipelines/__init__.py b/vulnerabilities/pipelines/__init__.py
index 0d3589b67..ff6c41631 100644
--- a/vulnerabilities/pipelines/__init__.py
+++ b/vulnerabilities/pipelines/__init__.py
@@ -10,11 +10,13 @@
 import logging
 from datetime import datetime
 from datetime import timezone
+from timeit import default_timer as timer
 from traceback import format_exc as traceback_format_exc
 from typing import Iterable
 
 from aboutcode.pipeline import BasePipeline
 from aboutcode.pipeline import LoopProgress
+from aboutcode.pipeline import humanize_time
 
 from vulnerabilities.importer import AdvisoryData
 from vulnerabilities.improver import MAX_CONFIDENCE
@@ -29,6 +31,57 @@
 class VulnerableCodePipeline(BasePipeline):
     pipeline_id = None  # Unique Pipeline ID
 
+    def on_failure(self):
+        """
+        Tasks to run in the event that pipeline execution fails.
+
+        Implement cleanup or other tasks that need to be performed
+        on pipeline failure, such as:
+            - Removing cloned repositories.
+            - Deleting downloaded archives.
+        """
+        pass
+
+    def execute(self):
+        """Execute each steps in the order defined on this pipeline class."""
+        self.log(f"Pipeline [{self.pipeline_name}] starting")
+
+        steps = self.pipeline_class.get_steps(groups=self.selected_groups)
+        steps_count = len(steps)
+        pipeline_start_time = timer()
+
+        for current_index, step in enumerate(steps, start=1):
+            step_name = step.__name__
+
+            if self.selected_steps and step_name not in self.selected_steps:
+                self.log(f"Step [{step_name}] skipped")
+                continue
+
+            self.set_current_step(f"{current_index}/{steps_count} {step_name}")
+            self.log(f"Step [{step_name}] starting")
+            step_start_time = timer()
+
+            try:
+                step(self)
+            except Exception as exception:
+                self.log("Pipeline failed")
+                on_failure_start_time = timer()
+                self.log(f"Running [on_failure] tasks")
+                self.on_failure()
+                on_failure_run_time = timer() - on_failure_start_time
+                self.log(f"Completed [on_failure] tasks in {humanize_time(on_failure_run_time)}")
+
+                return 1, self.output_from_exception(exception)
+
+            step_run_time = timer() - step_start_time
+            self.log(f"Step [{step_name}] completed in {humanize_time(step_run_time)}")
+
+        self.set_current_step("")  # Reset the `current_step` field on completion
+        pipeline_run_time = timer() - pipeline_start_time
+        self.log(f"Pipeline completed in {humanize_time(pipeline_run_time)}")
+
+        return 0, ""
+
     def log(self, message, level=logging.INFO):
         """Log the given `message` to the current module logger and execution_log."""
         now_local = datetime.now(timezone.utc).astimezone()
diff --git a/vulnerabilities/pipelines/gitlab_importer.py b/vulnerabilities/pipelines/gitlab_importer.py
index 87fef15d0..4f25c4d94 100644
--- a/vulnerabilities/pipelines/gitlab_importer.py
+++ b/vulnerabilities/pipelines/gitlab_importer.py
@@ -106,6 +106,9 @@ def clean_downloads(self):
             self.log(f"Removing cloned repository")
             self.vcs_response.delete()
 
+    def on_failure(self):
+        self.clean_downloads()
+
 
 def parse_advisory_path(base_path: Path, file_path: Path) -> Tuple[str, str, str]:
     """
diff --git a/vulnerabilities/pipelines/npm_importer.py b/vulnerabilities/pipelines/npm_importer.py
index 60a1d109c..7b6d3aba2 100644
--- a/vulnerabilities/pipelines/npm_importer.py
+++ b/vulnerabilities/pipelines/npm_importer.py
@@ -163,3 +163,6 @@ def clean_downloads(self):
         if self.vcs_response:
             self.log(f"Removing cloned repository")
             self.vcs_response.delete()
+
+    def on_failure(self):
+        self.clean_downloads()
diff --git a/vulnerabilities/pipelines/pypa_importer.py b/vulnerabilities/pipelines/pypa_importer.py
index bdda50c94..aebafacf4 100644
--- a/vulnerabilities/pipelines/pypa_importer.py
+++ b/vulnerabilities/pipelines/pypa_importer.py
@@ -68,3 +68,6 @@ def clean_downloads(self):
         if self.vcs_response:
             self.log(f"Removing cloned repository")
             self.vcs_response.delete()
+
+    def on_failure(self):
+        self.clean_downloads()