feat: enque pipelines based on priority

Signed-off-by: Keshav Priyadarshi <git@keshav.space>

Author: keshav-space

vulnerabilities/schedules.py

@@ -95,10 +95,21 @@ def update_pipeline_schedule():
     PipelineSchedule.objects.exclude(pipeline_id__in=pipelines.keys()).delete()
     for id, pipeline_class in pipelines.items():
         run_once = getattr(pipeline_class, "run_once", False)
+        run_interval = getattr(pipeline_class, "run_interval", 24)
+        run_priority = getattr(
+            pipeline_class, "run_priority", PipelineSchedule.ExecutionPriority.DEFAULT
+        )
 
-        PipelineSchedule.objects.get_or_create(
+        pipeline, created = PipelineSchedule.objects.get_or_create(
             pipeline_id=id,
             defaults={
                 "is_run_once": run_once,
+                "run_interval": run_interval,
+                "run_priority": run_priority,
             },
         )
+
+        if not created:
+            pipeline.run_priority = run_priority
+            pipeline.run_interval = run_interval
+            pipeline.save()

vulnerabilities/tasks.py

@@ -20,7 +20,13 @@
 
 logger = logging.getLogger(__name__)
 
-queue = django_rq.get_queue("default")
+default_queue = django_rq.get_queue("default")
+high_queue = django_rq.get_queue("high")
+
+queues = {
+    "default": django_rq.get_queue("default"),
+    "high": django_rq.get_queue("high"),
+}
 
 
 def execute_pipeline(pipeline_id, run_id):
@@ -112,6 +118,8 @@ def set_run_failure(job, connection, type, value, traceback):
 
 def enqueue_pipeline(pipeline_id):
     pipeline_schedule = models.PipelineSchedule.objects.get(pipeline_id=pipeline_id)
+    queue = queues.get(pipeline_schedule.get_priority_display())
+
     if pipeline_schedule.status in [
         models.PipelineRun.Status.RUNNING,
         models.PipelineRun.Status.QUEUED,
@@ -139,5 +147,7 @@ def enqueue_pipeline(pipeline_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)
+
+    for queue in queues.values():
+        if job_id in queue.jobs:
+            queue.remove(job_id)