Update the pipeline to use bulk

Signed-off-by: ziad hany <ziadhany2016@gmail.com>

Author: ziadhany

vulnerabilities/pipelines/v2_improvers/reference_collect_commits.py

@@ -8,18 +8,21 @@
 #
 
 from aboutcode.pipeline import LoopProgress
+from django.db.models import Prefetch
 from packageurl.contrib.purl2url import purl2url
 from packageurl.contrib.url2purl import url2purl
 
 from aboutcode.federated import get_core_purl
+from vulnerabilities.models import AdvisoryReference
 from vulnerabilities.models import AdvisoryV2
+from vulnerabilities.models import ImpactedPackage
 from vulnerabilities.models import PackageCommitPatch
-from vulnerabilities.pipelines import VulnerableCodeBaseImporterPipelineV2
-from vulnerabilities.pipes.advisory import VCS_URLS_SUPPORTED_TYPES
+from vulnerabilities.models import Patch
+from vulnerabilities.pipelines import VulnerableCodePipeline
 from vulnerabilities.utils import is_commit
 
 
-class CollectReferencesFixCommitsPipeline(VulnerableCodeBaseImporterPipelineV2):
+class CollectReferencesFixCommitsPipeline(VulnerableCodePipeline):
     """
     Improver pipeline to scout References/Patch and create PackageCommitPatch entries.
     """
@@ -30,45 +33,98 @@ class CollectReferencesFixCommitsPipeline(VulnerableCodeBaseImporterPipelineV2):
     def steps(cls):
         return (cls.collect_and_store_fix_commits,)
 
-    def get_vcs_commit(self, url):
-        """Extracts and VCS URL and commit hash from URL.
+    def get_vcs_data(self, url):
+        """Extracts a VCS URL and commit hash from URL.
         >> get_vcs_commit('https://github.com/aboutcode-org/vulnerablecode/commit/98e516011d6e096e25247b82fc5f196bbeecff10')
