Change Live Evaluation API to use RQ #1953

michaelehab · michaelehab · commit ba3b3c1f46aa · 2025-08-25T22:32:07.000+03:00
Signed-off-by: Michael Ehab Mikhail &lt;michael.ehab@hotmail.com&gt;
diff --git a/vulnerabilities/api_v2.py b/vulnerabilities/api_v2.py
@@ -44,6 +44,7 @@
 from vulnerabilities.models import VulnerabilityReference
 from vulnerabilities.models import VulnerabilitySeverity
 from vulnerabilities.models import Weakness
+from vulnerabilities.tasks import enqueue_ad_hoc_pipeline
 from vulnerabilities.throttling import PermissionBasedUserRateThrottle
 
 
@@ -1300,8 +1301,7 @@ def lookup(self, request):
 
 
 class LiveEvaluationSerializer(serializers.Serializer):
-    purl_string = serializers.CharField(help_text="PackageURL to evaluate")
-    no_threading = serializers.BooleanField(required=False, default=False)
+    purl = serializers.CharField(help_text="PackageURL to evaluate")
 
 
 class LiveEvaluationViewSet(viewsets.GenericViewSet):
@@ -1310,7 +1310,7 @@ class LiveEvaluationViewSet(viewsets.GenericViewSet):
     @extend_schema(
         request=LiveEvaluationSerializer,
         responses={
-            202: {"description": "Live evaluation done successfully"},
+            202: {"description": "Live evaluation enqueued successfully; returns Run IDs"},
             400: {"description": "Invalid request"},
             500: {"description": "Internal server error"},
         },
@@ -1324,8 +1324,7 @@ def evaluate(self, request):
                 status=status.HTTP_400_BAD_REQUEST,
             )
 
-        purl_string = serializer.validated_data.get("purl_string")
-        no_threading = serializer.validated_data.get("no_threading", False)
+        purl_string = serializer.validated_data.get("purl")
 
         try:
             purl = PackageURL.from_string(purl_string) if purl_string else None
@@ -1349,31 +1348,78 @@ def evaluate(self, request):
                 status=status.HTTP_400_BAD_REQUEST,
             )
 
-        results = []
+        # Create a single LivePipelineRun to represent this evaluation
+        from vulnerabilities.models import LivePipelineRun
 
-        def run_importer(importer):
+        live_run = LivePipelineRun.objects.create(purl=purl_string)
+        runs = []
+        for importer in importers:
             importer_name = getattr(importer, "pipeline_id", importer.__name__)
-            response_data = {"importer": importer_name, "purl": purl_string, "steps_completed": []}
+            run_id = enqueue_ad_hoc_pipeline(importer_name, inputs={"purl": purl})
+            # Attach each PipelineRun to the LivePipelineRun
+            from vulnerabilities.models import PipelineRun
+
             try:
