Source code for swh.osv.map_origins

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

import asyncio
import json
from pathlib import Path

import aiohttp
from google.protobuf.field_mask_pb2 import FieldMask
import grpc.aio

import swh.graph.grpc.swhgraph_pb2 as swhgraph
import swh.graph.grpc.swhgraph_pb2_grpc as swhgraph_grpc
from swh.model.model import Origin

from .utils import Cache

GRAPH_GRPC_URL = "localhost:50091"


SWH_API_SEMAPHORE = asyncio.Semaphore(10)
"""Prevents more than 10 concurrent requests to the SWH API"""


SWH_ORIGIN_URLS_BY_ECOSYSTEM = {
    "Ubuntu": {
        "deb://Ubuntu/packages/{package_name}",
        "deb://Ubuntu-Legacy/packages/{package_name}",
    },
    "npm": {"https://www.npmjs.com/package/{package_name}"},
    "PyPI": {"https://pypi.org/project/{package_name}/"},
}




class InvalidCommitId(ValueError):
    pass





class OriginDenormalizer:
    KNOWN_ORIGINS_CACHE_PATH = Path("./known_origins_urls.json")
    UNKNOWN_ORIGINS_CACHE_PATH = Path("./unknown_origins_urls.json")

    def __init__(
        self,
        stub: swhgraph_grpc.TraversalServiceStub,
        session: aiohttp.ClientSession,
    ):
        self._stub = stub
        self._session = session

        self._denormalization_cache = Cache(
            Path("./denormalize_urls.json"), dict[str, str](), json.dumps, json.loads
        )

        self._known_origins_cache = Cache(
            Path("./known_origins_urls.json"),
            set[str](),
            lambda data: json.dumps(list(data)),
            json.loads,
        )

        self._unknown_origins_cache = Cache(
            Path("./unknown_origins_urls.json"),
            set[str](),
            lambda data: json.dumps(list(data)),
            json.loads,
        )



    async def load(self) -> None:
        await self._denormalization_cache.load()
        await self._known_origins_cache.load()
        await self._unknown_origins_cache.load()




    async def is_origin_url_in_graph(self, url: str) -> bool:
        if (
            url not in self._known_origins_cache
            and url not in self._unknown_origins_cache
        ):
            swhid = Origin(url=url).swhid()
            try:
                await self._stub.GetNode(
                    swhgraph.GetNodeRequest(
                        swhid=str(swhid), mask=FieldMask(paths=["swhid"])
                    )
                )
            except grpc.aio.AioRpcError as e:
                if e.code() == grpc.StatusCode.NOT_FOUND:
                    self._unknown_origins_cache.add(url)
                    await self._unknown_origins_cache.save()
                else:
                    raise
            else:
                self._known_origins_cache.add(url)
                await self._known_origins_cache.save()
        if url in self._known_origins_cache:
            assert (
                url not in self._unknown_origins_cache
            ), f"{url} is both known and unknown"
            return True
        else:
            assert (
                url in self._unknown_origins_cache
            ), f"{url} is neither known nor unknown"
            return False


    def _normalize_origin_url(self, url: str) -> str:
        if url.startswith("https://pypi.org/project/"):
            url = url.replace("_", "-").replace(".", "-")
        return url.lower().removesuffix(".git")



    async def denormalize_origin_url(self, url: str) -> str | None:
        # most of the time, this succeeds because the URL in SWH matches the URL in OSV
        if await self.is_origin_url_in_graph(url):
            return url

        # simple heuristic that covers ~20% of the missing origins
        if url.endswith(".git"):
            if await self.is_origin_url_in_graph(url.removesuffix(".git")):
                return url.removesuffix(".git")

        # OSV URL not found in SWH. This is usually because OSV lowercased the URL
        # while SWH did not.
        # Use SWH's URL search to try to find the non-lowercase URL
        if url in self._denormalization_cache:
            assert url.lower().removesuffix(".git") == self._denormalization_cache[
                url
            ].lower().removesuffix(".git"), (
                url.lower(),
                self._denormalization_cache[url].lower(),
            )
            return self._denormalization_cache[url]
        if url in self._unknown_origins_cache:
            # Skip SWH API request, we already tried and it failed.
            return None

        url = f"https://archive.softwareheritage.org/api/1/origin/search/{url}?limit=10"
        for _ in range(10):  # retry at most 10 times on errors
            async with SWH_API_SEMAPHORE:
                async with self._session.get(url) as r:
                    if 500 <= r.status < 600:
                        print(f"{r.status} error for {url} ; retrying in 10s")
                        await asyncio.sleep(
                            10
                        )  # TODO: release the semaphore while sleeping
                        continue
                    for origin in await r.json():
                        if self._normalize_origin_url(
                            url
                        ) == self._normalize_origin_url(origin["url"]):
                            break
                    else:
                        self._unknown_origins_cache.add(url)
                        await self._unknown_origins_cache.save()
                        return None
            break

        # don't do this in the loop in order to release SWH_API_SEMAPHORE as soon as possible
        self._denormalization_cache[url] = origin["url"]
        await self._denormalization_cache.save()
        return origin["url"]