dnscms: hide drafts from search results

dnscms/dnscms/wagtail_hooks.py

@@ -1,6 +1,12 @@
+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
 
 
 @hooks.register("register_rich_text_features")
@@ -8,6 +14,36 @@ def enable_additional_rich_text_features(features):
     features.default_features.extend(["h5", "h6", "blockquote"])
 
 
+@hooks.register("register_schema_query")
+def filter_search_to_live_pages(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.
+    """
+    if not grapple_registry.class_models:
+        return
+
+    class SearchLivePublicMixin:
+        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
+
+    query_mixins.insert(0, SearchLivePublicMixin)
+
+
 @hooks.register("construct_page_action_menu")
 def make_publish_default_action(menu_items, request, context):
     for index, item in enumerate(menu_items):