diff --git a/vulnerabilities/data_source.py b/vulnerabilities/data_source.py index 05928a460..ca68d1b83 100644 --- a/vulnerabilities/data_source.py +++ b/vulnerabilities/data_source.py @@ -75,7 +75,7 @@ class InvalidConfigurationError(Exception): @dataclasses.dataclass class DataSourceConfiguration: - batch_size: int + pass class DataSource(ContextManager): @@ -105,8 +105,9 @@ def __init__( :param config: Optional dictionary with subclass-specific configuration """ config = config or {} + self.batch_size = batch_size try: - self.config = self.__class__.CONFIG_CLASS(batch_size, **config) + self.config = self.__class__.CONFIG_CLASS(**config) # These really should be declared in DataSourceConfiguration above but that would # prevent DataSource subclasses from declaring mandatory parameters (i.e. positional # arguments) @@ -183,7 +184,7 @@ def batch_advisories(self, advisories: List[Advisory]) -> Set[Advisory]: advisories = advisories[:] # copy the list as we are mutating it in the loop below while advisories: - b, advisories = advisories[:self.config.batch_size], advisories[self.config.batch_size:] + b, advisories = advisories[:self.batch_size], advisories[self.batch_size:] yield set(b) diff --git a/vulnerabilities/import_runner.py b/vulnerabilities/import_runner.py index c8176d1b5..9b049fd3d 100644 --- a/vulnerabilities/import_runner.py +++ b/vulnerabilities/import_runner.py @@ -21,6 +21,7 @@ # VulnerableCode is a free software code scanning tool from nexB Inc. and others. # Visit https://github.com/nexB/vulnerablecode/ for support and download. +import dataclasses import datetime import logging from typing import Dict @@ -77,6 +78,7 @@ def run(self, cutoff_date: datetime.datetime = None) -> None: _process_updated_advisories(data_source) self.importer.last_run = datetime.datetime.now(tz=datetime.timezone.utc) + self.importer.data_source_cfg = dataclasses.asdict(data_source.config) self.importer.save() logger.debug(f'Successfully finished import for {self.importer.name}.') diff --git a/vulnerabilities/importers/rust.py b/vulnerabilities/importers/rust.py index 6c085a899..dc5ac2e42 100644 --- a/vulnerabilities/importers/rust.py +++ b/vulnerabilities/importers/rust.py @@ -64,7 +64,7 @@ def _load_advisories(self, files) -> Set[Advisory]: files = [f for f in files if not f.endswith('-0000.toml')] # skip temporary files while files: - batch, files = files[:self.config.batch_size], files[self.config.batch_size:] + batch, files = files[:self.batch_size], files[self.batch_size:] advisories = set() diff --git a/vulnerabilities/tests/test_import_runner.py b/vulnerabilities/tests/test_import_runner.py index 44b1dd897..ce5572c8b 100644 --- a/vulnerabilities/tests/test_import_runner.py +++ b/vulnerabilities/tests/test_import_runner.py @@ -46,7 +46,7 @@ def updated_advisories(self): def _yield_advisories(self, advisories): while advisories: - b, advisories = advisories[:self.config.batch_size], advisories[self.config.batch_size:] + b, advisories = advisories[:self.batch_size], advisories[self.batch_size:] yield b