Fetch and apply annotations if using specific()

CHANGELOG.txt

@@ -44,6 +44,7 @@ Changelog
  * Fix: Accept unicode characters in slugs on the "copy page" form (François Poulain)
  * Fix: Remove top padding when `FieldRowPanel` is used inside a `MultiFieldPanel` (Jérôme Lebleu)
  * Fix: Add Wagtail User Bar back to page previews and ensure moderation actions are available (Coen van der Kamp)
+ * Fix: Resolve issue where queryset annotations were lost (e.g. `.annotate_score()`) when using specific models in page query (Dan Bentley)
 
 
 2.9 (04.05.2020)

CONTRIBUTORS.rst

@@ -462,6 +462,7 @@ Contributors
 * Caitlin White
 * Brylie Christopher Oxley
 * Lacey Williams Henschel
+* Dan Bentley
 
 Translators
 ===========

docs/releases/2.10.rst

@@ -62,6 +62,7 @@ Bug fixes
  * Support IPv6 domain (Alex Gleason, Coen van der Kamp)
  * Remove top padding when ``FieldRowPanel`` is used inside a ``MultiFieldPanel`` (Jérôme Lebleu)
  * Add Wagtail User Bar back to page previews and ensure moderation actions are available (Coen van der Kamp)
+ * Fix issue where queryset annotations were lost (e.g. ``.annotate_score()``) when using specific models in page query (Dan Bentley)
 
 
 Upgrade considerations

wagtail/core/query.py

@@ -373,7 +373,16 @@ def specific_iterator(qs, defer=False):
 
     This should be called from ``PageQuerySet.specific``
     """
-    pks_and_types = qs.values_list('pk', 'content_type')
+    annotation_aliases = qs.query.annotations.keys()
+    values = qs.values('pk', 'content_type', *annotation_aliases)
+
+    annotations_by_pk = defaultdict(list)
+    if annotation_aliases:
+        # Extract annotation results keyed by pk so we can reapply to fetched pages.
+        for data in values:
+            annotations_by_pk[data['pk']] = {k: v for k, v in data.items() if k in annotation_aliases}
+
+    pks_and_types = [[v['pk'], v['content_type']] for v in values]
     pks_by_type = defaultdict(list)
     for pk, content_type in pks_and_types:
         pks_by_type[content_type].append(pk)
@@ -390,6 +399,12 @@ def specific_iterator(qs, defer=False):
         model = content_types[content_type].model_class() or qs.model
         pages = model.objects.filter(pk__in=pks)
 
+        if annotation_aliases:
+            # Reapply annotations to pages.
+            for page in pages:
+                for annotation, value in annotations_by_pk.get(page.pk, {}).items():
+                    setattr(page, annotation, value)
+
         if defer:
             # Defer all specific fields
             from wagtail.core.models import Page

wagtail/core/tests/test_page_queryset.py

@@ -1,5 +1,5 @@
 from django.contrib.contenttypes.models import ContentType
-from django.db.models import Q
+from django.db.models import Count, Q
 from django.test import TestCase
 
 from wagtail.core.models import Page, PageViewRestriction, Site
@@ -607,6 +607,38 @@ class TestSpecificQuery(TestCase):
             Page.objects.get(url_path='/home/events/').specific,
             Page.objects.get(url_path='/home/about-us/').specific])
 
+    def test_specific_query_with_annotations_performs_no_additional_queries(self):
+
+        with self.assertNumQueries(5):
+            pages = list(Page.objects.live().specific())
+
+            self.assertEqual(len(pages), 7)
+
+        with self.assertNumQueries(5):
+            pages = list(Page.objects.live().specific().annotate(count=Count('pk')))
+
+            self.assertEqual(len(pages), 7)
+
+    def test_specific_query_with_annotation(self):
+        # Ensure annotations are reapplied to specific() page queries
+
+        pages = Page.objects.live()
+        pages.first().save_revision()
+        pages.last().save_revision()
+
+        results = Page.objects.live().specific().annotate(revision_count=Count('revisions'))
+
+        self.assertEqual(results.first().revision_count, 1)
+        self.assertEqual(results.last().revision_count, 1)
+
+    def test_specific_query_with_search_and_annotation(self):
+        # Ensure annotations are reapplied to specific() page queries
+
+        results = Page.objects.live().specific().search(MATCH_ALL).annotate_score('_score')
+
+        for result in results:
+            self.assertTrue(hasattr(result, '_score'))
+
     def test_specific_query_with_search(self):
         # 1276 - The database search backend didn't return results with the
         # specific type when searching a specific queryset.