Fix: custom document model not working with wagtail grapple

example/documents/migrations/0001_initial.py

@@ -0,0 +1,84 @@
+# Generated by Django 3.1.6 on 2021-02-06 17:44
+
+from django.conf import settings
+from django.db import migrations, models
+import django.db.models.deletion
+import taggit.managers
+import wagtail.core.models
+import wagtail.search.index
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ("wagtailcore", "0059_apply_collection_ordering"),
+        ("taggit", "0003_taggeditem_add_unique_index"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="CustomDocument",
+            fields=[
+                (
+                    "id",
+                    models.AutoField(
+                        auto_created=True,
+                        primary_key=True,
+                        serialize=False,
+                        verbose_name="ID",
+                    ),
+                ),
+                ("title", models.CharField(max_length=255, verbose_name="title")),
+                ("file", models.FileField(upload_to="documents", verbose_name="file")),
+                (
+                    "created_at",
+                    models.DateTimeField(auto_now_add=True, verbose_name="created at"),
+                ),
+                ("file_size", models.PositiveIntegerField(editable=False, null=True)),
+                (
+                    "file_hash",
+                    models.CharField(blank=True, editable=False, max_length=40),
+                ),
+                (
+                    "collection",
+                    models.ForeignKey(
+                        default=wagtail.core.models.get_root_collection_id,
+                        on_delete=django.db.models.deletion.CASCADE,
+                        related_name="+",
+                        to="wagtailcore.collection",
+                        verbose_name="collection",
+                    ),
+                ),
+                (
+                    "tags",
+                    taggit.managers.TaggableManager(
+                        blank=True,
+                        help_text=None,
+                        through="taggit.TaggedItem",
+                        to="taggit.Tag",
+                        verbose_name="tags",
+                    ),
+                ),
+                (
+                    "uploaded_by_user",
+                    models.ForeignKey(
+                        blank=True,
+                        editable=False,
+                        null=True,
+                        on_delete=django.db.models.deletion.SET_NULL,
+                        to=settings.AUTH_USER_MODEL,
+                        verbose_name="uploaded by user",
+                    ),
+                ),
+            ],
+            options={
+                "verbose_name": "document",
+                "verbose_name_plural": "documents",
+                "abstract": False,
+            },
+            bases=(wagtail.search.index.Indexed, models.Model),
+        ),
+    ]

example/documents/models.py

@@ -0,0 +1,14 @@
+from django.db import models
+
+from wagtail.documents.models import Document, AbstractDocument
+
+from grapple.models import GraphQLString
+
+
+class CustomDocument(AbstractDocument):
+    admin_form_fields = Document.admin_form_fields
+
+    def custom_document_property(self, info, **kwargs):
+        return "Document Model!"
+
+    graphql_fields = (GraphQLString("custom_document_property", required=True),)

example/example/settings/base.py

