Source code for swh.osv.luigi.cherrypicks

# Copyright (C) 2026  The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information


"""
Luigi tasks for listing cherry-picks
====================================
"""

from os import PathLike
from pathlib import Path
from typing import Dict, Iterator, Tuple

import luigi

from swh.graph.luigi.compressed_graph import LocalGraph
from swh.osv.luigi import DownloadVulnerabilities
from swh.osv.luigi.reachability import ComputeConnectedComponents
from swh.osv.luigi.utils import _BaseUploadToS3




class ListCherrypicks(luigi.Task):
    """Lists cherry-picks using subgraph node ids (without transitive reduction)."""

    local_graph_path = luigi.PathParameter()
    graph_name = luigi.StrParameter(default="graph")
    vulnerabilities_dir = luigi.PathParameter()
    include_vcs_matches = luigi.BoolParameter(default=True)

    @property
    def resources(self):
        """Declares the task uses every CPU available"""
        import socket

        hostname = socket.getfqdn()
        return {
            f"{hostname}_max_cpu": 1,
        }



    def requires(self) -> Dict[str, luigi.Task]:
        """Returns instances of :class:`swh.graph.luigi.compressed_graph.LocalGraph`,
        :class:`swh.osv.luigi.DownloadVulnerabilities` and
        :class:`swh.osv.luigi.reachability.ComputeConnectedComponents`."""
        return {
            "graph": LocalGraph(local_graph_path=self.local_graph_path),
            "sqlite": DownloadVulnerabilities(
                local_graph_path=self.local_graph_path,
                graph_name=self.graph_name,
                vulnerabilities_dir=self.vulnerabilities_dir,
            ),
            "subgraphwccs": ComputeConnectedComponents(
                local_graph_path=self.local_graph_path,
                graph_name=self.graph_name,
                vulnerabilities_dir=self.vulnerabilities_dir,
                include_vcs_matches=self.include_vcs_matches,
            ),
        }




    def output(self) -> luigi.LocalTarget:
        """``cherrypicks_subgraphnodeids`` directory containing cherrypick multimaps
        and their EF indices."""
        return luigi.LocalTarget(
            self.vulnerabilities_dir / "cherrypicks_subgraphnodeids"
        )




    def run(self) -> None:
        """Runs ``list-cherrypicks`` and ``build-ef`` on each multimap."""
        from .shell import Rust

        out_dir = Path(self.output().path)
        out_dir.mkdir(parents=True, exist_ok=True)

        cherrypicks_base = out_dir / "cherrypicks2originals"
        originals_base = out_dir / "originals2cherrypicks"
        # fmt: off
        Rust(
            "list-cherrypicks",
            "--graph", self.local_graph_path / self.graph_name,
            "--db", self.input()["sqlite"],
            "--subgraphwccs", self.input()["subgraphwccs"],
            "--cherrypicks2originals-out", cherrypicks_base,
            "--originals2cherrypicks-out", originals_base,
            "--subgraph-nodeids",
        ).run()
        Rust(
            "build-ef",
            "--graph", cherrypicks_base,
        ).run()
        Rust(
            "build-ef",
            "--graph", originals_base,
        ).run()
        # fmt: on
        (out_dir / ".stamp").touch()






class UploadCherrypicksToS3(_BaseUploadToS3):
    """Uploads the cherry-picks multimaps to S3."""

    include_vcs_matches = luigi.BoolParameter(default=True)



    def requires(self) -> Dict[str, luigi.Task]:
        """Returns an instance of :class:`ListCherrypicks`."""
        return {
            "cherrypicks": ListCherrypicks(
                local_graph_path=self.local_graph_path,
                graph_name=self.graph_name,
                vulnerabilities_dir=self.vulnerabilities_dir,
                include_vcs_matches=self.include_vcs_matches,
            ),
        }




    def output(self) -> luigi.Target:
        """S3 target for the cherrypick multimap directory stamp."""
        import luigi.contrib.s3

        return luigi.contrib.s3.S3Target(
            f"{self._s3_prefix()}/cherrypicks_subgraphnodeids/.stamp"
        )


    def _files_to_upload(self) -> Iterator[Tuple[PathLike, str]]:
        """Yields all cherrypick multimap files for upload."""
        local_dir = Path(self.input()["cherrypicks"].path)
        for path in sorted(local_dir.rglob("*")):
            if path.is_dir():
                continue
            yield (
                path,
                f"cherrypicks_subgraphnodeids/{path.relative_to(local_dir)}",
            )