Add: Add --chunk-size and --queue-size arguments

Allow for more flexible download by adjusting the size of the download/processing queue and the number of items to download in a single request. These arguments are mostly interesting for debugging.
greenbone · Jun 12, 2024 · bb32fb9 · bb32fb9
1 parent 1d9ba6b
commit bb32fb9
Show file tree

Hide file tree

Showing 4 changed files with 93 additions and 8 deletions.
diff --git a/greenbone/scap/cpe/cli/download.py b/greenbone/scap/cpe/cli/download.py
@@ -127,6 +127,23 @@ def parse_args(args: Sequence[str] | None = None) -> Namespace:
         help="Use a NVD API key for downloading the CPEs. Using an API key "
         "allows for downloading with extended rate limits.",
     )
+    parser.add_argument(
+        "--chunk-size",
+        help="Number of CPEs to download and process in one request. A lower "
+        "number allows for more frequent updates and feedback.",
+        type=int,
+        metavar="N",
+    )
+    parser.add_argument(
+        "--queue-size",
+        help="Size of the download queue. It sets the maximum number of CPEs "
+        "kept in the memory. The maximum number of CPEs is chunk size * queue "
+        "size. Default: %(default)s.",
+        type=int,
+        metavar="N",
+        default=DEFAULT_QUEUE_SIZE,
+    )
+
     return parser.parse_args(args)
 
 
@@ -135,10 +152,12 @@ def __init__(
         self,
         console: Console,
         *,
-        queue: asyncio.Queue[Sequence[CPE]] | None = None,
         verbose: int = 0,
+        chunk_size: int | None = None,
+        queue_size: int = DEFAULT_QUEUE_SIZE,
     ) -> None:
-        self.queue = queue or asyncio.Queue(DEFAULT_QUEUE_SIZE)
+        self.queue: asyncio.Queue[Sequence[CPE]] = asyncio.Queue(queue_size)
+        self.chunk_size = chunk_size
         self.console = console
         self.event = asyncio.Event()
         self.verbose = verbose
@@ -254,6 +273,7 @@ async def download(
                     request_results=request_results,
                     last_modified_start_date=last_modified_start_date,
                     last_modified_end_date=last_modified_end_date,
+                    results_per_page=self.chunk_size,
                 )
 
         result_count = len(results)  # type: ignore
@@ -305,6 +325,9 @@ async def download(console: Console, error_console: Console) -> None:
     number: int | None = args.number
     nvd_api_key: str | None = args.nvd_api_key or os.environ.get("NVD_API_KEY")
 