-                pipeline_instance = importer(purl=purl)
-                status_code, error = pipeline_instance.execute()
-                if status_code != 0:
-                    response_data["error"] = f"Importer {importer_name} failed: {error}"
-                else:
-                    response_data["steps_completed"].append("import")
-            except Exception as e:
-                response_data["error"] = f"Error running importer {importer_name}: {str(e)}"
-            return response_data
-
-        if not no_threading and len(importers) > 1:
-            with ThreadPoolExecutor(max_workers=len(importers)) as executor:
-                future_to_importer = {
-                    executor.submit(run_importer, importer): importer for importer in importers
+                run_obj = PipelineRun.objects.get(run_id=run_id)
+                run_obj.live_pipeline = live_run
+                run_obj.save()
+            except PipelineRun.DoesNotExist:
+                pass
+            runs.append(
+                {
+                    "importer": importer_name,
+                    "run_id": str(run_id) if run_id else None,
                 }
-                for future in as_completed(future_to_importer):
-                    results.append(future.result())
-        else:
-            for importer in importers:
-                results.append(run_importer(importer))
+            )
+        return Response(
+            {"live_run_id": str(live_run.run_id), "runs": runs}, status=status.HTTP_202_ACCEPTED
+        )
 
-        return Response(results, status=status.HTTP_202_ACCEPTED)
+    @extend_schema(
+        parameters=[
+            OpenApiParameter(
+                name="live_run_id",
+                description="UUID of the live run to check status for",
+                required=True,
+                type={"type": "string"},
+                location=OpenApiParameter.PATH,
+            )
+        ],
+        responses={200: "LivePipelineRun status and importers status"},
+    )
+    @action(detail=False, methods=["get"], url_path=r"status/(?P<live_run_id>[0-9a-f\-]{36})")
+    def status(self, request, live_run_id=None):
+        from vulnerabilities.models import LivePipelineRun
+        from vulnerabilities.models import PipelineRun
+
+        try:
+            live_run = LivePipelineRun.objects.get(run_id=live_run_id)
+        except LivePipelineRun.DoesNotExist:
+            return Response({"detail": "Live run not found."}, status=status.HTTP_404_NOT_FOUND)
+
+        live_run.update_status()
+
+        # Gather status for each importer run
+        importer_statuses = []
+        for run in live_run.pipelineruns.all():
+            importer_statuses.append(
+                {
+                    "importer": run.pipeline.pipeline_id,
+                    "run_id": str(run.run_id),
+                    "status": run.status,
+                    "run_start_date": run.run_start_date,
+                    "run_end_date": run.run_end_date,
+                    "run_exitcode": run.run_exitcode,
+                    "run_output": run.run_output,
+                }
+            )
+
+        response = {
+            "live_run_id": str(live_run.run_id),
+            "overall_status": live_run.status,
+            "created_date": live_run.created_date,
+            "completed_date": live_run.completed_date,
+            "purl": live_run.purl,
+            "importers": importer_statuses,
+        }
+        return Response(response)
diff --git a/vulnerabilities/migrations/0102_livepipelinerun_pipelinerun_live_pipeline.py b/vulnerabilities/migrations/0102_livepipelinerun_pipelinerun_live_pipeline.py
@@ -0,0 +1,48 @@
+# Generated by Django 4.2.22 on 2025-08-25 18:03
+
+from django.db import migrations, models
+import django.db.models.deletion
+import uuid
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("vulnerabilities", "0101_advisorytodov2_todorelatedadvisoryv2_and_more"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="LivePipelineRun",
+            fields=[
+                (
+                    "run_id",
+                    models.UUIDField(
+                        default=uuid.uuid4,
+                        editable=False,
+                        primary_key=True,
+                        serialize=False,
+                        unique=True,
+                    ),
+                ),
+                ("created_date", models.DateTimeField(auto_now_add=True, db_index=True)),
+                ("completed_date", models.DateTimeField(blank=True, editable=False, null=True)),
+                ("status", models.CharField(default="queued", max_length=20)),
+                ("purl", models.CharField(blank=True, max_length=300, null=True)),
+            ],
+            options={
+                "ordering": ["-created_date"],
+            },
+        ),
+        migrations.AddField(
+            model_name="pipelinerun",
+            name="live_pipeline",
+            field=models.ForeignKey(
+                blank=True,
+                null=True,
+                on_delete=django.db.models.deletion.CASCADE,
+                related_name="pipelineruns",
+                to="vulnerabilities.livepipelinerun",
+            ),
+        ),
+    ]
diff --git a/vulnerabilities/models.py b/vulnerabilities/models.py
@@ -1972,6 +1972,35 @@ class CodeFixV2(CodeChangeV2):
     )
 
 
+class LivePipelineRun(models.Model):
+    """Represents a single live evaluation run for all compatible importers."""
+
+    run_id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False, unique=True)
+    created_date = models.DateTimeField(auto_now_add=True, db_index=True)
+    completed_date = models.DateTimeField(blank=True, null=True, editable=False)
+    status = models.CharField(max_length=20, default="queued")
+    purl = models.CharField(max_length=300, blank=True, null=True)
+
+    def is_finished(self):
+        return self.status == "finished"
+
+    def update_status(self):
+        if not self.pipelineruns.exists():
+            self.status = "queued"
+        elif all(run.status == PipelineRun.Status.SUCCESS for run in self.pipelineruns.all()):
+            self.status = "finished"
+            self.completed_date = timezone.now()
+        elif any(run.status == PipelineRun.Status.FAILURE for run in self.pipelineruns.all()):
+            self.status = "failed"
+            self.completed_date = timezone.now()
+        else:
+            self.status = "running"
+        self.save()
+
+    class Meta:
+        ordering = ["-created_date"]
+
+
 class PipelineRun(models.Model):
     """The Database representation of a pipeline execution."""
 
@@ -1981,6 +2010,14 @@ class PipelineRun(models.Model):
         on_delete=models.CASCADE,
     )
 
