Ability for BLOCK features to prefetch all objects to reduce database queries

wagtail_editorjs/editorjs/features.py

@@ -89,7 +89,7 @@ def validate(self, data: Any):
         if data["data"]["style"] not in ["ordered", "unordered"]:
             raise ValueError("Invalid style value")
 
-    def render_block_data(self, block: EditorJSBlock, context = None) -> EditorJSElement:
+    def render_block_data(self, block: EditorJSBlock, prefetch_data: Any, context = None) -> EditorJSElement:
         element = "ol" if block["data"]["style"] == "ordered" else "ul"
         return parse_list(block["data"]["items"], element)
 
@@ -148,7 +148,7 @@ def validate(self, data: Any):
             if "text" not in item:
                 raise ValueError("Invalid text value")
 
-    def render_block_data(self, block: EditorJSBlock, context = None) -> EditorJSElement:
+    def render_block_data(self, block: EditorJSBlock, prefetch_data: Any, context = None) -> EditorJSElement:
         s = []
         for item in block["data"]["items"]:
             class_ = "checklist-item"
@@ -186,7 +186,7 @@ def validate(self, data: Any):
         if 'code' not in data['data']:
             raise ValueError('Invalid code value')
 
-    def render_block_data(self, block: EditorJSBlock, context = None) -> EditorJSElement:
+    def render_block_data(self, block: EditorJSBlock, prefetch_data: Any, context = None) -> EditorJSElement:
         return EditorJSElement("code", block["data"]["code"], attrs={"class": "code"})
 
     @classmethod
@@ -201,7 +201,7 @@ class DelimiterFeature(EditorJSFeature):
     allowed_tags = ["hr"]
     allowed_attributes = ["class"]
 
-    def render_block_data(self, block: EditorJSBlock, context = None) -> EditorJSElement:
+    def render_block_data(self, block: EditorJSBlock, prefetch_data: Any, context = None) -> EditorJSElement:
         return EditorJSElement("hr", close_tag=False, attrs={"class": "delimiter"})
 
     @classmethod
@@ -219,7 +219,7 @@ def validate(self, data: Any):
         if level > 6 or level < 1:
             raise ValueError("Invalid level value")
 