@@ -29,6 +29,7 @@
 INSTALLED_APPS = [
     "home",
     "images",
+    "documents",
     "search",
     "wagtail.contrib.forms",
     "wagtail.contrib.redirects",
@@ -159,6 +160,7 @@
 
 WAGTAIL_SITE_NAME = "example"
 WAGTAILIMAGES_IMAGE_MODEL = "images.CustomImage"
+WAGTAILDOCS_DOCUMENT_MODEL = "documents.CustomDocument"
 WAGTAILDOCS_SERVE_METHOD = "serve_view"
 
 # Base URL to use when referring to full URLs within the Wagtail admin backend -
@@ -169,7 +171,7 @@
 
 # Grapple Config:
 GRAPHENE = {"SCHEMA": "grapple.schema.schema"}
-GRAPPLE_APPS = {"images": "", "home": ""}
+GRAPPLE_APPS = {"images": "", "home": "", "documents": ""}
 GRAPPLE_ADD_SEARCH_HIT = True
 
 HEADLESS_PREVIEW_CLIENT_URLS = {"default": "http://localhost:8001/preview"}

example/example/tests/test_grapple.py

@@ -448,6 +448,7 @@ def test_query_documents_id(self):
         {
             documents {
                 id
+                customDocumentProperty
             }
         }
         """
@@ -460,6 +461,10 @@ def test_query_documents_id(self):
         self.assertEquals(
             executed["data"]["documents"][0]["id"], str(self.example_document.id)
         )
+        self.assertEquals(
+            executed["data"]["documents"][0]["customDocumentProperty"],
+            "Document Model!",
+        )
 
     def test_query_file_field(self):
         query = """

example/home/migrations/0002_create_homepage.py

@@ -48,6 +48,9 @@ def remove_homepage(apps, schema_editor):
 
 
 class Migration(migrations.Migration):
+    run_before = [
+        ("wagtailcore", "0053_locale_model"),
+    ]
 
     dependencies = [("home", "0001_initial")]
 

example/home/migrations/0021_auto_20210401_0935.py

@@ -0,0 +1,26 @@
+# Generated by Django 3.1.7 on 2021-04-01 09:35
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("documents", "0001_initial"),
+        ("home", "0020_correct_cover_model_fk"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="blogpage",
+            name="book_file",
+            field=models.ForeignKey(
+                blank=True,
+                null=True,
+                on_delete=django.db.models.deletion.SET_NULL,
+                related_name="+",
+                to="documents.customdocument",
+            ),
+        ),
+    ]

example/home/models.py

@@ -14,6 +14,7 @@
 from wagtail.snippets.edit_handlers import SnippetChooserPanel
 from wagtail.images.edit_handlers import ImageChooserPanel
 from wagtail.documents.edit_handlers import DocumentChooserPanel
+from wagtail.documents import get_document_model_string
 
 from wagtail_headless_preview.models import HeadlessPreviewMixin
 from wagtailmedia.edit_handlers import MediaChooserPanel
@@ -83,7 +84,7 @@ class BlogPage(HeadlessPreviewMixin, Page):
         related_name="+",
     )
     book_file = models.ForeignKey(
-        "wagtaildocs.Document",
+        get_document_model_string(),
         null=True,
         blank=True,
         on_delete=models.SET_NULL,

grapple/types/documents.py

@@ -3,13 +3,14 @@
 from graphene_django.types import DjangoObjectType
 
 from wagtail import VERSION as WAGTAIL_VERSION
-from wagtail.documents.models import Document as WagtailDocument
 
 if WAGTAIL_VERSION < (2, 9):
     from wagtail.documents.models import get_document_model
 else:
     from wagtail.documents import get_document_model
 
+from wagtail.documents.models import Document as WagtailDocument
+
 from ..registry import registry
 from ..utils import get_media_item_url, resolve_queryset
 from .collections import CollectionObjectType
@@ -42,10 +43,14 @@ def resolve_url(self, info, **kwargs):
         return get_media_item_url(self)
 
 
+def get_document_type():
+    mdl = get_document_model()
+    return registry.documents.get(mdl, DocumentObjectType)
+
+
 def DocumentsQuery():
-    registry.documents[WagtailDocument] = DocumentObjectType
     mdl = get_document_model()
-    mdl_type = registry.documents[mdl]
+    mdl_type = get_document_type()
 
     class Mixin:
         document = graphene.Field(mdl_type, id=graphene.ID())
@@ -57,6 +62,7 @@ class Mixin:
                 graphene.ID, description="Filter by collection id"
             ),
         )
+        document_type = graphene.String(required=True)
 
         def resolve_document(self, info, id, **kwargs):
             """Returns a document given the id, if in a public collection"""
@@ -73,12 +79,6 @@ def resolve_documents(self, info, **kwargs):
             return resolve_queryset(qs, info, **kwargs)
 
         def resolve_document_type(self, info, **kwargs):
-            return mdl_type
+            return get_document_type()
 
     return Mixin
-
-
-def get_document_type():
-    registry.documents[WagtailDocument] = DocumentObjectType
-    mdl = get_document_model()
-    return registry.documents[mdl]