Remap primary keys of child objects in revisions when copying page

wagtail/wagtailcore/models.py

@@ -2,6 +2,7 @@ from __future__ import unicode_literals
 
 import logging
 import json
+from collections import defaultdict
 
 from modelcluster.models import ClusterableModel, get_all_child_relations
 
@@ -810,18 +811,27 @@ class Page(six.with_metaclass(PageBase, MP_Node, ClusterableModel, index.Indexed
         else:
             page_copy = self.add_sibling(instance=page_copy)
 
+        # A dict that maps child objects to their new ids
+        # Used to remap child object ids in revisions
+        child_object_id_map = defaultdict(dict)
+
         # Copy child objects
         specific_self = self.specific
         for child_relation in get_all_child_relations(specific_self):
+            accessor_name = child_relation.get_accessor_name()
             parental_key_name = child_relation.field.attname
-            child_objects = getattr(specific_self, child_relation.get_accessor_name(), None)
+            child_objects = getattr(specific_self, accessor_name, None)
 
             if child_objects:
                 for child_object in child_objects.all():
+                    old_pk = child_object.pk
                     child_object.pk = None
                     setattr(child_object, parental_key_name, page_copy.id)
                     child_object.save()
 
+                    # Add mapping to new primary key (so we can apply this change to revisions)
+                    child_object_id_map[accessor_name][old_pk] = child_object.pk
+
         # Copy revisions
         if copy_revisions:
             for revision in self.revisions.all():
@@ -835,8 +845,9 @@ class Page(six.with_metaclass(PageBase, MP_Node, ClusterableModel, index.Indexed
                 revision_content['pk'] = page_copy.pk
 
                 for child_relation in get_all_child_relations(specific_self):
+                    accessor_name = child_relation.get_accessor_name()
                     try:
-                        child_objects = revision_content[child_relation.get_accessor_name()]
+                        child_objects = revision_content[accessor_name]
                     except KeyError:
                         # KeyErrors are possible if the revision was created
                         # before this child relation was added to the database
@@ -845,6 +856,11 @@ class Page(six.with_metaclass(PageBase, MP_Node, ClusterableModel, index.Indexed
                     for child_object in child_objects:
                         child_object[child_relation.field.name] = page_copy.pk
 
+                        # Remap primary key to copied versions
+                        # If the primary key is not recognised (eg, the child object has been deleted from the database)
+                        # set the primary key to None
+                        child_object['pk'] = child_object_id_map[accessor_name].get(child_object['pk'], None)
+
                 revision.content_json = json.dumps(revision_content)
 
                 # Save

wagtail/wagtailcore/tests/test_page_model.py

@@ -1,4 +1,3 @@
-import unittest
 import datetime
 import json
 
@@ -424,7 +423,6 @@ class TestCopyPage(TestCase):
         self.assertEqual(new_christmas_event.advert_placements.count(), 1, "Child objects defined on the superclass weren't copied")
         self.assertEqual(christmas_event.advert_placements.count(), 1, "Child objects defined on the superclass were removed from the original page")
 
-    @unittest.expectedFailure
     def test_copy_page_copies_revisions(self):
         christmas_event = EventPage.objects.get(url_path='/home/events/christmas/')
         christmas_event.save_revision()