-    def render_block_data(self, block: EditorJSBlock, context = None) -> EditorJSElement:
+    def render_block_data(self, block: EditorJSBlock, prefetch_data: Any, context = None) -> EditorJSElement:
         return EditorJSElement(
             "h" + str(block["data"]["level"]),
             block["data"].get("text")
@@ -245,7 +245,7 @@ def validate(self, data: Any):
         if "html" not in data["data"]:
             raise ValueError("Invalid html value")
 
-    def render_block_data(self, block: EditorJSBlock, context = None) -> EditorJSElement:
+    def render_block_data(self, block: EditorJSBlock, prefetch_data: Any, context = None) -> EditorJSElement:
         return EditorJSElement("div", block["data"]["html"], attrs={"class": "html"})
 
     @classmethod
@@ -269,7 +269,7 @@ def validate(self, data: Any):
         if "message" not in data["data"]:
             raise ValueError("Invalid message value")
 
-    def render_block_data(self, block: EditorJSBlock, context = None) -> EditorJSElement:
+    def render_block_data(self, block: EditorJSBlock, prefetch_data: Any, context = None) -> EditorJSElement:
         return EditorJSElement(
             "div",
             attrs={
@@ -333,6 +333,21 @@ def get_config(self, context: dict[str, Any]):
             f"editorjs-image-chooser-{context['widget']['attrs']['id']}"
         config["config"]["getImageUrl"] = reverse("wagtail_editorjs:image_for_id_fmt")
         return config
+
+    def get_prefetch_data(self, block: EditorJSBlock, context=None):
+        return block["data"]["imageId"]
+
+    def prefetch_data(self, data: list[tuple[int, EditorJSBlock, Any]], context=None):
+        ids = [i[2] for i in data]
+
+        images_qs = Image.objects.in_bulk(ids)
+
+        for j, (i, block, prefetch_data) in enumerate(data):
+            try:
+                prefetch_data = int(prefetch_data)
+            except ValueError:
+                pass
+            data[j] = (i, block, images_qs[prefetch_data])
 
     def validate(self, data: Any):
         super().validate(data)
@@ -350,10 +365,7 @@ def render_template(self, context: dict[str, Any] = None):
             {"id": widget_id}
         )
 
-    def render_block_data(self, block: EditorJSBlock, context = None) -> EditorJSElement:
-        image = block["data"].get("imageId")
-        image = Image.objects.get(id=image)
-
+    def render_block_data(self, block: EditorJSBlock, prefetch_data: Any, context = None) -> EditorJSElement:
         classlist = []
         styles = {}
         if block["data"].get("withBorder"):
@@ -373,7 +385,7 @@ def render_block_data(self, block: EditorJSBlock, context = None) -> EditorJSEle
         if styles:
             attrs["style"] = styles
 
-        url = image.file.url
+        url = prefetch_data.file.url
         if not any([url.startswith(i) for i in ["http://", "https://", "//"]]) and context:
             request = context.get("request")
             if request:
@@ -478,21 +490,34 @@ def validate(self, data: Any):
 
             if "title" not in image:
                 raise ValueError("Invalid title value")
-
-    def render_block_data(self, block: EditorJSBlock, context = None) -> EditorJSElement:
-        images = block["data"]["images"]
+
+    def get_prefetch_data(self, block: EditorJSBlock, context=None):
+        data = []
+        for image in block["data"]["images"]:
+            data.append(image["id"])
+        return data
+
+    def prefetch_data(self, data: list[tuple[int, EditorJSBlock, Any]], context=None):
         ids = []
-        for image in images:
-            ids.append(image["id"])
+        for i, block, prefetch_data in data:
+            ids.extend(prefetch_data)
+
+        images_qs = Image.objects.in_bulk(ids)
 
-        images = Image.objects.in_bulk(ids)
+        for j, (i, block, prefetch_data) in enumerate(data):
+            prefetch_data = []
+            for image in block["data"]["images"]:
+                try:
+                    id = int(image["id"])
+                except ValueError:
+                    pass
+                prefetch_data.append(images_qs[id])
+            data[j] = (i, block, prefetch_data)
+
+
+    def render_block_data(self, block: EditorJSBlock, prefetch_data: Any, context = None) -> EditorJSElement:
         s = []
-        for id in ids:
-            try:
-                id = int(id)
-            except ValueError:
-                pass
-            image = images[id]
+        for image in prefetch_data:
             url = image.file.url
             if not any([url.startswith(i) for i in ["http://", "https://", "//"]]) and context:
                 request = context.get("request")
@@ -542,7 +567,7 @@ def validate(self, data: Any):
         if "withHeadings" not in data["data"]:
             raise ValueError("Invalid withHeadings value")
 
-    def render_block_data(self, block: EditorJSBlock, context = None) -> EditorJSElement:
+    def render_block_data(self, block: EditorJSBlock, prefetch_data: Any, context = None) -> EditorJSElement:
         table = []
         for i, row in enumerate(block["data"]["content"]):
             tr = []
@@ -589,7 +614,7 @@ def validate(self, data: Any):
             raise ValueError("Invalid caption value")
 
 
-    def render_block_data(self, block: EditorJSBlock, context = None) -> EditorJSElement:
+    def render_block_data(self, block: EditorJSBlock, prefetch_data: Any, context = None) -> EditorJSElement:
         text = block["data"]["text"]
         caption = block["data"]["caption"]
         return EditorJSElement(
@@ -639,11 +664,23 @@ def validate(self, data: Any):
         if "title" not in data["data"]:
             raise ValueError("Invalid title value")
 
-    def render_block_data(self, block: EditorJSBlock, context = None) -> EditorJSElement:
+    def get_prefetch_data(self, block: EditorJSBlock, context=None):
+        return block["data"]["file"]["id"]
+
+    def prefetch_data(self, data: list[tuple[int, EditorJSBlock, Any]], context=None):
+        ids = [i[2] for i in data]
+        documents = Document.objects.in_bulk(ids)
+
+        for j, (i, block, prefetch_data) in enumerate(data):
+            try:
+                prefetch_data = int(prefetch_data)
+            except ValueError:
+                pass
+            data[j] = (i, block, documents[prefetch_data])
+
+    def render_block_data(self, block: EditorJSBlock, prefetch_data: Any, context = None) -> EditorJSElement:
 
-        document_id = block["data"]["file"]["id"]
-        document = Document.objects.get(pk=document_id)
-        url = document.url
+        url = prefetch_data.url
 
         if not any([url.startswith(i) for i in ["http://", "https://", "//"]]) and context:
             request = context.get("request")
@@ -653,8 +690,8 @@ def render_block_data(self, block: EditorJSBlock, context = None) -> EditorJSEle
         if block["data"]["title"]:
             title = block["data"]["title"]
         else:
-            if document:
-                title = document.title
+            if prefetch_data:
+                title = prefetch_data.title
             else:
                 title = url
 
@@ -672,7 +709,7 @@ def render_block_data(self, block: EditorJSBlock, context = None) -> EditorJSEle
                 ),
                 EditorJSElement(
                     "span",
-                    filesize_to_human_readable(document.file.size),
+                    filesize_to_human_readable(prefetch_data.file.size),
                     attrs={"class": "attaches-size"},
                 ),
                 EditorJSElement(

wagtail_editorjs/registry.py

@@ -192,7 +192,13 @@ def validate(self, data: Any):
         if "data" not in data:
             raise ValueError("Invalid data format")
 
-    def render_block_data(self, block: EditorJSBlock, context = None) -> "EditorJSElement":
+    def get_prefetch_data(self, block: EditorJSBlock, context = None):
+        return None
+
+    def prefetch_data(self, data: list[tuple[int, EditorJSBlock, Any]], context = None):
+        pass
+
+    def render_block_data(self, block: EditorJSBlock, prefetch_data, context = None) -> "EditorJSElement":
         return EditorJSElement(
             "p",
             block["data"].get("text")
@@ -395,18 +401,10 @@ def build_elements(self, inline_data: list, context: dict[str, Any] = None) -> l
         ids = []
         # element_soups = []
         for data in inline_data:
-            # soup: BeautifulSoup
-            # element: EditorJSElement
-            # matches: dict[bs4.elementType, dict[str, Any]]
-            # data: dict[str, Any] # Block data.
-            item, data = data
-
-            # # Store element and soup for later replacement of content.
-            # element_soups.append((soup, element))
+            item, attrs = data
 
-            # Item is bs4 tag, attrs are must_have_attrs
-
-            id = self.get_id(item, data, context)
+            # Item is bs4 tag, attrs are must_have_attrs and can_have_attrs
+            id = self.get_id(item, attrs, context)
             ids.append((item, id))
 
             # delete all attributes

wagtail_editorjs/render.py

@@ -6,6 +6,7 @@
 from .registry import (
     EditorJSElement,
     BaseInlineEditorJSFeature,
+    EditorJSFeature,
     InlineEditorJSFeature,
     EDITOR_JS_FEATURES,
 )
@@ -38,7 +39,7 @@ def render_editorjs_html(features: list[str], data: dict, context=None, clean: b
         if isinstance(feature, BaseInlineEditorJSFeature)
     ]
 
-    html = []
+    blocks: list[tuple[EditorJSFeature, dict, Any]] = []
     for block in data["blocks"]:
 
         feature: str = block["type"]
@@ -48,14 +49,27 @@ def render_editorjs_html(features: list[str], data: dict, context=None, clean: b
         if not feature_mapping:
             continue
 
-        # Build the actual block.
-        element: EditorJSElement = feature_mapping.render_block_data(block, context)
+        # Get the prefetch data for this block
+        prefetch_data = feature_mapping.get_prefetch_data(block, context)
+        blocks.append((feature_mapping, block, prefetch_data))
+
+    # Prefetch the data
+    prefetch_map: dict[EditorJSFeature, list[tuple[int, dict, Any]]] = defaultdict(list)
+    for i, (feature_mapping, block, prefetch_data) in enumerate(blocks):
+        prefetch_map[feature_mapping].append((i, block, prefetch_data))
+
+    for feature_mapping, data in prefetch_map.items():
+        feature_mapping.prefetch_data(data, context)
 
-        # if element.tag != "div":
-        #     new = EditorJSElement("div", [element], attrs=element.attrs)
-        #     element.attrs = {}
-        #     element = new
+        for i, block, prefetch_data in data:
+            blocks[i] = (feature_mapping, block, prefetch_data)
 
+    html = []
+    for feature_mapping, block, prefetch_data in blocks:
+        tunes: dict[str, Any] = block.get("tunes", {})
+
+        # Build the actual block.
+        element: EditorJSElement = feature_mapping.render_block_data(block, prefetch_data, context)
 
         # Tune the element.
         for tune_name, tune_value in tunes.items():

wagtail_editorjs/tests.py

@@ -30,7 +30,7 @@ def test_editorjs_features(self):
 
             for data in test_data_list:
                 if hasattr(feature, "render_block_data"):
-                    tpl = feature.render_block_data(data)
+                    tpl = feature.render_block_data(data, None)
                     html.append(tpl)
 
         rendered_1 = render_editorjs_html(