dnscms: exclude non-live pages from search results, rank results across page types

dnscms/dnscms/wagtail_hooks.py

@@ -1,10 +1,7 @@
-from django.apps import apps as django_apps
 from django.templatetags.static import static
 from django.utils.html import format_html
 from grapple.registry import registry as grapple_registry
 from wagtail import hooks
-from wagtail.documents import get_document_model
-from wagtail.images import get_image_model
 from wagtail.models import Page
 from wagtail.search.backends import get_search_backend
 
@@ -15,33 +12,37 @@ def enable_additional_rich_text_features(features):
 
 
 @hooks.register("register_schema_query")
-def filter_search_to_live_pages(query_mixins):
+def override_search_resolver(query_mixins):
     """
-    Grapple's default `search` resolver hits every page regardless of publish
-    state, exposing drafts on the public API. Prepend a mixin so MRO picks our
-    `resolve_search`, which restricts Page subclasses to live + public.
+    Override Grapple's `search` resolver. Two fixes vs. the upstream version:
+    1. Restrict pages to live + public so drafts and access-restricted pages
+       don't leak via the public API.
+    2. Run a single search across all `Page` subclasses (instead of iterating
+       per-model) so results are ranked by relevance across types rather than
+       grouped by content type. Specific instances are fetched in a second
+       bulk query and reordered to match the search ranking.
+
+    Documents and images are intentionally not searched. The upstream resolver
+    includes them, but the frontend search page only renders Page types and
+    discards everything else, so iterating those indexes is wasted work.
     """
     if not grapple_registry.class_models:
         return
 
-    class SearchLivePublicMixin:
+    class SearchOverrideMixin:
         def resolve_search(self, info, **kwargs):
             query = kwargs.get("query")
             if not query:
                 return None
             s = get_search_backend()
-            results = []
-            models = [get_document_model(), get_image_model()]
-            for app in grapple_registry.apps:
-                models += django_apps.all_models[app].values()
-            for model in models:
-                if issubclass(model, Page):
-                    results += s.search(query, model.objects.live().public())
-                else:
-                    results += s.search(query, model)
-            return results
+            ranked = list(s.search(query, Page.objects.live().public()))
+            if not ranked:
+                return []
+            ids = [p.id for p in ranked]
+            specific_map = {p.id: p for p in Page.objects.filter(id__in=ids).specific()}
+            return [specific_map[i] for i in ids if i in specific_map]
 
-    query_mixins.insert(0, SearchLivePublicMixin)
+    query_mixins.insert(0, SearchOverrideMixin)
 
 
 @hooks.register("construct_page_action_menu")