Source code for swh.lister.crates.lister

# Copyright (C) 2022-2023  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

import csv
from dataclasses import dataclass
from datetime import datetime
import json
import logging
from pathlib import Path
import tarfile
import tempfile
from typing import Any, Dict, Iterator, List, Optional
from urllib.parse import urlparse

import iso8601
from packaging.version import parse as parse_version

from swh.scheduler.interface import SchedulerInterface
from swh.scheduler.model import ListedOrigin

from ..pattern import CredentialsType, Lister

logger = logging.getLogger(__name__)

# Aliasing the page results returned by `get_pages` method from the lister.
CratesListerPage = List[Dict[str, Any]]

[docs] @dataclass class CratesListerState: """Store lister state for incremental mode operations. 'index_last_update' represents the UTC time the database dump was started """ index_last_update: Optional[datetime] = None
[docs] class CratesLister(Lister[CratesListerState, CratesListerPage]): """List origins from the "" forge. It downloads a tar.gz archive which contains database table content as csv files which is automatically generated every 24 hours. Parsing two csv files we can list all package names and their related versions. In incremental mode, it check each entry comparing their 'last_update' value with self.state.index_last_update """ LISTER_NAME = "crates" VISIT_TYPE = "crates" INSTANCE = "crates" BASE_URL = "" DB_DUMP_URL = "" CRATE_FILE_URL_PATTERN = ( "{crate}/{crate}-{version}.crate" ) CRATE_URL_PATTERN = "{crate}" def __init__( self, scheduler: SchedulerInterface, url: str = BASE_URL, instance: str = INSTANCE, credentials: CredentialsType = None, max_origins_per_page: Optional[int] = None, max_pages: Optional[int] = None, enable_origins: bool = True, ): super().__init__( scheduler=scheduler, credentials=credentials, url=url, instance=instance, max_origins_per_page=max_origins_per_page, max_pages=max_pages, enable_origins=enable_origins, ) self.index_metadata: Dict[str, str] = {}
[docs] def state_from_dict(self, d: Dict[str, Any]) -> CratesListerState: index_last_update = d.get("index_last_update") if index_last_update is not None: d["index_last_update"] = iso8601.parse_date(index_last_update) return CratesListerState(**d)
[docs] def state_to_dict(self, state: CratesListerState) -> Dict[str, Any]: d: Dict[str, Optional[str]] = {"index_last_update": None} index_last_update = state.index_last_update if index_last_update is not None: d["index_last_update"] = index_last_update.isoformat() return d
[docs] def is_new(self, dt_str: str): """Returns True when dt_str is greater than self.state.index_last_update """ dt = iso8601.parse_date(dt_str) last = self.state.index_last_update return not last or (last is not None and last < dt)
[docs] def get_and_parse_db_dump(self) -> Dict[str, Any]: """Download and parse csv files from db_dump_path. Returns a dict where each entry corresponds to a package name with its related versions. """ with tempfile.TemporaryDirectory() as tmpdir: file_name = self.DB_DUMP_URL.split("/")[-1] archive_path = Path(tmpdir) / file_name # Download the Db dump with self.http_request(self.DB_DUMP_URL, stream=True) as res: with open(archive_path, "wb") as out_file: for chunk in res.iter_content(chunk_size=1024): out_file.write(chunk) # Extract the Db dump db_dump_path = Path(str(archive_path).split(".tar.gz")[0]) members_to_extract = [] with as tf: for member in tf.getmembers(): if ("/data/crates.csv", "/data/versions.csv", "/metadata.json") ): members_to_extract.append(member) tf.extractall(members=members_to_extract, path=db_dump_path) csv.field_size_limit(1000000) (crates_csv_path,) = list(db_dump_path.glob("*/data/crates.csv")) (versions_csv_path,) = list(db_dump_path.glob("*/data/versions.csv")) (index_metadata_json_path,) = list(db_dump_path.rglob("*/metadata.json")) with"rb") as index_metadata_json: self.index_metadata = json.load(index_metadata_json) crates: Dict[str, Any] = {} with as crates_fd: crates_csv = csv.DictReader(crates_fd) for item in crates_csv: if self.is_new(item["updated_at"]): # crate 'id' as key crates[item["id"]] = { "name": item["name"], "updated_at": item["updated_at"], "versions": {}, } data: Dict[str, Any] = {} with as versions_fd: versions_csv = csv.DictReader(versions_fd) for version in versions_csv: if version["crate_id"] in crates.keys(): crate: Dict[str, Any] = crates[version["crate_id"]] crate["versions"][version["num"]] = version # crate 'name' as key data[crate["name"]] = crate return data
[docs] def page_entry_dict(self, entry: Dict[str, Any]) -> Dict[str, Any]: """Transform package version definition dict to a suitable page entry dict """ crate_file = self.CRATE_FILE_URL_PATTERN.format( crate=entry["name"], version=entry["version"] ) filename = urlparse(crate_file).path.split("/")[-1] return dict( name=entry["name"], version=entry["version"], checksum=entry["checksum"], yanked=True if entry["yanked"] == "t" else False, crate_file=crate_file, filename=filename, last_update=entry["updated_at"], )
[docs] def get_pages(self) -> Iterator[CratesListerPage]: """Each page is a list of crate versions with: - name: Name of the crate - version: Version - checksum: Checksum - yanked: Whether the package is yanked or not - crate_file: Url of the crate file - filename: File name of the crate file - last_update: Last update for that version """ # Fetch Db dump, then Parse the data. dataset = self.get_and_parse_db_dump() logger.debug("Found %s crates in crates_index", len(dataset)) # Each entry from dataset will correspond to a page for name, item in dataset.items(): page = [] # sort crate versions versions: list = sorted(item["versions"].keys(), key=parse_version) for version in versions: v = item["versions"][version] v["name"] = name v["version"] = version page.append(self.page_entry_dict(v)) yield page
[docs] def get_origins_from_page(self, page: CratesListerPage) -> Iterator[ListedOrigin]: """Iterate on all crate pages and yield ListedOrigin instances.""" assert is not None url = self.CRATE_URL_PATTERN.format(crate=page[0]["name"]) last_update = page[0]["last_update"] artifacts = [] crates_metadata = [] for entry in page: # Build an artifact entry following original-artifacts-json specification # # noqa: B950 artifacts.append( { "version": entry["version"], "filename": entry["filename"], "url": entry["crate_file"], "checksums": { "sha256": entry["checksum"], }, } ) crates_metadata.append( { "version": entry["version"], "yanked": entry["yanked"], "last_update": entry["last_update"], } ) yield ListedOrigin(, visit_type=self.VISIT_TYPE, url=url, last_update=iso8601.parse_date(last_update), extra_loader_arguments={ "artifacts": artifacts, "crates_metadata": crates_metadata, }, )
[docs] def finalize(self) -> None: last: datetime = iso8601.parse_date(self.index_metadata["timestamp"]) if not self.state.index_last_update: self.state.index_last_update = last self.updated = True