Source code for swh.osv.luigi.reachability

# 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 computing reachability information
==================================================
"""

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

import luigi

from swh.graph.luigi.compressed_graph import LocalGraph
from swh.osv.luigi import DownloadVulnerabilities, MapPackagesToVcs
from swh.osv.luigi.utils import _BaseUploadToS3




class ComputeConnectedComponents(luigi.Task):
    """Computes the connected components of the subgraph spanned by commits
    referenced in vulnerability reports."""

    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 an instance of :class:`swh.graph.luigi.compressed_graph.LocalGraph`
        and :class:`swh.osv.luigi.DownloadVulnerabilities`. If
        :attr:`include_vcs_matches` is True, also returns an instance of
        :class:`swh.osv.luigi.MapPackagesToVcs`."""
        result: Dict[str, luigi.Task] = {
            "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,
            ),
        }
        if self.include_vcs_matches:
            result["packages_to_vcs"] = MapPackagesToVcs(
                local_graph_path=self.local_graph_path,
                graph_name=self.graph_name,
                vulnerabilities_dir=self.vulnerabilities_dir,
            )
        return result




    def output(self) -> luigi.LocalTarget:
        """``connected_components.wccs`` file containing the epserde-serialized
        subgraph connected components."""
        return luigi.LocalTarget(self.vulnerabilities_dir / "connected_components.wccs")




    def run(self) -> None:
        """Runs ``compute-connected-components``."""
        from .shell import Rust

        # fmt: off
        args = [
            "compute-connected-components",
            "--graph", self.local_graph_path / self.graph_name,
            "--db", self.input()["sqlite"],
            "--subgraphwccs-out", self.output(),
        ]
        if self.include_vcs_matches:
            args += [
                "--package2vcs", self.input()["packages_to_vcs"],
                "--include-vcs-matches",
            ]
        Rust(*args).run()


        # fmt: on




class UploadSubgraphWccsToS3(_BaseUploadToS3):
    """Uploads the connected components WCCS file to S3."""

    include_vcs_matches = luigi.BoolParameter(default=True)



    def requires(self) -> Dict[str, luigi.Task]:
        """Returns an instance of :class:`ComputeConnectedComponents`."""
        return {
            "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.Target:
        """``connected_components.wccs`` on S3."""
        import luigi.contrib.s3

        return luigi.contrib.s3.S3Target(
            f"{self._s3_prefix()}/connected_components.wccs"
        )


    def _files_to_upload(self) -> Iterator[Tuple[PathLike, str]]:
        yield (self.input()["subgraphwccs"].path, "connected_components.wccs")