Fix out of memory exception in populate_history management command (#408)

Ross Mechanic · web-flow · commit ad523eb4b825 · 2018-06-14T16:25:40.000-04:00
* fix out-of-memory exceptions while generating historical models with management command

* Refactored bulk_create_history to be on history manager

* Added tests for bulk_history_create

* Fixed linting errors

* Added tests to command

* Added to AUTHORS.rst and CHANGES.rst files

* Use _default_manager when querying models

* Updated gh number
diff --git a/AUTHORS.rst b/AUTHORS.rst
@@ -34,6 +34,7 @@ Authors
 - Klaas van Schelven
 - Kris Neuharth
 - Maciej "RooTer" Urbański
+- Mark Davidoff
 - Martin Bachwerk
 - Marty Alchin
 - Mauricio de Abreu Antunes
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -1,6 +1,10 @@
 Changes
 =======
 
+Unreleased
+----------
+- Fixed out-of-memory exception when running populate_history management command (gh-408)
+
 2.1.0 (2018-06-04)
 ------------------
 - Add ability to specify custom history_reason field (gh-379)
diff --git a/simple_history/exceptions.py b/simple_history/exceptions.py
@@ -6,3 +6,8 @@
 class MultipleRegistrationsError(Exception):
     """The model has been registered to have history tracking more than once"""
     pass
+
+
+class NotHistoricalModelError(TypeError):
+    """No related history model found."""
+    pass
diff --git a/simple_history/management/commands/_populate_utils.py b/simple_history/management/commands/_populate_utils.py
diff --git a/simple_history/management/commands/populate_history.py b/simple_history/management/commands/populate_history.py
@@ -1,8 +1,9 @@
+import django
 from django.apps import apps
 from django.core.management.base import BaseCommand, CommandError
 
-from . import _populate_utils as utils
-from ... import models
+from ... import models, utils
+from ...exceptions import NotHistoricalModelError
 
 get_model = apps.get_model
 
@@ -39,10 +40,12 @@ def add_arguments(self, parser):
             default=200,
             type=int,
             help='Set a custom batch size when bulk inserting historical '
-                 'records.',
+                 'records.'
         )
 
     def handle(self, *args, **options):
+        self.verbosity = options['verbosity']
+
         to_process = set()
         model_strings = options.get('models', []) or args
 
@@ -54,14 +57,16 @@ def handle(self, *args, **options):
             for model in models.registered_models.values():
                 try:    # avoid issues with mutli-table inheritance
                     history_model = utils.get_history_model_for_model(model)
-                except utils.NotHistorical:
+                except NotHistoricalModelError:
                     continue
                 to_process.add((model, history_model))
             if not to_process:
-                self.stdout.write(self.NO_REGISTERED_MODELS)
+                if self.verbosity >= 1:
+                    self.stdout.write(self.NO_REGISTERED_MODELS)
 
         else:
-            self.stdout.write(self.COMMAND_HINT)
+            if self.verbosity >= 1:
+                self.stdout.write(self.COMMAND_HINT)
 
         self._process(to_process, batch_size=options['batchsize'])
 
@@ -94,11 +99,55 @@ def _model_from_natural_key(self, natural_key):
                              " < {model} >\n".format(model=natural_key))
         try:
             history_model = utils.get_history_model_for_model(model)
-        except utils.NotHistorical:
+        except NotHistoricalModelError:
             raise ValueError(self.MODEL_NOT_HISTORICAL +
                              " < {model} >\n".format(model=natural_key))
         return model, history_model
 
