Move export serializers to the pipes module

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

Author: keshav-space

vulnerabilities/pipelines/exporters/federate_vulnerabilities.py

@@ -14,19 +14,13 @@
 from operator import attrgetter
 from pathlib import Path
 
-import saneyaml
 from aboutcode.pipeline import LoopProgress
 from django.conf import settings
-from django.db.models import Prefetch
 from django.utils import timezone
 
 from aboutcode.federated import DataFederation
-from vulnerabilities.models import AdvisoryV2
-from vulnerabilities.models import ImpactedPackage
-from vulnerabilities.models import ImpactedPackageAffecting
-from vulnerabilities.models import ImpactedPackageFixedBy
-from vulnerabilities.models import PackageV2
 from vulnerabilities.pipelines import VulnerableCodePipeline
+from vulnerabilities.pipes import export
 from vulnerabilities.pipes import federatedcode
 from vulnerabilities.utils import load_json
 
@@ -99,7 +93,7 @@ def publish_package_related_advisories(self):
         chunk_size = 500
         files_to_commit = set()
 
-        packages_count, package_qs = package_prefetched_qs(self.checkpoint)
+        packages_count, package_qs = export.package_prefetched_qs(self.checkpoint)
         grouped_packages = itertools.groupby(
             package_qs.iterator(chunk_size=chunk_size),
             key=attrgetter("type", "namespace", "name", "version"),
@@ -112,11 +106,11 @@ def publish_package_related_advisories(self):
             logger=self.log,
         )
         for _, packages in progress.iter(grouped_packages):
-            purl, package_vulnerabilities = get_package_related_advisory(packages)
+            purl, package_vulnerabilities = export.get_package_related_advisory(packages)
             package_repo, datafile_path = self.data_cluster.get_datafile_repo_and_path(purl)
             package_vulnerability_path = f"packages/{package_repo}/{datafile_path}"
 
