Merge pull request #99 from macro1/populate-history-command

Populate history command

Author: treyhunner

setup.py

@@ -13,7 +13,7 @@
     author_email='[email protected]',
     maintainer='Trey Hunner',
     url='https://github.com/treyhunner/django-simple-history',
-    packages=["simple_history"],
+    packages=["simple_history", "simple_history.management", "simple_history.management.commands"],
     classifiers=[
         "Development Status :: 5 - Production/Stable",
         "Framework :: Django",

simple_history/management/__init__.py

@@ -0,0 +1,38 @@
+try:
+    from django.utils.timezone import now
+except ImportError:     # pragma: no cover
+    from datetime import datetime
+    now = datetime.now
+from django.db import transaction
+
+
+class NotHistorical(TypeError):
+    """No related history model found."""
+
+
+def get_history_model_for_model(model):
+    """Find the history model for a given app model."""
+    try:
+        manager_name = model._meta.simple_history_manager_attribute
+    except AttributeError:
+        raise NotHistorical("Cannot find a historical model for "
+                            "{model}.".format(model=model))
+    return getattr(model, manager_name).model
+
+
+def bulk_history_create(model, history_model):
+    """Save a copy of all instances to the historical model."""
+    historical_instances = [
+        history_model(
+            history_date=getattr(instance, '_history_date', now()),
+            history_user=getattr(instance, '_history_user', None),
+            **dict((field.attname, getattr(instance, field.attname))
+                   for field in instance._meta.fields)
+        ) for instance in model.objects.all()]
+    try:
+        history_model.objects.bulk_create(historical_instances)
+    except AttributeError:  # pragma: no cover
+        # bulk_create was added in Django 1.4, handle legacy versions
+        with transaction.commit_on_success():
+            for instance in historical_instances:
+                instance.save()

simple_history/management/commands/__init__.py

@@ -0,0 +1,101 @@
+from optparse import make_option
+from django.core.management.base import BaseCommand, CommandError
+from django.db.models.loading import get_model
+
+from ... import models
+from . import _populate_utils as utils
+
+
+class Command(BaseCommand):
+    args = "<app.model app.model ...>"
+    help = ("Populates the corresponding HistoricalRecords field with "
+            "the current state of all instances in a model")
+
+    COMMAND_HINT = "Please specify a model or use the --auto option"
+    MODEL_NOT_FOUND = "Unable to find model"
+    MODEL_NOT_HISTORICAL = "No history model found"
+    NO_REGISTERED_MODELS = "No registered models were found\n"
+    START_SAVING_FOR_MODEL = "Saving historical records for {model}\n"
+    DONE_SAVING_FOR_MODEL = "Finished saving historical records for {model}\n"
+    EXISTING_HISTORY_FOUND = "Existing history found, skipping model"
+    INVALID_MODEL_ARG = "An invalid model was specified"
+
+    option_list = BaseCommand.option_list + (
+        make_option(
+            '--auto',
+            action='store_true',
+            dest='auto',
+            default=False,
+            help="Automatically search for models with the "
+                 "HistoricalRecords field type",
+        ),
+    )
+
+    def handle(self, *args, **options):
+        to_process = set()
+
+        if args:
+            for model_pair in self._handle_model_list(*args):
+                to_process.add(model_pair)
+
+        elif options['auto']:
+            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:
+                    continue
+                to_process.add((model, history_model))
+            if not to_process:
+                self.stdout.write(self.NO_REGISTERED_MODELS)
+
+        else:
+            self.stdout.write(self.COMMAND_HINT)
+
+        self._process(to_process)
+
+    def _handle_model_list(self, *args):
+        failing = False
+        for natural_key in args:
+            try:
+                model, history = self._model_from_natural_key(natural_key)
+            except ValueError as e:
+                failing = True
+                self.stderr.write("{error}\n".format(error=e))
+            else:
+                if not failing:
+                    yield (model, history)
+        if failing:
+            raise CommandError(self.INVALID_MODEL_ARG)
+
+    def _model_from_natural_key(self, natural_key):
+        try:
+            app_label, model = natural_key.split(".", 1)
+        except ValueError:
+            model = None
+        else:
+            try:
+                model = get_model(app_label, model)
+                # Django 1.7 raises a LookupError
+            except LookupError:     # pragma: no cover
+                model = None
+        if not model:
+            raise ValueError(self.MODEL_NOT_FOUND +
+                             " < {model} >\n".format(model=natural_key))
+        try:
+            history_model = utils.get_history_model_for_model(model)
+        except utils.NotHistorical:
+            raise ValueError(self.MODEL_NOT_HISTORICAL +
+                             " < {model} >\n".format(model=natural_key))
+        return model, history_model
+
+    def _process(self, to_process):
+        for model, history_model in to_process:
+            if history_model.objects.count():
+                self.stderr.write("{msg} {model}\n".format(
+                    msg=self.EXISTING_HISTORY_FOUND,
+                    model=model,
+                ))
+                continue
+            self.stdout.write(self.START_SAVING_FOR_MODEL.format(model=model))
+            utils.bulk_history_create(model, history_model)
+            self.stdout.write(self.DONE_SAVING_FOR_MODEL.format(model=model))

simple_history/management/commands/_populate_utils.py

@@ -0,0 +1,2 @@
+from .test_models import *
+from .test_commands import *

simple_history/management/commands/populate_history.py

@@ -0,0 +1,103 @@
+from contextlib import contextmanager
+from six.moves import cStringIO as StringIO
+from datetime import datetime
+from django.test import TestCase
+from django.core import management
+from simple_history import models as sh_models
+from simple_history.management.commands import populate_history
+
+from .. import models
+
+
+@contextmanager
+def replace_registry(new_value=None):
+    hidden_registry = sh_models.registered_models
+    sh_models.registered_models = new_value or {}
+    try:
+        yield
+    except Exception:
+        raise
+    finally:
+        sh_models.registered_models = hidden_registry
+
+
+class TestPopulateHistory(TestCase):
+    command_name = 'populate_history'
+    command_error = (management.CommandError, SystemExit)
+
+    def test_no_args(self):
+        out = StringIO()
+        management.call_command(self.command_name,
+                                stdout=out, stderr=StringIO())
+        self.assertIn(populate_history.Command.COMMAND_HINT, out.getvalue())
+
+    def test_bad_args(self):
+        test_data = (
+            (populate_history.Command.MODEL_NOT_HISTORICAL, ("tests.place",)),
+            (populate_history.Command.MODEL_NOT_FOUND, ("invalid.model",)),
+            (populate_history.Command.MODEL_NOT_FOUND, ("bad_key",)),
+        )
+        for msg, args in test_data:
+            out = StringIO()
+            self.assertRaises(self.command_error, management.call_command,
+                              self.command_name, *args,
+                              stdout=StringIO(), stderr=out)
+            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()
+        management.call_command(self.command_name, auto=True,
+                                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())
+        models.Poll.history.all().delete()
+        models.Book.objects.create(isbn="9780007117116")
+        models.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)
+
+    def test_failing_wont_save(self):
+        models.Poll.objects.create(question="Will this populate?",
+                                   pub_date=datetime.now())
+        models.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)
+
+    def test_multi_table(self):
+        data = {'rating': 5, 'name': "Tea 'N More"}
+        models.Restaurant.objects.create(**data)
+        models.Restaurant.updates.all().delete()
+        management.call_command(self.command_name, 'tests.restaurant',
+                                stdout=StringIO(), stderr=StringIO())
+        update_record = models.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()
+        management.call_command(self.command_name, 'tests.restaurant',
+                                stdout=StringIO(), stderr=out)
+        self.assertEqual(models.Restaurant.updates.count(), pre_call_count)
+        self.assertIn(populate_history.Command.EXISTING_HISTORY_FOUND,
+                      out.getvalue())
+
+    def test_no_historical(self):
+        out = StringIO()
+        with replace_registry():
+            management.call_command(self.command_name, auto=True,
+                                    stdout=out)
+        self.assertIn(populate_history.Command.NO_REGISTERED_MODELS,
+                      out.getvalue())

simple_history/tests/tests/__init__.py

@@ -17,13 +17,13 @@
 
 from simple_history.models import HistoricalRecords
 from simple_history import register
-from .models import (
+from ..models import (
     AdminProfile, Bookcase, MultiOneToOne, Poll, Choice, Restaurant, Person,
     FileModel, Document, Book, HistoricalPoll, Library, State, AbstractBase,
     ConcreteAttr, ConcreteUtil, SelfFK, Temperature, WaterLevel,
     ExternalModel1, ExternalModel3, UnicodeVerboseName
 )
-from .external.models import ExternalModel2, ExternalModel4
+from ..external.models import ExternalModel2, ExternalModel4
 
 today = datetime(2021, 1, 1, 10, 0)
 tomorrow = today + timedelta(days=1)