+    def _bulk_history_create(self, model, batch_size):
+        """Save a copy of all instances to the historical model.
+
+        :param model: Model you want to bulk create
+        :param batch_size: number of models to create at once.
+        :return:
+        """
+
+        instances = []
+        history = utils.get_history_manager_for_model(model)
+        if self.verbosity >= 1:
+            self.stdout.write(
+                "Starting bulk creating history models for {} instances {}-{}"
+                .format(model, 0, batch_size)
+            )
+
+        iterator_kwargs = {'chunk_size': batch_size} \
+            if django.VERSION >= (2, 0, 0) else {}
+        for index, instance in enumerate(model._default_manager.iterator(
+            **iterator_kwargs
+        )):
+            # Can't Just pass batch_size to bulk_create as this can lead to
+            # Out of Memory Errors as we load too many models into memory after
+            # creating them. So we only keep batch_size worth of models in
+            # historical_instances and clear them after we hit batch_size
+            if index % batch_size == 0:
+
+                history.bulk_history_create(instances, batch_size=batch_size)
+
+                instances = []
+
+                if self.verbosity >= 1:
+                    self.stdout.write(
+                        "Finished bulk creating history models for {} "
+                        "instances {}-{}, starting next {}"
+                        .format(model, index - batch_size, index, batch_size)
+                    )
+
+            instances.append(instance)
+
+        # create any we didn't get in the last loop
+        if instances:
+            history.bulk_history_create(instances, batch_size=batch_size)
+
     def _process(self, to_process, batch_size):
         for model, history_model in to_process:
             if history_model.objects.count():
@@ -107,6 +156,12 @@ def _process(self, to_process, batch_size):
                     model=model,
                 ))
                 continue
-            self.stdout.write(self.START_SAVING_FOR_MODEL.format(model=model))
-            utils.bulk_history_create(model, history_model, batch_size)
-            self.stdout.write(self.DONE_SAVING_FOR_MODEL.format(model=model))
+            if self.verbosity >= 1:
+                self.stdout.write(self.START_SAVING_FOR_MODEL.format(
+                    model=model
+                ))
+            self._bulk_history_create(model, batch_size)
+            if self.verbosity >= 1:
+                self.stdout.write(self.DONE_SAVING_FOR_MODEL.format(
+                    model=model
+                ))
diff --git a/simple_history/manager.py b/simple_history/manager.py
@@ -1,6 +1,7 @@
 from __future__ import unicode_literals
 
 from django.db import models
+from django.utils.timezone import now
 
 
 class HistoryDescriptor(object):
@@ -88,3 +89,19 @@ def _as_of_set(self, date):
                               history_type='-').exists():
                 continue
             yield last_change.instance
+
+    def bulk_history_create(self, objs, batch_size=None):
+        """Bulk create the history for the objects specified by objs"""
+
+        historical_instances = [
+            self.model(
+                history_date=getattr(instance, '_history_date', now()),
+                history_user=getattr(instance, '_history_user', None),
+                **{
+                    field.attname: getattr(instance, field.attname)
+                    for field in instance._meta.fields
+                }
+            ) for instance in objs]
+
+        return self.model.objects.bulk_create(historical_instances,
+                                              batch_size=batch_size)
diff --git a/simple_history/tests/tests/test_commands.py b/simple_history/tests/tests/test_commands.py
@@ -1,14 +1,13 @@
 from contextlib import contextmanager
 from datetime import datetime
 
-from six.moves import cStringIO as StringIO
-from django.test import TestCase
 from django.core import management
+from django.test import TestCase
+from six.moves import cStringIO as StringIO
 
 from simple_history import models as sh_models
 from simple_history.management.commands import populate_history
-
-from .. import models
+from ..models import Book, Poll, Restaurant
 
 
 @contextmanager
@@ -47,60 +46,60 @@ def test_bad_args(self):
             self.assertIn(msg, out.getvalue())
 
     def test_auto_populate(self):
-        models.Poll.objects.create(question="Will this populate?",
-                                   pub_date=datetime.now())
-        models.Poll.history.all().delete()
+        Poll.objects.create(question="Will this populate?",
+                            pub_date=datetime.now())
+        Poll.history.all().delete()
         management.call_command(self.command_name, auto=True,
                                 stdout=StringIO(), stderr=StringIO())
-        self.assertEqual(models.Poll.history.all().count(), 1)
+        self.assertEqual(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()
+        Poll.objects.create(question="Will this populate?",
+                            pub_date=datetime.now())
+        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)
+        self.assertEqual(Poll.history.all().count(), 1)
 
     def test_specific_populate(self):
-        models.Poll.objects.create(question="Will this populate?",
-                                   pub_date=datetime.now())
-        models.Poll.history.all().delete()
-        models.Book.objects.create(isbn="9780007117116")
-        models.Book.history.all().delete()
+        Poll.objects.create(question="Will this populate?",
+                            pub_date=datetime.now())
+        Poll.history.all().delete()
+        Book.objects.create(isbn="9780007117116")
+        Book.history.all().delete()
         management.call_command(self.command_name, "tests.book",
                                 stdout=StringIO(), stderr=StringIO())
