Add query count tests and remove N+1 queries in page rendering and search

apps/core/models/content.py

@@ -2,14 +2,15 @@
 
 from bs4 import BeautifulSoup
 from django.http import HttpResponse
-from django.template import Context, Template
 from django.utils.functional import cached_property
 from django.utils.html import format_html
 from django.utils.text import slugify
 from wagtail.admin.panels import FieldPanel
 from wagtail.api import APIField
+from wagtail.blocks import StructValue
 from wagtail.fields import StreamField
 from wagtail.models import Page
+from wagtail.rich_text import RichText
 from wagtail.search import index
 from wagtail_ai.panels import AITitleFieldPanel
 
@@ -19,10 +20,22 @@
 from ..blocks import CONTENT_BLOCKS
 
 
+def get_rich_text_sources(body):
+    for block in body:
+        if isinstance(block.value, StructValue):
+            values = block.value.values()
+        else:
+            values = [block.value]
+        for value in values:
+            if isinstance(value, RichText):
+                yield value.source
+
+
 def create_table_of_contents(body):
-    template = Template("{% load wagtailcore_tags %}{% include_block body %}")
-    content = template.render(Context({"body": body}))
-    soup = BeautifulSoup(content, "html.parser")
+    # Headings and their ids are stored in the rich text HTML, so read them
+    # from there rather than rendering the body a second time, which would
+    # repeat the database lookups for every link in it.
+    soup = BeautifulSoup("".join(get_rich_text_sources(body)), "html.parser")
     headings = soup.select("h2,h3")
     toc = ""
     if headings:

apps/core/tests/test_content_page.py

@@ -66,6 +66,45 @@ def test_escapes_existing_id(self):
             '<ul><li><a href="#section&quot; onclick=&quot;alert(1)">Title</a></li></ul>',
         )
 
+    def test_create_table_of_contents_annotated_text(self):
+        self.content_page.body = json.dumps(
+            [
+                {
+                    "type": "text_annotated",
+                    "value": {
+                        "content": '<h2 id="new">New</h2><h3>Details</h3>',
+                        "version": "7.4",
+                        "change_type": "added",
+                    },
+                }
+            ]
+        )
+
+        self.assertEqual(
+            self.content_page.table_of_contents,
+            '<ul><li><a href="#new">New</a></li><li><a href="#details">Details</a></li></ul>',
+        )
+
+    def test_create_table_of_contents_makes_no_queries(self):
+        self.content_page.body = json.dumps(
+            [
+                {
+                    "type": "text",
+                    "value": (
+                        f'<h2 id="links">Links</h2><p><a linktype="page" '
+                        f'id="{self.content_page.pk}">A page</a></p>'
+                    ),
+                }
+            ]
+        )
+
+        with self.assertNumQueries(0):
+            table_of_contents = self.content_page.table_of_contents
+
+        self.assertEqual(
+            table_of_contents, '<ul><li><a href="#links">Links</a></li></ul>'
+        )
+
     def test_table_of_contents_is_cached_on_page_instance(self):
         with patch(
             "apps.core.models.content.create_table_of_contents",

apps/core/tests/test_page_response.py

@@ -40,5 +40,5 @@ def test_content_page_template_used(self):
             self.content_page_response, template_name="core/content_page.html", count=1
         )
         self.assertTemplateUsed(
-            self.content_page_response, template_name="core/blocks/text.html", count=4
+            self.content_page_response, template_name="core/blocks/text.html", count=2
         )

apps/core/tests/test_query_counts.py

@@ -0,0 +1,95 @@
+import json
+import uuid
+
+from django.core.cache import cache
+from django.db import connection
+from django.test import TestCase, override_settings
+from django.test.utils import CaptureQueriesContext
+
+from apps.core.factories import ContentPageFactory, HomePageFactory
+
+# Production uses Redis, so cache reads and writes aren't database queries
+# there. Use an in-memory cache so the counts only reflect page rendering.
+LOCMEM_CACHES = {
+    "default": {"BACKEND": "django.core.cache.backends.locmem.LocMemCache"},
+}
+
+
+@override_settings(CACHES=LOCMEM_CACHES)
+class TestPageQueryCounts(TestCase):
+    def setUp(self):
+        cache.clear()
+        self.home = HomePageFactory()
+        self.page = ContentPageFactory(parent=self.home)
+
+    def get_query_count(self, url):
+        # Start from an empty cache so the cached header and language
+        # selector fragments are rendered too.
+        cache.clear()
+        with CaptureQueriesContext(connection) as ctx:
+            response = self.client.get(url)
+        self.assertEqual(response.status_code, 200)
+        return len(ctx.captured_queries)
+
+    def set_home_sections(self, pages):
+        self.home.sections = json.dumps(
+            [
+                {
+                    "type": "section_grid",
+                    "value": [
+                        {
+                            "type": "item",
+                            "id": str(uuid.uuid4()),
+                            "value": {
+                                "section": "tutorial",
+                                "title": page.title,
+                                "text": "Section text",
+                                "page": page.pk,
+                            },
+                        }
+                        for page in pages
+                    ],
+                }
+            ]
+        )
+        self.home.save_revision().publish()
+
+    def test_content_page(self):
+        # Warm the cached fragments, as for most requests in production.
+        self.client.get(self.page.url)
+        with self.assertNumQueries(17):
+            self.client.get(self.page.url)
+
+    def test_home_page(self):
+        self.set_home_sections([self.page])
+        self.client.get(self.home.url)
+        with self.assertNumQueries(12):
+            self.client.get(self.home.url)
+
+    def test_child_pages_do_not_add_queries(self):
+        ContentPageFactory(parent=self.page)
+        expected = self.get_query_count(self.page.url)
+
+        for _ in range(4):
+            ContentPageFactory(parent=self.page)
+
+        self.assertEqual(self.get_query_count(self.page.url), expected)
+
+    def test_menu_pages_do_not_add_queries(self):
+        ContentPageFactory(parent=self.home)
+        expected = self.get_query_count(self.page.url)
+
+        for _ in range(4):
+            section = ContentPageFactory(parent=self.home)
+            ContentPageFactory(parent=section)
+
+        self.assertEqual(self.get_query_count(self.page.url), expected)
+
+    def test_home_page_sections_do_not_add_queries(self):
+        pages = [ContentPageFactory(parent=self.home) for _ in range(5)]
+        self.set_home_sections(pages[:1])
+        expected = self.get_query_count(self.home.url)
+
+        self.set_home_sections(pages)
+
+        self.assertEqual(self.get_query_count(self.home.url), expected)

