Use scheduler to explicitly queue pipeline execution jobs

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

Author: keshav-space

vulnerabilities/management/commands/run_scheduler.py

@@ -18,7 +18,9 @@
 
 def init_pipeline_scheduled():
     """Initialize schedule jobs for active PipelineSchedule."""
-    active_pipeline_qs = models.PipelineSchedule.objects.filter(is_active=True)
+    active_pipeline_qs = models.PipelineSchedule.objects.filter(is_active=True).order_by(
+        "created_date"
+    )
     for pipeline_schedule in active_pipeline_qs:
         if scheduled_job_exists(pipeline_schedule.schedule_work_id):
             continue

vulnerabilities/schedules.py

@@ -7,34 +7,37 @@
 # See https://aboutcode.org for more information about nexB OSS projects.
 #
 
+import datetime
 import logging
 
 import django_rq
 from redis.exceptions import ConnectionError
 
-from vulnerabilities.tasks import execute_pipeline
+from vulnerabilities.tasks import enqueue_pipeline
 from vulnerablecode.settings import VULNERABLECODE_PIPELINE_TIMEOUT
 
 log = logging.getLogger(__name__)
 scheduler = django_rq.get_scheduler()
 
 
-def schedule_execution(pipeline_schedule):
+def schedule_execution(pipeline_schedule, execute_now=False):
     """
     Takes a `PackageSchedule` object as input and schedule a
     recurring job using `rq_scheduler` to execute the pipeline.
     """
-    first_execution = pipeline_schedule.next_run_date
+    first_execution = datetime.datetime.now(tz=datetime.timezone.utc)
+    if not execute_now:
+        first_execution = pipeline_schedule.next_run_date
+
     interval_in_seconds = pipeline_schedule.run_interval * 24 * 60 * 60
 
     job = scheduler.schedule(
         scheduled_time=first_execution,
-        func=execute_pipeline,
+        func=enqueue_pipeline,
         args=[pipeline_schedule.pipeline_id],
         interval=interval_in_seconds,
-        result_ttl=interval_in_seconds,  # Remove job results after next run
         timeout=VULNERABLECODE_PIPELINE_TIMEOUT,
-        repeat=None,  # None for repeat forever
+        repeat=None,
     )
     return job._id
 
@@ -89,6 +92,6 @@ def update_pipeline_schedule():
     from vulnerabilities.models import PipelineSchedule
 
     pipeline_ids = [*IMPORTERS_REGISTRY.keys(), *IMPROVERS_REGISTRY.keys()]
-    # pipeline_ids = ["nvd_importer", "vulnerabilities.importers.curl.CurlImporter"]
+
     PipelineSchedule.objects.exclude(pipeline_id__in=pipeline_ids).delete()
     [PipelineSchedule.objects.get_or_create(pipeline_id=id) for id in pipeline_ids]

vulnerabilities/tasks.py

@@ -12,28 +12,26 @@
 from io import StringIO
 from traceback import format_exc as traceback_format_exc
 
-from rq import get_current_job
+import django_rq
 
 from vulnerabilities import models
 from vulnerabilities.importer import Importer
 from vulnerabilities.improver import Improver
+from vulnerablecode.settings import VULNERABLECODE_PIPELINE_TIMEOUT
 
 logger = logging.getLogger(__name__)
 
+queue = django_rq.get_queue("default")
 
-def execute_pipeline(pipeline_id):
+
+def execute_pipeline(pipeline_id, run_id):
     from vulnerabilities.pipelines import VulnerableCodePipeline
 
     logger.info(f"Enter `execute_pipeline` {pipeline_id}")
 
-    pipeline_schedule = models.PipelineSchedule.objects.get(pipeline_id=pipeline_id)
-    job = get_current_job()
-
-    run = models.PipelineRun.objects.create(
-        pipeline=pipeline_schedule,
-        run_id=job.id,
+    run = models.PipelineRun.objects.get(
+        run_id=run_id,
     )
-
     run.set_vulnerablecode_version_and_commit()
     run.set_run_started()
 
@@ -104,3 +102,18 @@ def set_run_failure(job, connection, type, value, traceback):
         return
 
     run.set_run_ended(exitcode=1, output=f"value={value} trace={traceback}")
+
+
+def enqueue_pipeline(pipeline_id):
+    pipeline_schedule = models.PipelineSchedule.objects.get(pipeline_id=pipeline_id)
+    run = models.PipelineRun.objects.create(
+        pipeline=pipeline_schedule,
+    )
+    job = queue.enqueue(
+        execute_pipeline,
+        pipeline_id,
+        run.run_id,
+        job_id=str(run.run_id),
+        on_failure=set_run_failure,
+        job_timeout=VULNERABLECODE_PIPELINE_TIMEOUT,
+    )