+    live_pipeline = models.ForeignKey(
+        "LivePipelineRun",
+        related_name="pipelineruns",
+        on_delete=models.CASCADE,
+        blank=True,
+        null=True,
+    )
+
     run_id = models.UUIDField(
         primary_key=True,
         default=uuid.uuid4,
@@ -2245,8 +2282,9 @@ def append_to_log(self, message, is_multiline=False):
         if not is_multiline:
             message = message.replace("\n", "").replace("\r", "")
 
-        self.log = self.log + message + "\n"
-        self.save(update_fields=["log"])
+        new_log = (self.log or "") + message + "\n"
+        type(self).objects.filter(run_id=self.run_id).update(log=new_log)
+        self.log = new_log
 
     def dequeue(self):
         from vulnerabilities.tasks import dequeue_job
@@ -2342,12 +2380,15 @@ def save(self, *args, **kwargs):
     def pipeline_class(self):
         """Return the pipeline class."""
         from vulnerabilities.importers import IMPORTERS_REGISTRY
+        from vulnerabilities.importers import LIVE_IMPORTERS_REGISTRY
         from vulnerabilities.improvers import IMPROVERS_REGISTRY
 
         if self.pipeline_id in IMPROVERS_REGISTRY:
             return IMPROVERS_REGISTRY.get(self.pipeline_id)
         if self.pipeline_id in IMPORTERS_REGISTRY:
             return IMPORTERS_REGISTRY.get(self.pipeline_id)
+        if self.pipeline_id in LIVE_IMPORTERS_REGISTRY:
+            return LIVE_IMPORTERS_REGISTRY.get(self.pipeline_id)
 
     @property
     def description(self):
diff --git a/vulnerabilities/tasks.py b/vulnerabilities/tasks.py
@@ -21,10 +21,10 @@
 
 logger = logging.getLogger(__name__)
 
-queue = django_rq.get_queue("default")
+default_queue = django_rq.get_queue("default")
 
 
-def execute_pipeline(pipeline_id, run_id):
+def execute_pipeline(pipeline_id, run_id, inputs=None):
     from vulnerabilities.pipelines import VulnerableCodePipeline
 
     logger.info(f"Enter `execute_pipeline` {pipeline_id}")
@@ -39,7 +39,8 @@ def execute_pipeline(pipeline_id, run_id):
     exitcode = 0
     run_class = run.pipeline_class
     if issubclass(run_class, VulnerableCodePipeline):
-        pipeline_instance = run_class(run_instance=run)
+        inputs = inputs or {}
+        pipeline_instance = run_class(run_instance=run, **inputs)
         exitcode, output = pipeline_instance.execute()
     elif issubclass(run_class, Importer) or issubclass(run_class, Improver):
         exitcode, output = legacy_runner(run_class=run_class, run=run)
@@ -121,7 +122,7 @@ def enqueue_pipeline(pipeline_id):
     run = models.PipelineRun.objects.create(
         pipeline=pipeline_schedule,
     )
-    job = queue.enqueue(
+    job = default_queue.enqueue(
         execute_pipeline,
         pipeline_id,
         run.run_id,
@@ -131,7 +132,35 @@ def enqueue_pipeline(pipeline_id):
     )
 
 
+def enqueue_ad_hoc_pipeline(pipeline_id, *, inputs=None):
+    """Enqueue a one-off execution for the given pipeline_id with optional inputs.
+
+    Returns the created run_id or None if the pipeline cannot be enqueued.
+    """
+    try:
+        pipeline_schedule = models.PipelineSchedule.objects.get(pipeline_id=pipeline_id)
+    except models.PipelineSchedule.DoesNotExist:
+        pipeline_schedule = models.PipelineSchedule.objects.create(
+            pipeline_id=pipeline_id,
+            is_active=False,
+        )
+
+    run = models.PipelineRun.objects.create(pipeline=pipeline_schedule)
+
+    live_queue = django_rq.get_queue("live")
+    job = live_queue.enqueue(
+        execute_pipeline,
+        pipeline_id,
+        run.run_id,
+        inputs or {},
+        job_id=str(run.run_id),
+        on_failure=set_run_failure,
+        job_timeout=f"{pipeline_schedule.execution_timeout}h",
+    )
+    return run.run_id
+
+
 def dequeue_job(job_id):
     """Remove a job from queue if it hasn't been executed yet."""
-    if job_id in queue.jobs:
-        queue.remove(job_id)
+    if job_id in default_queue.jobs:
+        default_queue.remove(job_id)
diff --git a/vulnerabilities/tests/test_api_v2.py b/vulnerabilities/tests/test_api_v2.py
diff --git a/vulnerablecode/settings.py b/vulnerablecode/settings.py