apps/core/tests/test_text_block_annotated.py

@@ -65,7 +65,7 @@ def test_removed_change_type_uses_base_badge_class(self):
     def test_uses_annotated_template(self):
         page = self._page_with_annotated_block("<p>Content</p>", change_type="added")
         response = self.client.get(page.url)
-        self.assertTemplateUsed(response, "core/blocks/text_annotated.html", count=2)
+        self.assertTemplateUsed(response, "core/blocks/text_annotated.html", count=1)
 
     def test_version_validation_accepts_x_y_format(self):
         version_block = AnnotatedTextBlock().child_blocks["version"]

apps/search/tests/test_views.py

@@ -0,0 +1,64 @@
+from django.db import connection
+from django.test import TestCase
+from django.test.utils import CaptureQueriesContext
+
+from apps.core.factories import ContentPageFactory, HomePageFactory
+
+SEARCH_JSON_URL = "/en/search_json/"
+
+
+class TestSearchJson(TestCase):
+    def setUp(self):
+        self.home = HomePageFactory()
+        self.section = ContentPageFactory(parent=self.home, title="Section")
+
+    def search(self, query):
+        response = self.client.get(SEARCH_JSON_URL, {"query": query})
+        self.assertEqual(response.status_code, 200)
+        return response.json()
+
+    def get_query_count(self, query):
+        with CaptureQueriesContext(connection) as ctx:
+            self.search(query)
+        return len(ctx.captured_queries)
+
+    def test_no_query(self):
+        response = self.client.get(SEARCH_JSON_URL)
+
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(response.json(), [])
+
+    def test_parent_section(self):
+        page = ContentPageFactory(parent=self.section, title="Zebra page")
+        ContentPageFactory(parent=page, title="Zebra subpage")
+        ContentPageFactory(parent=self.home, title="Zebra section")
+
+        results = self.search("zebra")
+
+        self.assertEqual(
+            sorted((result["title"], result["parent_section"]) for result in results),
+            [
+                ("Zebra page", "Section"),
+                ("Zebra section", "Home"),
+                ("Zebra subpage", "Section"),
+            ],
+        )
+
+    def test_full_url(self):
+        page = ContentPageFactory(parent=self.section, title="Zebra page")
+
+        results = self.search("zebra")
+
+        self.assertEqual([result["full_url"] for result in results], [page.full_url])
+
+    def test_results_do_not_add_queries(self):
+        ContentPageFactory(parent=self.section, title="Zebra 0")
+        self.search("zebra")
+        expected = self.get_query_count("zebra")
+
+        other_section = ContentPageFactory(parent=self.home, title="Other")
+        for i in range(1, 5):
+            ContentPageFactory(parent=other_section, title=f"Zebra {i}")
+
+        self.assertEqual(len(self.search("zebra")), 5)
+        self.assertEqual(self.get_query_count("zebra"), expected)

apps/search/views.py

@@ -47,17 +47,36 @@ def search(request):
     )
 
 
+# Sections are the children of each locale's homepage, under the tree root.
+SECTION_DEPTH = 3
+
+
+def get_section_path(page):
+    return page.path[: Page.steplen * SECTION_DEPTH]
+
+
+def get_section_titles(pages):
+    """Titles of the sections the pages are in, keyed by section path."""
+    paths = {get_section_path(page) for page in pages if page.depth > SECTION_DEPTH}
+    return dict(Page.objects.filter(path__in=paths).values_list("path", "title"))
+
+
 class PageSerializer(serializers.ModelSerializer):
+    full_url = serializers.SerializerMethodField("get_full_url")
     parent_section = serializers.SerializerMethodField("get_parent_section")
 
     class Meta:
         model = Page
         fields = ["id", "title", "search_description", "full_url", "parent_section"]
 
+    def get_full_url(self, page):
+        # With the request, Wagtail looks up the site root paths once for all
+        # results, rather than once per result.
+        return page.get_full_url(self.context["request"])
+
     def get_parent_section(self, page):
-        ancestors = page.get_ancestors()
-        if len(ancestors) >= 3:
-            return ancestors[2].title
+        if page.depth > SECTION_DEPTH:
+            return self.context["section_titles"][get_section_path(page)]
         else:
             return _("Home")
 
@@ -77,5 +96,14 @@ def search_json(request):
         query.add_hit()
     else:
         search_results = Page.objects.none()
-    serializer = PageSerializer(search_results, many=True)
+    # Look up all the parent sections at once, rather than once per result.
+    search_results = list(search_results)
+    serializer = PageSerializer(
+        search_results,
+        many=True,
+        context={
+            "request": request,
+            "section_titles": get_section_titles(search_results),
+        },
+    )
     return Response(serializer.data)