Source code for swh.osv.download_nvd_osv

# Copyright (C) 2025-2026  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 contextlib
import multiprocessing.dummy
from pathlib import Path
import tarfile
import threading

from google.cloud import storage
import pyzstd
import tqdm




def download_to_tarballs(path: Path) -> None:
    storage_client = storage.Client.create_anonymous_client()
    bucket = storage_client.bucket(bucket_name="cve-osv-conversion")
    try:
        path.mkdir()
    except FileExistsError:
        pass
    blobs = list(
        tqdm.tqdm(
            bucket.list_blobs(prefix="osv-output/"),
            desc="Listing blobs",
        )
    )
    with contextlib.ExitStack() as stack:
        threadlocal = threading.local()
        with multiprocessing.dummy.Pool(
            128, initializer=lambda: _open_tarfile(threadlocal, stack, path)
        ) as p:
            for _ in tqdm.tqdm(
                p.imap_unordered(lambda blob: _download_blob(blob, threadlocal), blobs),
                total=len(blobs),
                desc="Downloading blobs",
            ):
                pass



def _open_tarfile(
    threadlocal: threading.local, stack: contextlib.ExitStack, path: Path
) -> None:
    fd = stack.enter_context(
        pyzstd.open(
            str(path / f"shard_{threading.current_thread().ident}.tar.zst"), "w"
        )
    )
    threadlocal.tarfile = stack.enter_context(tarfile.open(fileobj=fd, mode="w"))


def _download_blob(blob: storage.blob.Blob, threadlocal: threading.local) -> None:
    if blob.name == "osv-output/.placeholder":
        return
    tarinfo = tarfile.TarInfo(name=blob.name)
    tarinfo.size = blob.size
    threadlocal.tarfile.addfile(tarinfo, blob.open("rb"))


if __name__ == "__main__":
    download_to_tarballs(Path("./nvd_cve"))