-        self.assertEqual(models.Book.history.all().count(), 1)
-        self.assertEqual(models.Poll.history.all().count(), 0)
+        self.assertEqual(Book.history.all().count(), 1)
+        self.assertEqual(Poll.history.all().count(), 0)
 
     def test_failing_wont_save(self):
-        models.Poll.objects.create(question="Will this populate?",
-                                   pub_date=datetime.now())
-        models.Poll.history.all().delete()
+        Poll.objects.create(question="Will this populate?",
+                            pub_date=datetime.now())
+        Poll.history.all().delete()
         self.assertRaises(self.command_error,
                           management.call_command, self.command_name,
                           "tests.poll", "tests.invalid_model",
                           stdout=StringIO(), stderr=StringIO())
-        self.assertEqual(models.Poll.history.all().count(), 0)
+        self.assertEqual(Poll.history.all().count(), 0)
 
     def test_multi_table(self):
         data = {'rating': 5, 'name': "Tea 'N More"}
-        models.Restaurant.objects.create(**data)
-        models.Restaurant.updates.all().delete()
+        Restaurant.objects.create(**data)
+        Restaurant.updates.all().delete()
         management.call_command(self.command_name, 'tests.restaurant',
                                 stdout=StringIO(), stderr=StringIO())
-        update_record = models.Restaurant.updates.all()[0]
+        update_record = Restaurant.updates.all()[0]
         for attr, value in data.items():
             self.assertEqual(getattr(update_record, attr), value)
 
     def test_existing_objects(self):
         data = {'rating': 5, 'name': "Tea 'N More"}
         out = StringIO()
-        models.Restaurant.objects.create(**data)
-        pre_call_count = models.Restaurant.updates.count()
+        Restaurant.objects.create(**data)
+        pre_call_count = Restaurant.updates.count()
         management.call_command(self.command_name, 'tests.restaurant',
                                 stdout=StringIO(), stderr=out)
-        self.assertEqual(models.Restaurant.updates.count(), pre_call_count)
+        self.assertEqual(Restaurant.updates.count(), pre_call_count)
         self.assertIn(populate_history.Command.EXISTING_HISTORY_FOUND,
                       out.getvalue())
 
@@ -111,3 +110,35 @@ def test_no_historical(self):
                                     stdout=out)
         self.assertIn(populate_history.Command.NO_REGISTERED_MODELS,
                       out.getvalue())
+
+    def test_batch_processing_with_batch_size_less_than_total(self):
+        data = [
+            Poll(id=1, question='Question 1', pub_date=datetime.now()),
+            Poll(id=2, question='Question 2', pub_date=datetime.now()),
+            Poll(id=3, question='Question 3', pub_date=datetime.now()),
+            Poll(id=4, question='Question 4', pub_date=datetime.now()),
+        ]
+        Poll.objects.bulk_create(data)
+
+        management.call_command(self.command_name, auto=True, batchsize=3,
+                                stdout=StringIO(), stderr=StringIO())
+
+        self.assertEqual(Poll.history.count(), 4)
+
+    def test_stdout_not_printed_when_verbosity_is_0(self):
+        out = StringIO()
+        Poll.objects.create(question='Question 1', pub_date=datetime.now())
+
+        management.call_command(self.command_name, auto=True, batchsize=3,
+                                stdout=out, stderr=StringIO(), verbosity=0)
+
+        self.assertEqual(out.getvalue(), '')
+
+    def test_stdout_printed_when_verbosity_is_not_specified(self):
+        out = StringIO()
+        Poll.objects.create(question='Question 1', pub_date=datetime.now())
+
+        management.call_command(self.command_name, auto=True, batchsize=3,
+                                stdout=out, stderr=StringIO())
+
+        self.assertNotEqual(out.getvalue(), '')
diff --git a/simple_history/tests/tests/test_manager.py b/simple_history/tests/tests/test_manager.py
diff --git a/simple_history/utils.py b/simple_history/utils.py