+    chunk_size: int | None = args.chunk_size
+    queue_size: int = args.queue_size
+
     if since_from_file:
         if since_from_file.exists():
             since = datetime.fromisoformat(
@@ -366,7 +389,9 @@ async def download(console: Console, error_console: Console) -> None:
     if verbose:
         console.log(f"Using PostgreSQL database {cpe_database_name}")
 
-    cli = CPECli(console, verbose=verbose)
+    cli = CPECli(
+        console, verbose=verbose, chunk_size=chunk_size, queue_size=queue_size
+    )
 
     with Progress(console=console) as progress:
         async with (

diff --git a/greenbone/scap/cve/cli/download.py b/greenbone/scap/cve/cli/download.py
@@ -137,6 +137,22 @@ def parse_args(args: Sequence[str] | None = None) -> Namespace:
         help="Use a NVD API key for downloading the CVEs. Using an API key "
         "allows for downloading with extended rate limits.",
     )
+    parser.add_argument(
+        "--chunk-size",
+        help="Number of CVEs to download and process in one request. A lower "
+        "number allows for more frequent updates and feedback.",
+        type=int,
+        metavar="N",
+    )
+    parser.add_argument(
+        "--queue-size",
+        help="Size of the download queue. It sets the maximum number of CVEs "
+        "kept in the memory. The maximum number of CVEs is chunk size * queue "
+        "size. Default: %(default)s.",
+        type=int,
+        metavar="N",
+        default=DEFAULT_QUEUE_SIZE,
+    )
     return parser.parse_args(args)
 
 
@@ -145,10 +161,12 @@ def __init__(
         self,
         console: Console,
         *,
-        queue: asyncio.Queue[Sequence[CVE]] | None = None,
         verbose: int = 0,
+        chunk_size: int | None = None,
+        queue_size: int = DEFAULT_QUEUE_SIZE,
     ) -> None:
-        self.queue = queue or asyncio.Queue(DEFAULT_QUEUE_SIZE)
+        self.queue: asyncio.Queue[Sequence[CVE]] = asyncio.Queue(queue_size)
+        self.chunk_size = chunk_size
         self.console = console
         self.event = asyncio.Event()
         self.verbose = verbose
@@ -260,6 +278,7 @@ async def download(
                     request_results=request_results,
                     last_modified_start_date=last_modified_start_date,
                     last_modified_end_date=last_modified_end_date,
+                    results_per_page=self.chunk_size,
                 )
 
         result_count = len(results)  # type: ignore
@@ -311,6 +330,9 @@ async def download(console: Console, error_console: Console):
     nvd_api_key: str | None = args.nvd_api_key or os.environ.get("NVD_API_KEY")
     updated_cves_file: Path | None = args.store_updated_cves
 
+    chunk_size: int | None = args.chunk_size
+    queue_size: int = args.queue_size
+
     until = now() if run_time_file else None
 
     if since_from_file:
@@ -374,7 +396,9 @@ async def download(console: Console, error_console: Console):
     if verbose:
         console.log(f"Using PostgreSQL database {cve_database_name} for CVEs")
 
-    cli = CVECli(console, verbose=verbose)
+    cli = CVECli(
+        console, verbose=verbose, chunk_size=chunk_size, queue_size=queue_size
+    )
 
     with Progress(console=console) as progress:
         async with (

diff --git a/tests/cpe/cli/test_download.py b/tests/cpe/cli/test_download.py
@@ -10,7 +10,7 @@
 
 from pontos.testing import temp_directory
 
-from greenbone.scap.cpe.cli.download import parse_args
+from greenbone.scap.cpe.cli.download import DEFAULT_QUEUE_SIZE, parse_args
 
 
 class ParseArgsTestCase(unittest.TestCase):
@@ -31,6 +31,8 @@ def test_defaults(self):
         self.assertIsNone(args.nvd_api_key)
         self.assertIsNone(args.since)
         self.assertIsNone(args.since_from_file)
+        self.assertIsNone(args.chunk_size)
+        self.assertEqual(args.queue_size, DEFAULT_QUEUE_SIZE)
 
     def test_database(self):
         args = parse_args(
@@ -151,3 +153,19 @@ def test_since_conflicts_since_from_file(self):
                     "2024-01-01T15:24:17.000000+00:00",
                 ]
             )
+
+    def test_chunk_size(self):
+        args = parse_args(["--chunk-size", "42"])
+
+        self.assertEqual(args.chunk_size, 42)
+
+        with self.assertRaises(SystemExit), redirect_stderr(StringIO()):
+            parse_args(["--chunk-size", "foo"])
+
+    def test_queue_size(self):
+        args = parse_args(["--queue-size", "42"])
+
+        self.assertEqual(args.queue_size, 42)
+
+        with self.assertRaises(SystemExit), redirect_stderr(StringIO()):
+            parse_args(["--queue-size", "foo"])
diff --git a/tests/cve/cli/test_download.py b/tests/cve/cli/test_download.py
@@ -10,7 +10,7 @@
 
 from pontos.testing import temp_directory
 
-from greenbone.scap.cve.cli.download import parse_args
+from greenbone.scap.cve.cli.download import DEFAULT_QUEUE_SIZE, parse_args
 
 
 class ParseArgsTestCase(unittest.TestCase):
@@ -32,6 +32,8 @@ def test_defaults(self):
         self.assertIsNone(args.since_from_file)
         self.assertIsNone(args.store_runtime)
         self.assertIsNone(args.store_updated_cves)
+        self.assertIsNone(args.chunk_size)
+        self.assertEqual(args.queue_size, DEFAULT_QUEUE_SIZE)
 
     def test_cve_database(self):
         args = parse_args(
@@ -152,3 +154,19 @@ def test_since_conflicts_since_from_file(self):
                     "2024-01-01T15:24:17.000000+00:00",
                 ]
             )
+
+    def test_chunk_size(self):
+        args = parse_args(["--chunk-size", "42"])
+
+        self.assertEqual(args.chunk_size, 42)
+
+        with self.assertRaises(SystemExit), redirect_stderr(StringIO()):
+            parse_args(["--chunk-size", "foo"])
+
+    def test_queue_size(self):
+        args = parse_args(["--queue-size", "42"])
+
+        self.assertEqual(args.queue_size, 42)
+
+        with self.assertRaises(SystemExit), redirect_stderr(StringIO()):
+            parse_args(["--queue-size", "foo"])