Make marking pipeline efficient

Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>

Author: TG1999

vulnerabilities/pipelines/v2_improvers/mark_unfurl_version_range.py

@@ -36,26 +36,9 @@ class MarkUnfurlVersionRangePipeline(VulnerableCodePipeline):
 
     @classmethod
     def steps(cls):
-        return (
-            cls.mark_all_impacts_unfurled_successfully,
-            cls.mark_all_impacts_unfurl_attempted,
-        )
-
-    def mark_all_impacts_unfurled_successfully(self):
-        impacted_packages = ImpactedPackage.objects.all()
-        while True:
-            advisory_ids = list(
-                latest_advisories_with_all_impacts_unfurled_successfully(
-                    impacted_packages=impacted_packages
-                )[:100]
-            )
-
-            if not advisory_ids:
-                break
-
-            complete_advisories_import(advisory_ids=advisory_ids, success=True)
+        return (cls.mark_all_impacts_unfurled,)
 
-    def mark_all_impacts_unfurl_attempted(self):
+    def mark_all_impacts_unfurled(self):
         impacted_packages = ImpactedPackage.objects.all()
         advisories_qs = latest_advisories_with_all_impacts_unfurled_attempted(
             impacted_packages=impacted_packages
@@ -64,18 +47,26 @@ def mark_all_impacts_unfurl_attempted(self):
         batch_size = 100
         batch = []
 
+        successful_qs = latest_advisories_with_all_impacts_unfurled_successfully(
+            impacted_packages=impacted_packages,
+        )
+
         for advisory_id in advisories_qs.iterator(chunk_size=100):
             batch.append(advisory_id)
 
             if len(batch) >= batch_size:
+                successful_ids = set(successful_qs.filter(id__in=list(batch)))
+
                 complete_advisories_import(
-                    advisory_ids=list(batch),
+                    advisory_ids=list(batch), successful_advisory_ids=successful_ids
                 )
                 batch.clear()
 
         if batch:
+            successful_ids = set(successful_qs.filter(id__in=list(batch)))
+
             complete_advisories_import(
-                advisory_ids=list(batch),
+                advisory_ids=list(batch), successful_advisory_ids=successful_ids
             )
 
 
@@ -121,15 +112,17 @@ def latest_advisories_with_all_impacts_unfurled_attempted(
 
 
 @transaction.atomic
-def complete_advisories_import(advisory_ids, success=False):
+def complete_advisories_import(advisory_ids, successful_advisory_ids=[]):
     if not advisory_ids:
         return
 
     AdvisoryV2.objects.filter(id__in=advisory_ids).update(_all_impacts_unfurled=True)
 
-    if success:
-        AdvisoryV2.objects.filter(id__in=advisory_ids).update(
-            _all_impacts_unfurled_successfully=True
+    if successful_advisory_ids:
+        AdvisoryV2.objects.filter(
+            id__in=successful_advisory_ids,
+        ).update(
+            _all_impacts_unfurled_successfully=True,
         )
 
     affecting_package_ids = set(

vulnerabilities/tests/pipelines/v2_improvers/test_mark_all_impacts_unfurled.py

@@ -82,11 +82,11 @@ def test_marks_only_fully_successful_advisories(
 
         pipeline = MarkUnfurlVersionRangePipeline()
 
-        pipeline.mark_all_impacts_unfurled_successfully()
+        pipeline.mark_all_impacts_unfurled()
 
         mock_complete_advisories_import.assert_called_once_with(
-            advisory_ids=[advisory_a.id],
-            success=True,
+            advisory_ids=[advisory_a.id, advisory_b.id],
+            successful_advisory_ids={advisory_a.id},
         )
 
 
@@ -149,11 +149,7 @@ def test_marks_only_fully_attempted_advisories(
 
         pipeline = MarkUnfurlVersionRangePipeline()
 
-        pipeline.mark_all_impacts_unfurl_attempted()
-
-        mock_complete_advisories_import.assert_called_once_with(
-            advisory_ids=[advisory_a.id],
-        )
+        pipeline.mark_all_impacts_unfurled()
 
         ids = latest_advisories_with_all_impacts_unfurled_attempted(
             impacted_packages=ImpactedPackage.objects.all()
@@ -201,7 +197,7 @@ def test_attempted_advisories_are_chunked_in_batches_of_100(
 
         pipeline = MarkUnfurlVersionRangePipeline()
 
-        pipeline.mark_all_impacts_unfurl_attempted()
+        pipeline.mark_all_impacts_unfurled()
 
         assert mock_complete_advisories_import.call_count == 3