-            write_file(
+            export.write_file(
                 repo_path=self.repo_path,
                 file_path=package_vulnerability_path,
                 data=package_vulnerabilities,
@@ -156,7 +150,7 @@ def publish_advisories(self):
         batch_size = 4000
         chunk_size = 1000
         files_to_commit = set()
-        advisory_qs = advisory_prefetched_qs(self.checkpoint)
+        advisory_qs = export.advisory_prefetched_qs(self.checkpoint)
         advisory_count = advisory_qs.count()
 
         self.log(f"Exporting {advisory_count} advisory.")
@@ -166,9 +160,9 @@ def publish_advisories(self):
             logger=self.log,
         )
         for advisory in progress.iter(advisory_qs.iterator(chunk_size=chunk_size)):
-            advisory_data = serialize_advisory(advisory)
+            advisory_data = export.serialize_advisory(advisory)
             adv_file = f"advisories/{advisory.avid}.yml"
-            write_file(
+            export.write_file(
                 repo_path=self.repo_path,
                 file_path=adv_file,
                 data=advisory_data,
@@ -232,135 +226,3 @@ def commit_message(
             commit_count=commit_count,
             total_commit_count=total_commit_count,
         )
-
-
-def package_prefetched_qs(checkpoint):
-    count = None
-    qs = (
-        PackageV2.objects.order_by("type", "namespace", "name", "version")
-        .only("package_url", "type", "namespace", "name", "version")
-        .prefetch_related(
-            Prefetch(
-                "affected_in_impacts",
-                queryset=ImpactedPackage.objects.only("advisory_id").prefetch_related(
-                    Prefetch(
-                        "advisory",
-                        queryset=AdvisoryV2.objects.only("avid"),
-                    )
-                ),
-            ),
-            Prefetch(
-                "fixed_in_impacts",
-                queryset=ImpactedPackage.objects.only("advisory_id").prefetch_related(
-                    Prefetch(
-                        "advisory",
-                        queryset=AdvisoryV2.objects.only("avid"),
-                    )
-                ),
-            ),
-        )
-    )
-
-    if checkpoint:
-        affected_package_ids_qs = (
-            ImpactedPackageAffecting.objects.filter(created_at__gte=checkpoint)
-            .values_list("package_id", flat=True)
-            .distinct()
-        )
-        fixing_package_ids_qs = (
-            ImpactedPackageFixedBy.objects.filter(created_at__gte=checkpoint)
-            .values_list("package_id", flat=True)
-            .distinct()
-        )
-
-        updated_packages = affected_package_ids_qs.union(fixing_package_ids_qs)
-        count = updated_packages.count()
-        qs = qs.filter(id__in=updated_packages)
-
-    count = qs.count() if not count else count
-
-    return count, qs
-
-
-def get_package_related_advisory(packages):
-    package_vulnerabilities = []
-    for package in packages:
-        affected_by_vulnerabilities = [
-            impact.advisory.avid for impact in package.affected_in_impacts.all()
-        ]
-        fixing_vulnerabilities = [impact.advisory.avid for impact in package.fixed_in_impacts.all()]
-
-        package_vulnerability = {
-            "purl": package.package_url,
-            "affected_by_advisories": sorted(affected_by_vulnerabilities),
-            "fixing_advisories": sorted(fixing_vulnerabilities),
-        }
-        package_vulnerabilities.append(package_vulnerability)
-
-    return package.package_url, package_vulnerabilities
-
-
-def advisory_prefetched_qs(checkpoint):
-    qs = AdvisoryV2.objects.order_by("date_collected").prefetch_related(
-        "impacted_packages",
-        "aliases",
-        "references",
-        "severities",
-        "weaknesses",
-    )
-
-    return qs.filter(date_collected__gte=checkpoint) if checkpoint else qs
-
-
-def serialize_severity(sev):
-    return {
-        "score": sev.value,
-        "scoring_system": sev.scoring_system,
-        "scoring_elements": sev.scoring_elements,
-        "published_at": str(sev.published_at),
-        "url": sev.url,
-    }
-
-
-def serialize_references(reference):
-    return {
-        "url": reference.url,
-        "reference_type": reference.reference_type,
-        "reference_id": reference.reference_id,
-    }
-
-
-def serialize_advisory(advisory):
-    """Return a plain data mapping serialized from advisory object."""
-    aliases = sorted([a.alias for a in advisory.aliases.all()])
-    severities = [serialize_severity(sev) for sev in advisory.severities.all()]
-    weaknesses = [wkns.cwe for wkns in advisory.weaknesses.all()]
-    references = [serialize_references(ref) for ref in advisory.references.all()]
-    impacts = [
-        {
-            "purl": impact.base_purl,
-            "affected_versions": impact.affecting_vers,
-            "fixed_versions": impact.fixed_vers,
-        }
-        for impact in advisory.impacted_packages.all()
-    ]
-
-    return {
-        "advisory_id": advisory.advisory_id,
-        "datasource_id": advisory.avid,
-        "datasource_url": advisory.url,
-        "aliases": aliases,
-        "summary": advisory.summary,
-        "impacted_packages": impacts,
-        "severities": severities,
-        "weaknesses": weaknesses,
-        "references": references,
-    }
-
-
-def write_file(repo_path, file_path, data):
-    """Write ``data`` as YAML to ``repo_path``."""
-    write_to = repo_path / file_path
-    write_to.parent.mkdir(parents=True, exist_ok=True)
-    with open(write_to, encoding="utf-8", mode="w") as f:
-        f.write(saneyaml.dump(data))

vulnerabilities/pipes/export.py

@@ -0,0 +1,149 @@
+# Copyright (c) nexB Inc. and others. All rights reserved.
+# VulnerableCode is a trademark of nexB Inc.
+# SPDX-License-Identifier: Apache-2.0
+# See http://www.apache.org/licenses/LICENSE-2.0 for the license text.
+# See https://github.com/aboutcode-org/vulnerablecode for support or download.
+# See https://aboutcode.org for more information about nexB OSS projects.
+#
+
+import saneyaml
+from django.db.models import Prefetch
+
+from vulnerabilities.models import AdvisoryV2
+from vulnerabilities.models import ImpactedPackage
+from vulnerabilities.models import ImpactedPackageAffecting
+from vulnerabilities.models import ImpactedPackageFixedBy
+from vulnerabilities.models import PackageCommitPatch
+from vulnerabilities.models import PackageV2
+
+
+def package_prefetched_qs(checkpoint):
+    count = None
+    qs = (
+        PackageV2.objects.order_by("type", "namespace", "name", "version")
+        .only("package_url", "type", "namespace", "name", "version")
+        .prefetch_related(
+            Prefetch(
+                "affected_in_impacts",
+                queryset=ImpactedPackage.objects.only("advisory_id").prefetch_related(
+                    Prefetch(
+                        "advisory",
+                        queryset=AdvisoryV2.objects.only("avid"),
+                    )
+                ),
+            ),
+            Prefetch(
+                "fixed_in_impacts",
+                queryset=ImpactedPackage.objects.only("advisory_id").prefetch_related(
+                    Prefetch(
+                        "advisory",
+                        queryset=AdvisoryV2.objects.only("avid"),
+                    )
+                ),
+            ),
+        )
+    )
+
+    if checkpoint:
+        affected_package_ids_qs = (
+            ImpactedPackageAffecting.objects.filter(created_at__gte=checkpoint)
+            .values_list("package_id", flat=True)
+            .distinct()
+        )
+        fixing_package_ids_qs = (
+            ImpactedPackageFixedBy.objects.filter(created_at__gte=checkpoint)
+            .values_list("package_id", flat=True)
+            .distinct()
+        )
+
+        updated_packages = affected_package_ids_qs.union(fixing_package_ids_qs)
+        count = updated_packages.count()
+        qs = qs.filter(id__in=updated_packages)
+
+    count = qs.count() if not count else count
+
+    return count, qs
+
+
+def get_package_related_advisory(packages):
+    package_vulnerabilities = []
+    for package in packages:
+        affected_by_vulnerabilities = [
+            impact.advisory.avid for impact in package.affected_in_impacts.all()
+        ]
+        fixing_vulnerabilities = [impact.advisory.avid for impact in package.fixed_in_impacts.all()]
+
+        package_vulnerability = {
+            "purl": package.package_url,
+            "affected_by_advisories": sorted(affected_by_vulnerabilities),
+            "fixing_advisories": sorted(fixing_vulnerabilities),
+        }
+        package_vulnerabilities.append(package_vulnerability)
+
+    return package.package_url, package_vulnerabilities
+
+
+def advisory_prefetched_qs(checkpoint):
+    qs = AdvisoryV2.objects.order_by("date_collected").prefetch_related(
+        "impacted_packages",
+        "aliases",
+        "references",
+        "severities",
+        "weaknesses",
+    )
+
+    return qs.filter(date_collected__gte=checkpoint) if checkpoint else qs
+
+
+def serialize_severity(sev):
+    return {
+        "score": sev.value,
+        "scoring_system": sev.scoring_system,
+        "scoring_elements": sev.scoring_elements,
+        "published_at": str(sev.published_at),
+        "url": sev.url,
+    }
+
+
+def serialize_references(reference):
+    return {
+        "url": reference.url,
+        "reference_type": reference.reference_type,
+        "reference_id": reference.reference_id,
+    }
+
+
+def serialize_advisory(advisory):
+    """Return a plain data mapping serialized from advisory object."""
+    aliases = sorted([a.alias for a in advisory.aliases.all()])
+    severities = [serialize_severity(sev) for sev in advisory.severities.all()]
+    weaknesses = [wkns.cwe for wkns in advisory.weaknesses.all()]
+    references = [serialize_references(ref) for ref in advisory.references.all()]
+    impacts = [
+        {
+            "purl": impact.base_purl,
+            "affected_versions": impact.affecting_vers,
+            "fixed_versions": impact.fixed_vers,
+        }
+        for impact in advisory.impacted_packages.all()
+    ]
+
+    return {
+        "advisory_id": advisory.advisory_id,
+        "datasource_id": advisory.avid,
+        "datasource_url": advisory.url,
+        "aliases": aliases,
+        "summary": advisory.summary,
+        "impacted_packages": impacts,
+        "severities": severities,
+        "weaknesses": weaknesses,
+        "references": references,
+    }
+
+
+def write_file(repo_path, file_path, data):
+    """Write ``data`` as YAML to ``repo_path``."""
+    write_to = repo_path / file_path
+    write_to.parent.mkdir(parents=True, exist_ok=True)
+    with open(write_to, encoding="utf-8", mode="w") as f:
+        f.write(saneyaml.dump(data))