Improve mapping performance by prefetching relations (#39)

* Initial attempt to improve mapping speed by prefetching related models that are used in the mappings.

* More performance improvements and add test to measure amount of queries.

* Get it down to 19 queries. Update fixtures this version will require an oscar version update.

* The public children and browsable categories are now manager methods, makes it easier to use.

* undo apps.py change

* Pass the model instance to extra_attrs and save it as a property on the resource, so later resources that use this resource can access the model instance (original sroouce)

* Implement the 'prefetch_attribute_values' method from oscar. Pass the model instance to extra_attrs of the resource and save it as a property so future resources can have access to this.

* Up num query by one, as we added one more prefetch in django-oscar.

* Create a prefetch registry, as before it was almost impossible to register new prefetches, replace prefetches with custom queryset, remove prefetches etc.

* Dont use Prefetch class if its not needed.

* Add unit tests for the PrefetchRegistry

---------

Co-authored-by: Joey Jurjens <joey@highbiza.nl>

oscar_odin/apps.py

@@ -18,3 +18,8 @@ def ready(self):
 
         # Register the Django model field resolver
         registration.register_field_resolver(ModelFieldResolver, Model)
+
+        # Register the default prefetches for the product queryset
+        from oscar_odin.mappings.prefetching.prefetch import register_default_prefetches
+
+        register_default_prefetches()

oscar_odin/fixtures/oscar_odin/order.json

@@ -67,7 +67,8 @@
         "unit_price_incl_tax": "10.00",
         "unit_price_excl_tax": "10.00",
         "tax_code": null,
-        "status": "Pending"
+        "status": "Pending",
+        "num_allocated": 0
     }
 },
 {
@@ -92,7 +93,8 @@
         "unit_price_incl_tax": "10.00",
         "unit_price_excl_tax": "10.00",
         "tax_code": null,
-        "status": "Pending"
+        "status": "Pending",
+        "num_allocated": 0
     }
 },
 {

oscar_odin/fixtures/oscar_odin/test_discount/order.json

@@ -155,7 +155,8 @@
         "unit_price_incl_tax": "349.99",
         "unit_price_excl_tax": "289.25",
         "tax_code": "high",
-        "status": "paid"
+        "status": "paid",
+        "num_allocated": 0
     }
 },
 {
@@ -180,7 +181,8 @@
         "unit_price_incl_tax": "69.99",
         "unit_price_excl_tax": "57.84",
         "tax_code": "high",
-        "status": "paid"
+        "status": "paid",
+        "num_allocated": 0
     }
 },
 {
@@ -205,7 +207,8 @@
         "unit_price_incl_tax": "39.99",
         "unit_price_excl_tax": "36.69",
         "tax_code": "low",
-        "status": "paid"
+        "status": "paid",
+        "num_allocated": 0
     }
 },
 {
@@ -230,7 +233,8 @@
         "unit_price_incl_tax": "34.99",
         "unit_price_excl_tax": "28.92",
         "tax_code": "high",
-        "status": "paid"
+        "status": "paid",
+        "num_allocated": 0
     }
 },
 {
@@ -255,7 +259,8 @@
         "unit_price_incl_tax": "14.99",
         "unit_price_excl_tax": "13.75",
         "tax_code": "low",
-        "status": "paid"
+        "status": "paid",
+        "num_allocated": 0
     }
 },
 {
@@ -280,7 +285,8 @@
         "unit_price_incl_tax": "349.99",
         "unit_price_excl_tax": "289.25",
         "tax_code": "high",
-        "status": "paid"
+        "status": "paid",
+        "num_allocated": 0
     }
 },
 {

oscar_odin/mappings/_common.py

@@ -1,8 +1,9 @@
 """Common code between mappings."""
 from typing import Any, Dict, Optional, Type
+from django.db.models import QuerySet, Model
+from django.db.models.manager import BaseManager
 
 import odin
-from django.db.models import QuerySet
 from odin.mapping import ImmediateResult, MappingBase, MappingMeta
 
 
@@ -14,8 +15,7 @@ def map_queryset(
 ) -> list:
     """Map a queryset to a list of resources.
 
-    This method will call ``QuerySet.all()`` to ensure that the queryset can
-    be directly iterated.
+    This method will ensure that the queryset can be directly iterated.
 
     :param mapping: The mapping type to use.
     :param queryset: The queryset to map.
@@ -26,8 +26,12 @@ def map_queryset(
         raise ValueError(
             f"Mapping {mapping} cannot map queryset of type {queryset.model}"
         )
+
+    if isinstance(queryset, BaseManager):
+        queryset = queryset.all()
+
     return list(
-        mapping.apply(queryset.all(), context=context, mapping_result=ImmediateResult)
+        mapping.apply(list(queryset), context=context, mapping_result=ImmediateResult)
     )
 
 
@@ -42,8 +46,24 @@ def create_object(self, **field_values):
             for key, field_value in field_values.items():
                 setattr(new_obj, key, field_value)
 
+            self.pass_model_instance(new_obj)
+
             return new_obj
         except AttributeError:
-            return super().create_object(**field_value)
+            new_obj = super().create_object(**field_value)
+            self.pass_model_instance(new_obj)
+            return new_obj
+
+    def pass_model_instance(self, obj):
+        """
+        Passes the model instance (original source) into the extra_attrs method of the resource.
+        The resource can then use this to store the model instance as a property for later resources.
+        This is useful, as the base resource for each model, isn't saving every model field on the resource.
+        Though, later resources that gets mapped from the base resource, might need to access a certain fields
+        from the model instance, this way they can access it without doing a separate query, which is good for performance.
+        """
+        if isinstance(self.source, Model):
+            obj.extra_attrs({"model_instance": self.source})
+        return obj
 
     register_mapping = False

oscar_odin/mappings/catalogue.py

@@ -9,6 +9,7 @@
 from django.db.models import QuerySet
 from django.db.models.fields.files import ImageFieldFile
 from django.http import HttpRequest
+from odin.mapping import ImmediateResult
 from oscar.apps.partner.strategy import Default as DefaultStrategy
 from oscar.core.loading import get_class, get_model
 
@@ -18,6 +19,7 @@
 from ._common import map_queryset, OscarBaseMapping
 from ._model_mapper import ModelMapping
 from ..utils import validate_resources
+from .prefetching.prefetch import prefetch_product_queryset
 
 from .context import ProductModelMapperContext
 from .constants import ALL_CATALOGUE_FIELDS, MODEL_IDENTIFIERS_MAPPING
@@ -163,7 +165,12 @@ def images(self) -> List[resources.catalogue.Image]:
     def categories(self):
         """Map related categories."""
         items = self.source.get_categories()
-        return map_queryset(CategoryToResource, items, context=self.context)
+        # Note: categories are prefetched with the 'to_attr' method, this means it's a list and not a queryset.
+        return list(
+            CategoryToResource.apply(
+                items, context=self.context, mapping_result=ImmediateResult
+            )
+        )
 
     @odin.assign_field
     def product_class(self) -> str:
@@ -333,7 +340,7 @@ def product_to_resource_with_strategy(
 ):
     """Map a product model to a resource.
 
-    This method will except either a single product or an iterable of product
+    This method will accept either a single product or an iterable of product
     models (eg a QuerySet), and will return the corresponding resource(s).
     The request and user are optional, but if provided they are supplied to the
     partner strategy selector.
@@ -361,7 +368,7 @@ def product_to_resource(
 ) -> Union[resources.catalogue.Product, Iterable[resources.catalogue.Product]]:
     """Map a product model to a resource.
 
-    This method will except either a single product or an iterable of product
+    This method will accept either a single product or an iterable of product
     models (eg a QuerySet), and will return the corresponding resource(s).
     The request and user are optional, but if provided they are supplied to the
     partner strategy selector.
@@ -400,12 +407,10 @@ def product_queryset_to_resources(
     :param kwargs: Additional keyword arguments to pass to the strategy selector.
     """
 
-    query_set = queryset.prefetch_related(
-        "images", "product_class", "product_class__options"
-    )
+    queryset = prefetch_product_queryset(queryset, include_children)
 
     return product_to_resource(
-        query_set,
+        queryset,
         request,
         user,
         include_children,

oscar_odin/mappings/prefetching/prefetch.py

@@ -0,0 +1,94 @@
+from django.db.models import Prefetch
+
+from oscar.core.loading import get_class, get_model
+
+from .registry import prefetch_registry
+
+ProductQuerySet = get_class("catalogue.managers", "ProductQuerySet")
+
+ProductModel = get_model("catalogue", "Product")
+
+
+def prefetch_product_queryset(
+    queryset: ProductQuerySet, include_children: bool = False, **kwargs
+) -> ProductQuerySet:
+    """
+    Optimize the product queryset with registered select_related and prefetch_related operations.
+
+    Args:
+        queryset (ProductQuerySet): The initial queryset to optimize.
+        include_children (bool): Whether to include prefetches for children.
+
+    Returns:
+        ProductQuerySet: The optimized queryset.
+    """
+    callable_kwargs = {"include_children": include_children, **kwargs}
+
+    select_related_fields = prefetch_registry.get_select_related()
+    queryset = queryset.select_related(*select_related_fields)
+
+    prefetches = prefetch_registry.get_prefetches()
+    for prefetch in prefetches.values():
+        if isinstance(prefetch, (str, Prefetch)):
+            queryset = queryset.prefetch_related(prefetch)
+        elif callable(prefetch):
+            queryset = prefetch(queryset, **callable_kwargs)
+
+    if include_children:
+        children_prefetches = prefetch_registry.get_children_prefetches()
+        for prefetch in children_prefetches.values():
+            if isinstance(prefetch, (str, Prefetch)):
+                queryset = queryset.prefetch_related(prefetch)
+            elif callable(prefetch):
+                queryset = prefetch(queryset, **callable_kwargs)
+
+    return queryset
+
+
+def register_default_prefetches():
+    # ProductToResource.product_class -> get_product_class
+    prefetch_registry.register_select_related(["product_class", "parent"])
+
+    # ProductToResource.images -> get_all_images
+    prefetch_registry.register_prefetch("images")
+
+    # ProducToResource.map_stock_price -> fetch_for_product
+    prefetch_registry.register_prefetch("stockrecords")
+
+    # This gets prefetches somewhere (.categories.all()), it's not in get_categories as that does
+    # .browsable() and that's where the prefetch_browsable_categories is for. But if we remove this,
+    # the amount of queries will be more again. ToDo: Figure out where this is used and document it.
+    prefetch_registry.register_prefetch("categories")
+
+    # The parent and its related fields are prefetched in numerous places in the resource.
+    # ProductToResource.product_class -> get_product_class (takes parent product_class if itself has no product_class)
+    # ProductToResource.images -> get_all_images (takes parent images if itself has no images)
+    prefetch_registry.register_prefetch("parent__product_class")
+    prefetch_registry.register_prefetch("parent__images")
+
+    # ProducToResource.attributes -> get_attribute_values
+    def prefetch_attribute_values(queryset: ProductQuerySet, **kwargs):
+        return queryset.prefetch_attribute_values(
+            include_parent_children_attributes=kwargs.get("include_children", False)
+        )
+
+    prefetch_registry.register_prefetch(prefetch_attribute_values)
+
+    # ProductToResource.categories -> get_categories
+    # ProductToResource.categories -> get_categories -> looks up the parent categories if child
+    def prefetch_browsable_categories(queryset: ProductQuerySet, **kwargs):
+        return queryset.prefetch_browsable_categories()
+
+    prefetch_registry.register_prefetch(prefetch_browsable_categories)
+
+    # ProductToResource.map_stock_price -> fetch_for_parent -> product.children.public() -> stockrecords
+    def prefetch_public_children_stockrecords(queryset: ProductQuerySet, **kwargs):
+        return queryset.prefetch_public_children(
+            queryset=ProductModel.objects.public().prefetch_related("stockrecords")
+        )
+
+    prefetch_registry.register_prefetch(prefetch_public_children_stockrecords)
+
+    # Register children prefetches
+    prefetch_registry.register_children_prefetch("children__images")
+    prefetch_registry.register_children_prefetch("children__stockrecords")