-        ('https://github.com/aboutcode-org/vulnerablecode', '98e516011d6e096e25247b82fc5f196bbeecff10')
+        ("pkg:github/aboutcode-org/vulnerablecode", 'https://github.com/aboutcode-org/vulnerablecode', '98e516011d6e096e25247b82fc5f196bbeecff10')
         >> get_vcs_commit('https://github.com/aboutcode-org/vulnerablecode/pull/1974')
         None
         """
-        purl = url2purl(url)
-        if not purl or purl.type not in VCS_URLS_SUPPORTED_TYPES:
-            return None
+        try:
+            purl = url2purl(url)
+            if not purl:
+                return
 
-        version = getattr(purl, "version", None)
-        if not version or not is_commit(version):
-            return None
-
-        vcs_url = purl2url(get_core_purl(purl).to_string())
-        return (vcs_url, version) if vcs_url else None
+            version = purl.version
+            if not version or not is_commit(version):
+                return
+            base_purl = get_core_purl(purl)
+            vcs_url = purl2url(base_purl.to_string())
+            if base_purl and vcs_url and version:
+                return base_purl, vcs_url, version
+        except Exception as e:
+            self.log(f"Invalid URL: url:{url} error:{e}")
 
     def collect_and_store_fix_commits(self):
-        impacted_packages_advisories = (
-            AdvisoryV2.objects.filter(impacted_packages__isnull=False)
-            .prefetch_related("references", "patches", "impacted_packages")
-            .distinct()
+        advisories = AdvisoryV2.objects.only("id").prefetch_related(
+            Prefetch("references", queryset=AdvisoryReference.objects.only("url")),
+            Prefetch("patches", queryset=Patch.objects.only("patch_url")),
         )
 
-        progress = LoopProgress(
-            total_iterations=impacted_packages_advisories.count(), logger=self.log
-        )
-        for adv in progress.iter(impacted_packages_advisories.paginated(per_page=500)):
+        progress = LoopProgress(total_iterations=advisories.count(), logger=self.log)
+
+        commit_batch = []
+        updated_pkg_patch_commit_count = 0
+        batch_size = 1000
+        for adv in progress.iter(advisories.paginated(per_page=batch_size)):
             urls = {r.url for r in adv.references.all()} | {p.patch_url for p in adv.patches.all()}
-            impacted_packages = list(adv.impacted_packages.all())
 
             for url in urls:
-                vcs_data = self.get_vcs_commit(url)
+                vcs_data = self.get_vcs_data(url)
                 if not vcs_data:
                     continue
+                base_purl, vcs_url, commit_hash = vcs_data
+                commit_batch.append((str(base_purl), vcs_url, commit_hash, adv.id))
+
+            if len(commit_batch) >= batch_size:
+                updated_pkg_patch_commit_count += self.bulk_commit_batch_update(commit_batch)
+                commit_batch.clear()
+
+        if commit_batch:
+            updated_pkg_patch_commit_count += self.bulk_commit_batch_update(commit_batch)
+            commit_batch.clear()
+
+        self.log(f"Successfully processed pkg patch commit {updated_pkg_patch_commit_count:,d}")
+
+    def bulk_commit_batch_update(self, vcs_data_table):
+        impact_data = {(row[0], row[3]) for row in vcs_data_table}  # base_purl, adv_id
+        commit_data = {(row[1], row[2]) for row in vcs_data_table}  # vcs_url, commit_hash
+
+        adv_ids = {aid for _, aid in impact_data}
+        existing_impacts = ImpactedPackage.objects.filter(advisory_id__in=adv_ids)
+        existing_impact_pairs = {(ip.base_purl, ip.advisory_id) for ip in existing_impacts}
+
+        new_impacts = impact_data - existing_impact_pairs
+        if new_impacts:
+            ImpactedPackage.objects.bulk_create(
+                [ImpactedPackage(base_purl=bp, advisory_id=aid) for bp, aid in new_impacts]
+            )
+
+        PackageCommitPatch.objects.bulk_create(
+            [
+                PackageCommitPatch(vcs_url=vcs_url, commit_hash=commit_hash)
+                for vcs_url, commit_hash in commit_data
+            ],
+            ignore_conflicts=True,
+        )
+
+        adv_ids = {adv_id for _, adv_id in impact_data}
+        fetched_impacts = {
+            (impacted_pkg.base_purl, impacted_pkg.advisory_id): impacted_pkg
+            for impacted_pkg in ImpactedPackage.objects.filter(advisory_id__in=adv_ids)
+        }
+
+        commit_hashes = {commit_hash for _, commit_hash in commit_data}
+        fetched_commits = {
+            (pkg_commit_patch.vcs_url, pkg_commit_patch.commit_hash): pkg_commit_patch
+            for pkg_commit_patch in PackageCommitPatch.objects.filter(commit_hash__in=commit_hashes)
+        }
+
+        for base_purl, vcs_url, commit_hash, adv_id in vcs_data_table:
+            impacted_package = fetched_impacts.get((base_purl, adv_id))
+            package_commit_obj = fetched_commits.get((vcs_url, commit_hash))
+
+            if impacted_package and package_commit_obj:
+                package_commit_obj.fixed_in_impacts.add(impacted_package)
 
-                vcs_url, commit_hash = vcs_data
-                package_commit_obj, _ = PackageCommitPatch.objects.get_or_create(
-                    vcs_url=vcs_url, commit_hash=commit_hash
-                )
-                package_commit_obj.fixed_in_impacts.add(*impacted_packages)
+        return len(vcs_data_table)

vulnerabilities/tests/pipelines/v2_improvers/test_reference_collect_commits_v2.py

@@ -30,28 +30,25 @@ def test_collect_fix_commits_pipeline_creates_entry():
         unique_content_id="11111",
         date_collected=datetime.now(),
     )
-    package = PackageV2.objects.create(
-        type="foo",
-        name="testpkg",
-        version="1.0",
-    )
+
     reference = AdvisoryReference.objects.create(
         url="https://github.com/test/testpkg/commit/6bd301819f8f69331a55ae2336c8b111fc933f3d"
     )
-    impact = ImpactedPackage.objects.create(advisory=advisory)
-    impact.affecting_packages.add(package)
     advisory.references.add(reference)
 
     pipeline = CollectReferencesFixCommitsPipeline()
     pipeline.collect_and_store_fix_commits()
 
     package_commit_patch = PackageCommitPatch.objects.all()
+    impacted_packages = advisory.impacted_packages.all()
 
     assert package_commit_patch.count() == 1
+    assert impacted_packages.count() == 1
+
     fix = package_commit_patch.first()
     assert fix.commit_hash == "6bd301819f8f69331a55ae2336c8b111fc933f3d"
     assert fix.vcs_url == "https://github.com/test/testpkg"
-    assert impact.fixed_by_package_commit_patches.count() == 1
+    assert impacted_packages.first().fixed_by_package_commit_patches.count() == 1
 
 
 @pytest.mark.django_db
@@ -64,13 +61,6 @@ def test_collect_fix_commits_pipeline_skips_non_commit_urls():
         unique_content_id="11111",
         date_collected=datetime.now(),
     )
-    package = PackageV2.objects.create(
-        type="pypi",
-        name="otherpkg",
-        version="2.0",
-    )
-    impact = ImpactedPackage.objects.create(advisory=advisory)
-    impact.affecting_packages.add(package)
 
     reference = AdvisoryReference.objects.create(
         url="https://github.com/test/testpkg/issues/12"