Fix previews for unpublished pages (#277)

example/example/tests/test_previews.py

@@ -0,0 +1,47 @@
+from datetime import date
+
+from django.test import TestCase
+from home.models import BlogPage, HomePage
+
+from grapple.types.pages import get_preview_page
+
+
+class TestUnsavedPagePreview(TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        self.home = HomePage.objects.first()
+        self.page = BlogPage(
+            title="Unsaved Page", slug="unsaved-page", date=date(2020, 1, 1)
+        )
+        # mock an unsaved page
+        # there must be a better way
+        # using the form libs
+        self.page._cached_parent_obj = self.home
+        self.page.path = "000100010001"
+        self.page.steplen = 4
+        self.preview = self.page.create_page_preview()
+        self.token = self.preview.token
+
+    def test_get_preview_page(self):
+        self.assertEqual(self.page.pk, None)
+        preview = get_preview_page(self.token)
+        self.assertEqual(preview.title, self.page.title)
+        self.assertEqual(preview.title, self.page.title)
+        self.assertEqual(preview.slug, self.page.slug)
+
+
+class TestSavedPagePreview(TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        self.home = HomePage.objects.first()
+        self.page = BlogPage(title="Saved Page", slug="Saved-page", date="2020-01-01")
+        self.home.add_child(instance=self.page)
+        self.preview = self.page.create_page_preview()
+        self.token = self.preview.token
+
+    def test_get_preview_page(self):
+        self.assertEqual(self.page.pk, 4)
+        preview = get_preview_page(self.token)
+        self.assertEqual(preview.title, self.page.title)
+        self.assertEqual(preview.slug, self.page.slug)
+        self.assertEqual(preview.id, self.page.id)

grapple/types/pages.py

@@ -169,6 +169,66 @@ class Meta:
         interfaces = (PageInterface,)
 
 
+def get_preview_page(token):
+    """
+    Get a preview page from a token.
+    """
+    try:
+        token_params_list = token.split(":").pop(0)
+        token_params_kvstr = token_params_list.split(";")
+
+        params = {}
+        for arg in token_params_kvstr:
+            key, value = arg.split("=")
+            params[key] = value
+
+        id = params.get("id")
+        if id:
+            """
+            This is a page that had already been saved. Lookup the class and call get_page_from_preview_token.
+
+            TODO: update headless preview to always send page_type in the token so we can always
+            the if content_type branch and elimiate the if id branch.
+            """
+            page = WagtailPage.objects.get(pk=id).specific
+            if page:
+                cls = type(page)
+                """
+                get_page_from_preview_token is added by wagtail-headless-preview,
+                this condition checks that headless-preview is installed and enabled
+                for the model.
+                """
+                if hasattr(cls, "get_page_from_preview_token"):
+                    """we assume that get_page_from_preview_token validates the token"""
+                    return cls.get_page_from_preview_token(token)
+
+        content_type = params.get("page_type")
+        if content_type:
+            """
+            this is a page which has not been saved yet. lookup the content_type to get the class
+            and call get_page_from_preview_token.
+            """
+            app_label, model = content_type.lower().split(".")
+            ctype = ContentType.objects.get(app_label=app_label, model=model)
+            if ctype:
+                cls = ctype.model_class()
+                """
+                get_page_from_preview_token is added by wagtail-headless-preview,
+                this condition checks that headless-preview is installed and enabled
+                for the model.
+                """
+                if hasattr(cls, "get_page_from_preview_token"):
+                    """we assume that get_page_from_preview_token validates the token"""
+                    return cls.get_page_from_preview_token(token)
+    except Exception:
+        """
+        catch and suppress errors. we don't want to expose any information about unpublished content
+        accidentally.
+        TODO: consider logging here.
+        """
+        return None
+
+
 def get_specific_page(
     id=None, slug=None, url_path=None, token=None, content_type=None, site=None
 ):
@@ -177,9 +237,13 @@ def get_specific_page(
     """
     page = None
     try:
+        if token:
+            return get_preview_page(token)
+
         # Everything but the special RootPage
         qs = WagtailPage.objects.live().public().filter(depth__gt=1).specific()
         ctype = None
+
         if site:
             qs = qs.in_site(site)
 
@@ -210,17 +274,6 @@ def get_specific_page(
             if qs.exists():
                 page = qs.first()
 
-        # If token provided then get draft/preview
-        if token:
-            if page:
-                page_type = type(page)
-                if hasattr(page_type, "get_page_from_preview_token"):
-                    page = page_type.get_page_from_preview_token(token)
-
-            elif ctype:
-                cls = ctype.model_class()
-                if hasattr(cls, "get_page_from_preview_token"):
-                    page = cls.get_page_from_preview_token(token)
     except BaseException:
         page = None