From efe57694e33a8d898739630a024f3bbd8b427f3b Mon Sep 17 00:00:00 2001 From: Shivam Sandbhor Date: Tue, 26 May 2020 18:25:12 +0530 Subject: [PATCH] Add CLI option to specify batch_size for the ImportRunner Signed-off-by: Shivam Sandbhor --- vulnerabilities/management/commands/import.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/vulnerabilities/management/commands/import.py b/vulnerabilities/management/commands/import.py index 6489f2cbe..6a8c5eacd 100644 --- a/vulnerabilities/management/commands/import.py +++ b/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}'))