From 09d1078dce54b9d139d1331d062310911f892397 Mon Sep 17 00:00:00 2001
From: Jonas Braathen <jonas@ponas.no>
Date: Tue, 26 May 2026 00:01:37 +0200
Subject: [PATCH] dnscms: hide drafts from search results

---
 dnscms/dnscms/wagtail_hooks.py | 36 ++++++++++++++
 dnscms/tests/test_search.py    | 88 ++++++++++++++++++++++++++++++++++
 2 files changed, 124 insertions(+)
 create mode 100644 dnscms/tests/test_search.py

diff --git a/dnscms/dnscms/wagtail_hooks.py b/dnscms/dnscms/wagtail_hooks.py
index a71a06c..bef3e4e 100644
--- a/dnscms/dnscms/wagtail_hooks.py
+++ b/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):
diff --git a/dnscms/tests/test_search.py b/dnscms/tests/test_search.py
new file mode 100644
index 0000000..6a183af
--- /dev/null
+++ b/dnscms/tests/test_search.py
@@ -0,0 +1,88 @@
+from wagtail.search.backends import get_search_backend
+
+from tests.conftest import EventPageFactory, GenericPageFactory
+
+
+SEARCH_QUERY = """
+query Search($query: String) {
+  results: search(query: $query) {
+    __typename
+    ... on PageInterface {
+      title
+    }
+  }
+}
+"""
+
+
+def _index(page):
+    # Wagtail's post_save signal enqueues indexing via django-tasks, which isn't
+    # drained synchronously in tests. Call the backend directly so the page is
+    # findable through the live search code path.
+    get_search_backend().add(page)
+
+
+def _titles_for(body, typename):
+    return [r["title"] for r in body["data"]["results"] if r["__typename"] == typename]
+
+
+def test_search_returns_live_generic_page(home_page, graphql_post):
+    page = GenericPageFactory(
+        parent=home_page,
+        title="PublishedGenericSearchToken",
+        slug="published-generic-search",
+    )
+    _index(page)
+
+    response, body = graphql_post(SEARCH_QUERY, {"query": "PublishedGenericSearchToken"})
+
+    assert response.status_code == 200
+    assert "errors" not in body, body
+    assert "PublishedGenericSearchToken" in _titles_for(body, "GenericPage")
+
+
+def test_search_excludes_draft_generic_page(home_page, graphql_post):
+    page = GenericPageFactory(
+        parent=home_page,
+        title="DraftGenericSearchToken",
+        slug="draft-generic-search",
+        live=False,
+    )
+    _index(page)
+
+    response, body = graphql_post(SEARCH_QUERY, {"query": "DraftGenericSearchToken"})
+
+    assert response.status_code == 200
+    assert "errors" not in body, body
+    assert "DraftGenericSearchToken" not in _titles_for(body, "GenericPage")
+
+
+def test_search_returns_live_event_page(home_page, event_index, graphql_post):
+    page = EventPageFactory(
+        parent=event_index,
+        title="PublishedEventSearchToken",
+        slug="published-event-search",
+    )
+    _index(page)
+
+    response, body = graphql_post(SEARCH_QUERY, {"query": "PublishedEventSearchToken"})
+
+    assert response.status_code == 200
+    assert "errors" not in body, body
+    assert "PublishedEventSearchToken" in _titles_for(body, "EventPage")
+
+
+def test_search_excludes_draft_event_page(home_page, event_index, graphql_post):
+    page = EventPageFactory(
+        parent=event_index,
+        title="DraftEventSearchToken",
+        slug="draft-event-search",
+        live=False,
+    )
+    _index(page)
+
+    response, body = graphql_post(SEARCH_QUERY, {"query": "DraftEventSearchToken"})
+
+    assert response.status_code == 200
+    assert "errors" not in body, body
+    assert "DraftEventSearchToken" not in _titles_for(body, "EventPage")