Merge pull request #186 from sbs2001/cli_batchsize

Add CLI option to specify batch_size for the ImportRunner

Author: sbs2001

vulnerabilities/management/commands/import.py

@@ -51,11 +51,17 @@ def add_arguments(self, parser):
         parser.add_argument('sources', nargs='*',
                             help='Data sources from which to import')
 
+        parser.add_argument(
+            '--batch_size', help='The batch size to be used for bulk inserting data')
+
     def handle(self, *args, **options):
         if options['list']:
             self.list_sources()
             return
 
+        if options['batch_size']:
+            self.batch_size = options['batch_size']
+
         if options['all']:
             self._import_data(Importer.objects.all(), options['cutoff_date'])
             return
@@ -94,8 +100,7 @@ def import_data(self, names, cutoff_date):
     def _import_data(self, importers, cutoff_date):
         for importer in importers:
             self.stdout.write(f'Importing data from {importer.name}')
-            # TODO add cmdline param for batch_size
-            ImportRunner(importer, 10).run(cutoff_date=cutoff_date)
-
+            batch_size = getattr(self, 'batch_size', 10)
+            ImportRunner(importer, batch_size).run(cutoff_date=cutoff_date)
             self.stdout.write(
                 self.style.SUCCESS(f'Successfully imported data from {importer.name}'))