diff --git a/vulnerabilities/management/commands/run_scheduler.py b/vulnerabilities/management/commands/run_scheduler.py index 476eb4816..8602ad237 100644 --- a/vulnerabilities/management/commands/run_scheduler.py +++ b/vulnerabilities/management/commands/run_scheduler.py @@ -12,6 +12,7 @@ from vulnerabilities import models from vulnerabilities.schedules import clear_zombie_pipeline_schedules +from vulnerabilities.schedules import refresh_runs from vulnerabilities.schedules import scheduled_job_exists from vulnerabilities.schedules import update_pipeline_schedule @@ -38,5 +39,6 @@ class Command(rqscheduler.Command): def handle(self, *args, **kwargs): clear_zombie_pipeline_schedules() update_pipeline_schedule() + refresh_runs() init_pipeline_scheduled() super(Command, self).handle(*args, **kwargs) diff --git a/vulnerabilities/pipelines/v2_importers/redhat_importer.py b/vulnerabilities/pipelines/v2_importers/redhat_importer.py index 31a6e8d64..80591d341 100644 --- a/vulnerabilities/pipelines/v2_importers/redhat_importer.py +++ b/vulnerabilities/pipelines/v2_importers/redhat_importer.py @@ -30,8 +30,8 @@ from vulnerabilities.pipelines import VulnerableCodeBaseImporterPipelineV2 from vulnerabilities.pipes import extractcode_utils from vulnerabilities.severity_systems import REDHAT_AGGREGATE +from vulnerabilities.utils import get_item from vulnerabilities.utils import load_json -from vulntotal import vulntotal_utils class RedHatImporterPipeline(VulnerableCodeBaseImporterPipelineV2): @@ -135,7 +135,7 @@ def parse_advisory(self, record): ) ) - impacts = get_item(advisory, "product_tree", "branches", 0, "branches", default=[]) + impacts = get_item(advisory, "product_tree", "branches", 0, "branches") or [] for impact in impacts: if impact["category"] == "product_family": continue @@ -145,7 +145,6 @@ def parse_advisory(self, record): "product", "product_identification_helper", "purl", - default=None, ): if not purl.startswith("pkg:rpm/"): continue @@ -188,11 +187,3 @@ def clean_download(self): def on_failure(self): self.clean_download() - - -def get_item(entity, *attributes, default=None): - try: - result = vulntotal_utils.get_item(entity, *attributes) - except (KeyError, IndexError, TypeError) as e: - result = default - return result diff --git a/vulnerabilities/pipelines/v2_improvers/relate_severities.py b/vulnerabilities/pipelines/v2_improvers/relate_severities.py index 1d7083f98..1fa8018cc 100644 --- a/vulnerabilities/pipelines/v2_improvers/relate_severities.py +++ b/vulnerabilities/pipelines/v2_improvers/relate_severities.py @@ -103,7 +103,6 @@ def relate_severities(self): matches = advisory_id_map.get(advisory.advisory_id, set()) for target_id in matches: if target_id != advisory.id: - self.log(f"Relating advisory {advisory.avid} to {target_id}") relations.append( through( from_advisoryv2_id=target_id, diff --git a/vulnerabilities/pipes/group_advisories.py b/vulnerabilities/pipes/group_advisories.py index 22457f3b5..5d1aa98e9 100644 --- a/vulnerabilities/pipes/group_advisories.py +++ b/vulnerabilities/pipes/group_advisories.py @@ -23,8 +23,6 @@ @transaction.atomic def delete_and_save_advisory_set(groups, package, relation=None): - print(f"Grouping advisories for package: {package.purl}") - AdvisorySet.objects.filter( package=package, relation_type=relation, @@ -101,8 +99,6 @@ def delete_and_save_advisory_set(groups, package, relation=None): batch_size=10000, ) - print(f"Successfully saved advisory sets for package: {package.purl}") - def group_advisory_for_package(package, logger=None): """ @@ -133,7 +129,6 @@ def group_advisory_for_package(package, logger=None): group_single_package_with_provided_advisories( package, affecting_advisories, fixed_by_advisories ) - logger(f"Successfully rebuilt advisory sets for package {package.purl}") except Exception as e: if logger: logger(f"Failed rebuilding advisory sets for package {package.purl}: {e!r}") diff --git a/vulnerabilities/schedules.py b/vulnerabilities/schedules.py index 1b358cbaa..038e9a927 100644 --- a/vulnerabilities/schedules.py +++ b/vulnerabilities/schedules.py @@ -116,3 +116,46 @@ def update_pipeline_schedule(): pipeline.run_priority = run_priority pipeline.run_interval = run_interval pipeline.save() + + +def mark_stale_runs(): + """Mark unfinished pipeline runs as stale.""" + from django.db.models import Q + + from vulnerabilities.models import PipelineRun + + stale_jobs = PipelineRun.objects.filter(run_start_date__isnull=False).filter( + Q(run_end_date__isnull=True) | Q(run_exitcode__isnull=True) + ) + stale_jobs_count = stale_jobs.count() + + for job in stale_jobs.iterator(chunk_size=1000): + job.set_run_staled() + + log.info(f"Marked {stale_jobs_count} unfinished jobs as stale.") + + +def requeue_missing_jobs(): + """Requeue pipeline runs that have not started and are missing from the queue.""" + + from vulnerabilities.models import PipelineRun + from vulnerabilities.tasks import enqueue_run + from vulnerabilities.tasks import is_job_in_queue + + missing_jobs_count = 0 + for job in PipelineRun.objects.filter(run_start_date__isnull=True).iterator(chunk_size=1000): + if not is_job_in_queue(job_id=job.run_id): + pipeline_latest_run = job.pipeline.latest_run + if pipeline_latest_run and pipeline_latest_run.run_id == job.run_id: + enqueue_run(run=job) + missing_jobs_count += 1 + else: + job.set_run_staled() + + log.info(f"Requeued {missing_jobs_count} missing jobs.") + + +def refresh_runs(): + """Refresh the stale and missing jobs.""" + mark_stale_runs() + requeue_missing_jobs() diff --git a/vulnerabilities/tasks.py b/vulnerabilities/tasks.py index ee9ae00e2..0188511ed 100644 --- a/vulnerabilities/tasks.py +++ b/vulnerabilities/tasks.py @@ -115,9 +115,22 @@ def set_run_failure(job, connection, type, value, traceback): run.set_run_ended(exitcode=1, output=f"value={value} trace={traceback}") +def enqueue_run(run): + pipeline_schedule = run.pipeline + queue = queues.get(pipeline_schedule.get_run_priority_display()) + + job = queue.enqueue( + execute_pipeline, + pipeline_schedule.pipeline_id, + run.run_id, + job_id=str(run.run_id), + on_failure=set_run_failure, + job_timeout=f"{pipeline_schedule.execution_timeout}h", + ) + + def enqueue_pipeline(pipeline_id): pipeline_schedule = models.PipelineSchedule.objects.get(pipeline_id=pipeline_id) - queue = queues.get(pipeline_schedule.get_run_priority_display()) if pipeline_schedule.status in [ models.PipelineRun.Status.RUNNING, @@ -134,14 +147,8 @@ def enqueue_pipeline(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=f"{pipeline_schedule.execution_timeout}h", - ) + + enqueue_run(run) def dequeue_job(job_id): @@ -152,6 +159,12 @@ def dequeue_job(job_id): queue.remove(job_id) +def is_job_in_queue(job_id): + """Return whether a job exists in any queue.""" + + return any(job_id in queue.jobs for queue in queues.values()) + + def compute_queue_load_factor(): """ Compute worker load per queue. diff --git a/vulnerabilities/tests/pipelines/v2_importers/test_redhat_importer_v2.py b/vulnerabilities/tests/pipelines/v2_importers/test_redhat_importer_v2.py index e9fce5869..57c5ac844 100644 --- a/vulnerabilities/tests/pipelines/v2_importers/test_redhat_importer_v2.py +++ b/vulnerabilities/tests/pipelines/v2_importers/test_redhat_importer_v2.py @@ -17,6 +17,7 @@ from vulnerabilities.models import PackageV2 from vulnerabilities.pipelines.v2_importers.redhat_importer import RedHatImporterPipeline from vulnerabilities.tests import util_tests +from vulnerabilities.tests.pipelines import TestLogger from vulnerabilities.utils import normalize_list TEST_DATA = Path(__file__).parent.parent.parent / "test_data" / "redhat" / "csaf_2_0" @@ -25,10 +26,15 @@ class TestArchLinuxImporterPipeline(TestCase): @patch("vulnerabilities.pipelines.v2_importers.redhat_importer.RedHatImporterPipeline.fetch") def test_redhat_advisories_v2(self, mock_fetch): + logger = TestLogger() + mock_fetch.__name__ = "fetch" pipeline = RedHatImporterPipeline() pipeline.location = TEST_DATA - pipeline.execute() + pipeline.log = logger.write + exit_code, _ = pipeline.execute() + + self.assertEqual(exit_code, 0) self.assertEqual(6, AdvisoryV2.objects.count()) self.assertEqual(93, PackageV2.objects.count()) expected_file = TEST_DATA.parent / "redhat_advisoryv2-expected.json" diff --git a/vulnerabilities/tests/test_models.py b/vulnerabilities/tests/test_models.py index 6e327ed6e..ec6fad7d8 100644 --- a/vulnerabilities/tests/test_models.py +++ b/vulnerabilities/tests/test_models.py @@ -688,6 +688,23 @@ def test_pipelinerun_success(self): self.run1.set_run_ended(exitcode=0) self.assertEqual(self.run1.status, self.run1.Status.SUCCESS) + def test_pipelinerun_schedules_mark_stale_runs_unfinished_job(self): + from vulnerabilities.schedules import mark_stale_runs + + self.run1.set_run_started() + self.assertEqual(self.run1.run_exitcode, None) + mark_stale_runs() + self.run1.refresh_from_db() + self.assertEqual(self.run1.run_exitcode, 88) + + def test_pipelinerun_schedules_mark_stale_runs_job_not_yet_started(self): + from vulnerabilities.schedules import mark_stale_runs + + self.assertEqual(self.run1.run_exitcode, None) + mark_stale_runs() + self.run1.refresh_from_db() + self.assertEqual(self.run1.run_exitcode, None) + class TestPipelineScheduleModel(DjangoTestCase): def setUp(self):