Speed up import process

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

Author: TG1999

vulnerabilities/pipelines/v2_improvers/compute_package_risk.py

@@ -16,6 +16,7 @@
 from vulnerabilities.models import AdvisoryV2
 from vulnerabilities.models import PackageV2
 from vulnerabilities.pipelines import VulnerableCodePipeline
+from vulnerabilities.pipes.risk_score import bulk_update
 from vulnerabilities.risk import compute_vulnerability_risk_factors
 
 
@@ -178,14 +179,3 @@ def compute_and_store_package_risk_score(self):
         )
         self.log(f"Successfully added risk score for {updated:,d} package")
 
-
-def bulk_update(model, items, fields, logger):
-    item_count = 0
-    if items:
-        try:
-            model.objects.bulk_update(objs=items, fields=fields)
-            item_count += len(items)
-        except Exception as e:
-            logger(f"Error updating {model.__name__}: {e}")
-        items.clear()
-    return item_count

vulnerabilities/pipes/advisory.py

@@ -48,7 +48,7 @@
 from vulnerabilities.models import VulnerabilitySeverity
 from vulnerabilities.models import Weakness
 from vulnerabilities.pipes.group_advisories import group_advisory_for_package
-from vulnerabilities.pipes.risk_score import compute_advisory_risk_score
+from vulnerabilities.pipes.risk_score import compute_advisory_risk_score, compute_package_risk_score_bulk
 from vulnerabilities.pipes.risk_score import compute_package_risk_score
 from vulnerabilities.pipes.univers_utils import get_exact_purls_v2
 
@@ -300,6 +300,7 @@ def insert_advisory_v2(
 ):
     from vulnerabilities.models import ImpactedPackage
     from vulnerabilities.models import PackageV2
+    from vulnerabilities.utils import TYPES_WITH_MULTIPLE_IMPORTERS
     from vulnerabilities.utils import compute_content_id_v2
 
     advisory_obj = None
@@ -398,9 +399,19 @@ def insert_advisory_v2(
             purls=package_fixed_purls
         )
 
-        add_packages_to_advisory(
-            logger, advisory_obj, impact, affected_packages_v2, fixed_packages_v2
-        )
+        if affected_pkg.package.type in TYPES_WITH_MULTIPLE_IMPORTERS:
+            group_advisories_for_package_with_lock(
+                advisory_obj, impact, affected_packages_v2, fixed_packages_v2, logger=logger
+            )
+        else:
+            impact.affecting_packages.add(*affected_packages_v2)
+            impact.fixed_by_packages.add(*fixed_packages_v2)
+
+        compute_package_risk_score_bulk(affected_packages_v2)
+        if affected_packages_v2:
+            affected_packages_v2[0].calculate_version_rank
+        elif fixed_packages_v2:
+            fixed_packages_v2[0].calculate_version_rank 
 
         introduced_commit_v2 = get_or_create_advisory_package_commit_patches(
             affected_pkg.introduced_by_commit_patches
@@ -414,7 +425,9 @@ def insert_advisory_v2(
 
 
 @transaction.atomic
-def add_packages_to_advisory(logger, advisory_obj, impact, affected_packages_v2, fixed_packages_v2):
+def group_advisories_for_package_with_lock(
+    advisory_obj, impact, affected_packages_v2, fixed_packages_v2, logger=print
+):
     from vulnerabilities.models import PackageV2
 
     all_packages = sorted(
@@ -453,35 +466,6 @@ def add_packages_to_advisory(logger, advisory_obj, impact, affected_packages_v2,
                 f"Failed to group advisories for package {pkg.purl}: {e!r} \n {traceback_format_exc()}",
                 level=logging.ERROR,
             )
-        try:
-            risk_score = compute_package_risk_score(
-                pkg, current_advisory_risk_score=advisory_obj.risk_score
-            )
-            logger(f"Computed risk score {risk_score} for package {pkg.purl}")
-            pkg.risk_score = risk_score
-            pkg.save()
-            pkg.calculate_version_rank
-        except Exception as e:
-            logger(
-                f"Failed to compute risk score for package {pkg.purl}: {e!r} \n {traceback_format_exc()}",
-                level=logging.ERROR,
-            )
-
-    for pkg in list(fixed_packages_v2):
-        logger(f"Grouping advisories for package: {pkg.purl}")
-        try:
-            group_advisory_for_package(
-                pkg,
-                logger=logger,
-                current_advisory=advisory_obj,
-                current_advisory_relation="fixing",
-            )
-        except Exception as e:
-            logger(
-                f"Failed to group advisories for package {pkg.purl}: {e!r} \n {traceback_format_exc()}",
-                level=logging.ERROR,
-            )
-        pkg.calculate_version_rank
 
 
 @transaction.atomic

vulnerabilities/pipes/risk_score.py

@@ -13,6 +13,7 @@
 from django.db.models import Max
 
 from vulnerabilities.models import AdvisoryV2
+from vulnerabilities.models import PackageV2
 
 
 def quantize_1(value):
@@ -40,6 +41,47 @@ def compute_package_risk_score(package, current_advisory_risk_score=None):
     return round(float(max_risk), 1)
 
 
+def compute_package_risk_score_bulk(packages):
+    """Calculate the risk score for a single PackageV2 object."""
+    purls = packages.values_list("package_url", flat=True)
+    advisories = AdvisoryV2.objects.latest_affecting_advisories_for_purls(purls).only(
+        "id", "risk_score"
+    )
+    qs = (
+        PackageV2.objects.filter(
+            id__in=packages.values_list("id", flat=True),
+            affected_in_impacts__advisory__risk_score__isnull=False,
+            affected_in_impacts__advisory__in=advisories,
+        )
+        .distinct()
+        .annotate(computed_risk=Max("affected_in_impacts__advisory__risk_score"))
+        .only("id")
+    )
+
+    batch = []
+    batch_size = 5000
+    updated = 0
+
+    for pkg in qs.iterator(chunk_size=batch_size):
+        pkg.risk_score = round(float(pkg.computed_risk), 1)
+        batch.append(pkg)
+
+        if len(batch) >= batch_size:
+            updated += bulk_update(
+                model=PackageV2,
+                items=batch,
+                fields=["risk_score"],
+            )
+            batch.clear()
+
+    updated += bulk_update(
+        model=PackageV2,
+        items=batch,
+        fields=["risk_score"],
+    )
+        
+
+
 def compute_advisory_risk_score(advisory):
     """
     Calculate the risk score for a single AdvisoryV2 object.
@@ -77,3 +119,16 @@ def compute_advisory_risk_score(advisory):
         risk_score = None
 
     return quantize_1(weighted_severity), quantize_1(exploitability), quantize_1(risk_score)
+
+
+def bulk_update(model, items, fields, logger=None):
+    item_count = 0
+    if items:
+        try:
+            model.objects.bulk_update(objs=items, fields=fields)
+            item_count += len(items)
+        except Exception as e:
+            if logger:
+                logger(f"Error updating {model.__name__}: {e}")
+        items.clear()
+    return item_count

vulnerabilities/views.py

@@ -326,7 +326,6 @@ def get_context_data(self, **kwargs):
             return context
 
         is_grouped = models.AdvisorySet.objects.filter(package=package).exists()
-        print(f"Is package {package.purl} grouped? {is_grouped}")
 
         if is_grouped:
             context["grouped"] = True