Merge pull request #263 from treyhunner/mikeengland-master

Added batch size option to the management command for populating the history

Author: macro1

AUTHORS.rst

@@ -33,6 +33,7 @@ Authors
 - Ulysses Vilela
 - vnagendra
 - Lucas Wiman
+- Michael England
 
 Background
 ==========

CHANGES.rst

@@ -1,6 +1,10 @@
 Changes
 =======
 
+tip (unreleased)
+----------------
+- Added --batchsize option to the populate_history management command.
+
 1.8.2 (2017-01-19)
 ------------------
 - Add Polish locale.

docs/usage.rst

@@ -83,6 +83,9 @@ initial change for preexisting model instances:
 
     $ python manage.py populate_history --auto
 
+By default, history rows are inserted in batches of 200. This can be changed if needed for large tables
+by using the ``--batchsize`` option, for example ``--batchsize 500``.
+
 .. _admin_integration:
 
 Integration with Django Admin

simple_history/management/commands/_populate_utils.py

@@ -15,7 +15,7 @@ def get_history_model_for_model(model):
     return getattr(model, manager_name).model
 
 
-def bulk_history_create(model, history_model):
+def bulk_history_create(model, history_model, batch_size):
     """Save a copy of all instances to the historical model."""
     historical_instances = [
         history_model(
@@ -26,4 +26,4 @@ def bulk_history_create(model, history_model):
                 for field in instance._meta.fields
             }
         ) for instance in model.objects.all()]
-    history_model.objects.bulk_create(historical_instances)
+    history_model.objects.bulk_create(historical_instances, batch_size=batch_size)

simple_history/management/commands/populate_history.py

@@ -30,6 +30,7 @@ class Command(BaseCommand):
     if hasattr(BaseCommand, 'option_list'):  # Django < 1.8
         option_list = BaseCommand.option_list + (
             make_option('--auto', action='store_true', dest='auto', default=False),
+            make_option('--batchsize', action='store', dest='batchsize', default=200, type=int),
         )
 
     def add_arguments(self, parser):
@@ -43,6 +44,14 @@ def add_arguments(self, parser):
             help='Automatically search for models with the '
                  'HistoricalRecords field type',
         )
+        parser.add_argument(
+            '--batchsize',
+            action='store',
+            dest='batchsize',
+            default=200,
+            type=int,
+            help='Set a custom batch size when bulk inserting historical records.',
+        )
 
     def handle(self, *args, **options):
         to_process = set()
@@ -65,7 +74,7 @@ def handle(self, *args, **options):
         else:
             self.stdout.write(self.COMMAND_HINT)
 
-        self._process(to_process)
+        self._process(to_process, batch_size=options['batchsize'])
 
     def _handle_model_list(self, *args):
         failing = False
@@ -101,7 +110,7 @@ def _model_from_natural_key(self, natural_key):
                              " < {model} >\n".format(model=natural_key))
         return model, history_model
 
-    def _process(self, to_process):
+    def _process(self, to_process, batch_size):
         for model, history_model in to_process:
             if history_model.objects.count():
                 self.stderr.write("{msg} {model}\n".format(
@@ -110,5 +119,5 @@ def _process(self, to_process):
                 ))
                 continue
             self.stdout.write(self.START_SAVING_FOR_MODEL.format(model=model))
-            utils.bulk_history_create(model, history_model)
+            utils.bulk_history_create(model, history_model, batch_size)
             self.stdout.write(self.DONE_SAVING_FOR_MODEL.format(model=model))

simple_history/tests/tests/test_commands.py

@@ -54,6 +54,14 @@ def test_auto_populate(self):
                                 stdout=StringIO(), stderr=StringIO())
         self.assertEqual(models.Poll.history.all().count(), 1)
 
+    def test_populate_with_custom_batch_size(self):
+        models.Poll.objects.create(question="Will this populate?",
+                                   pub_date=datetime.now())
+        models.Poll.history.all().delete()
+        management.call_command(self.command_name, auto=True, batchsize=500,
+                                stdout=StringIO(), stderr=StringIO())
+        self.assertEqual(models.Poll.history.all().count(), 1)
+
     def test_specific_populate(self):
         models.Poll.objects.create(question="Will this populate?",
                                    pub_date=datetime.now())