From 873220ce62b24db948235cb84925493c8343e5ae Mon Sep 17 00:00:00 2001 From: Joey Date: Wed, 11 Sep 2024 08:58:41 +0200 Subject: [PATCH] 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 --- oscar_odin/apps.py | 5 + oscar_odin/fixtures/oscar_odin/order.json | 6 +- .../oscar_odin/test_discount/order.json | 18 +- oscar_odin/mappings/_common.py | 30 ++- oscar_odin/mappings/catalogue.py | 19 +- oscar_odin/mappings/prefetching/__init__.py | 0 oscar_odin/mappings/prefetching/prefetch.py | 94 ++++++++++ oscar_odin/mappings/prefetching/registry.py | 143 ++++++++++++++ oscar_odin/resources/_base.py | 13 ++ tests/mappings/test_catalogue.py | 34 ++++ tests/mappings/test_prefetch_registry.py | 177 ++++++++++++++++++ 11 files changed, 519 insertions(+), 20 deletions(-) create mode 100644 oscar_odin/mappings/prefetching/__init__.py create mode 100644 oscar_odin/mappings/prefetching/prefetch.py create mode 100644 oscar_odin/mappings/prefetching/registry.py create mode 100644 tests/mappings/test_prefetch_registry.py diff --git a/oscar_odin/apps.py b/oscar_odin/apps.py index 5325649..087c851 100644 --- a/oscar_odin/apps.py +++ b/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() diff --git a/oscar_odin/fixtures/oscar_odin/order.json b/oscar_odin/fixtures/oscar_odin/order.json index 5f90469..94d5fbd 100644 --- a/oscar_odin/fixtures/oscar_odin/order.json +++ b/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 } }, { diff --git a/oscar_odin/fixtures/oscar_odin/test_discount/order.json b/oscar_odin/fixtures/oscar_odin/test_discount/order.json index 7d120a1..e2802b5 100644 --- a/oscar_odin/fixtures/oscar_odin/test_discount/order.json +++ b/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 } }, { diff --git a/oscar_odin/mappings/_common.py b/oscar_odin/mappings/_common.py index f30196e..35edf5b 100644 --- a/oscar_odin/mappings/_common.py +++ b/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 diff --git a/oscar_odin/mappings/catalogue.py b/oscar_odin/mappings/catalogue.py index 7e49a21..9f03b9e 100644 --- a/oscar_odin/mappings/catalogue.py +++ b/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, diff --git a/oscar_odin/mappings/prefetching/__init__.py b/oscar_odin/mappings/prefetching/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/oscar_odin/mappings/prefetching/prefetch.py b/oscar_odin/mappings/prefetching/prefetch.py new file mode 100644 index 0000000..e5e3f9e --- /dev/null +++ b/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") diff --git a/oscar_odin/mappings/prefetching/registry.py b/oscar_odin/mappings/prefetching/registry.py new file mode 100644 index 0000000..8bdaa64 --- /dev/null +++ b/oscar_odin/mappings/prefetching/registry.py @@ -0,0 +1,143 @@ +from typing import Dict, Any, Union, Callable, List, Set + +from django.db.models import Prefetch + +from oscar.core.loading import get_class + +ProductQuerySet = get_class("catalogue.managers", "ProductQuerySet") + +PrefetchType = Union[ + str, Prefetch, Callable[[ProductQuerySet, Dict[str, Any]], ProductQuerySet] +] +SelectRelatedType = Union[str, List[str]] + + +class PrefetchRegistry: + """ + This class enables the flexibility to register prefetch_related and select_related operations + for the product_queryset_to_resources method. By default, django-oscar-odin prefetches all related + fields that are used by the default mapping. + + However, it's likely that you have your own resource(s) that makes use of the product_queryset_to_resources + method while doing additional queries. To make it easier to add your own prefetches, to prevent n+1 queries + you can register your own prefetches and select_related operations. + + You can also unregister default ones, and replace them with your own. For example when your way of + getting stockrecords does a different query, it's better to unregister the default one and register your own. + This way, you're also not doing a useless query. + """ + + def __init__(self): + self.prefetches: Dict[str, PrefetchType] = {} + self.children_prefetches: Dict[str, PrefetchType] = {} + self.select_related: Set[str] = set() + + def register_prefetch(self, prefetch: PrefetchType): + """ + Register a prefetch_related operation. + + Args: + prefetch (PrefetchType): The prefetch to register. Can be a string, Prefetch object, or a method. + """ + key = self._get_key(prefetch) + self.prefetches[key] = prefetch + + def register_children_prefetch(self, prefetch: PrefetchType): + """ + Register a children prefetch_related operation. Children as is, the children of a parent. + + Args: + prefetch (PrefetchType): The child prefetch to register. Can be a string, Prefetch object, or a method. + """ + key = self._get_key(prefetch) + self.children_prefetches[key] = prefetch + + def register_select_related(self, select: SelectRelatedType): + """ + Register a select_related operation. + + Args: + select (SelectRelatedType): The select_related to register. Can be a string or a list of strings. + """ + if isinstance(select, str): + self.select_related.add(select) + elif isinstance(select, list): + self.select_related.update(select) + + def unregister_prefetch(self, prefetch: Union[str, Callable]): + """ + Unregister a prefetch_related operation. + + Args: + prefetch (Union[str, Callable]): The prefetch to remove. Can be a string or a method. + """ + key = self._get_key(prefetch) + self.prefetches.pop(key, None) + + def unregister_children_prefetch(self, prefetch: Union[str, Callable]): + """ + Unregister a children-specific prefetch_related operation. Children as is, the children of a parent. + + Args: + prefetch (Union[str, Callable]): The child prefetch to remove. Can be a string or a method. + """ + key = self._get_key(prefetch) + self.children_prefetches.pop(key, None) + + def unregister_select_related(self, select: str): + """ + Unregister a select_related operation. + + Args: + select (str): The select_related to remove. + """ + self.select_related.discard(select) + + def get_prefetches(self) -> Dict[str, PrefetchType]: + """ + Get all registered prefetch_related operations. + + Returns: + Dict[str, PrefetchType]: A dictionary of prefetch keys to their prefetch. + """ + return self.prefetches + + def get_children_prefetches(self) -> Dict[str, PrefetchType]: + """ + Get all registered children-specific prefetch_related operations. Children as is, the children of a parent. + + Returns: + Dict[str, PrefetchType]: A dictionary of child prefetch keys to their prefetch. + """ + return self.children_prefetches + + def get_select_related(self) -> List[str]: + """ + Get all registered select_related operations. + + Returns: + List[str]: A list of select_related fields. + """ + return list(self.select_related) + + def _get_key(self, operation: Union[PrefetchType, SelectRelatedType]) -> str: + """ + Get the key for an operation. + + Args: + operation (Union[PrefetchType, SelectRelatedType]): The operation to get the key for. + + Returns: + str: The key for the operation. + """ + if isinstance(operation, str): + return operation + elif isinstance(operation, Prefetch): + return operation.prefetch_to + elif callable(operation): + return operation.__name__ + else: + raise ValueError(f"Unsupported operation type: {type(operation)}") + + +prefetch_registry = PrefetchRegistry() diff --git a/oscar_odin/resources/_base.py b/oscar_odin/resources/_base.py index 1b8e42c..f26bc2e 100644 --- a/oscar_odin/resources/_base.py +++ b/oscar_odin/resources/_base.py @@ -5,6 +5,19 @@ class OscarResource(odin.AnnotatedResource, abstract=True): """Base resource for Oscar models.""" + @property + def model_instance(self): + return self._model_instance + + @model_instance.setter + def model_instance(self, value): + self._model_instance = value + + def extra_attrs(self, attrs): + model_instance = attrs.get("model_instance") + if model_instance is not None: + self.model_instance = model_instance + class Meta: allow_field_shadowing = True namespace = "oscar" diff --git a/tests/mappings/test_catalogue.py b/tests/mappings/test_catalogue.py index 8c183a4..7dfa66a 100644 --- a/tests/mappings/test_catalogue.py +++ b/tests/mappings/test_catalogue.py @@ -1,4 +1,7 @@ +from odin.codecs import dict_codec + from django.test import TestCase + from oscar.core.loading import get_model from oscar_odin.mappings import catalogue @@ -41,3 +44,34 @@ def test_mapping__where_is_a_parent_product_include_children(self): self.assertEqual(product.title, actual.title) self.assertIsNotNone(actual.children) self.assertEqual(3, len(actual.children)) + + def test_queryset_to_resources(self): + queryset = Product.objects.all() + product_resources = catalogue.product_queryset_to_resources(queryset) + + self.assertEqual(queryset.count(), len(product_resources)) + + def test_queryset_to_resources_num_queries(self): + queryset = Product.objects.all() + self.assertEqual(queryset.count(), 210) + + # Without all the prefetching, the queries would be 1000+ + # For future reference; It's fine if this test fails after some changes. + # However, the query shouldn't increase too much, if it does, it means you got a + # n+1 query problem and that should be fixed instead by prefetching, annotating etc. + with self.assertNumQueries(14): + resources = catalogue.product_queryset_to_resources( + queryset, include_children=False + ) + dict_codec.dump(resources, include_type_field=False) + + def test_queryset_to_resources_include_children_num_queries(self): + queryset = Product.objects.all() + self.assertEqual(queryset.count(), 210) + + # It should only go up by a few queries. + with self.assertNumQueries(20): + resources = catalogue.product_queryset_to_resources( + queryset, include_children=True + ) + dict_codec.dump(resources, include_type_field=False) diff --git a/tests/mappings/test_prefetch_registry.py b/tests/mappings/test_prefetch_registry.py new file mode 100644 index 0000000..921dd1a --- /dev/null +++ b/tests/mappings/test_prefetch_registry.py @@ -0,0 +1,177 @@ +import unittest +from unittest.mock import Mock, patch, call +from django.db.models import Prefetch +from oscar.core.loading import get_class +from oscar_odin.mappings.prefetching.registry import PrefetchRegistry, prefetch_registry +from oscar_odin.mappings.prefetching.prefetch import prefetch_product_queryset + +ProductQuerySet = get_class("catalogue.managers", "ProductQuerySet") + + +class TestPrefetchSystem(unittest.TestCase): + def setUp(self): + self.registry = PrefetchRegistry() + self.mock_queryset = Mock(spec=ProductQuerySet) + self.mock_queryset.select_related.return_value = self.mock_queryset + self.mock_queryset.prefetch_related.return_value = self.mock_queryset + + def tearDown(self): + # Clear the registry after each test + prefetch_registry.prefetches.clear() + prefetch_registry.children_prefetches.clear() + prefetch_registry.select_related.clear() + + def test_register_string_prefetch(self): + self.registry.register_prefetch("test_prefetch") + self.assertIn("test_prefetch", self.registry.prefetches) + self.assertEqual(self.registry.prefetches["test_prefetch"], "test_prefetch") + + def test_register_prefetch_object(self): + prefetch = Prefetch("test_prefetch") + self.registry.register_prefetch(prefetch) + self.assertIn("test_prefetch", self.registry.prefetches) + self.assertEqual(self.registry.prefetches["test_prefetch"], prefetch) + + def test_register_callable_prefetch(self): + def test_prefetch(queryset, **kwargs): + return queryset + + self.registry.register_prefetch(test_prefetch) + self.assertIn("test_prefetch", self.registry.prefetches) + self.assertEqual(self.registry.prefetches["test_prefetch"], test_prefetch) + + def test_register_children_prefetch(self): + self.registry.register_children_prefetch("test_children_prefetch") + self.assertIn("test_children_prefetch", self.registry.children_prefetches) + + def test_register_select_related_string(self): + self.registry.register_select_related("test_select") + self.assertIn("test_select", self.registry.select_related) + + def test_register_select_related_list(self): + self.registry.register_select_related(["test_select1", "test_select2"]) + self.assertIn("test_select1", self.registry.select_related) + self.assertIn("test_select2", self.registry.select_related) + + def test_unregister_prefetch(self): + self.registry.register_prefetch("test_prefetch") + self.registry.unregister_prefetch("test_prefetch") + self.assertNotIn("test_prefetch", self.registry.prefetches) + + def test_unregister_children_prefetch(self): + self.registry.register_children_prefetch("test_children_prefetch") + self.registry.unregister_children_prefetch("test_children_prefetch") + self.assertNotIn("test_children_prefetch", self.registry.children_prefetches) + + def test_unregister_select_related(self): + self.registry.register_select_related("test_select") + self.registry.unregister_select_related("test_select") + self.assertNotIn("test_select", self.registry.select_related) + + def test_get_prefetches(self): + self.registry.register_prefetch("test_prefetch") + prefetches = self.registry.get_prefetches() + self.assertIn("test_prefetch", prefetches) + + def test_get_children_prefetches(self): + self.registry.register_children_prefetch("test_children_prefetch") + children_prefetches = self.registry.get_children_prefetches() + self.assertIn("test_children_prefetch", children_prefetches) + + def test_get_select_related(self): + self.registry.register_select_related("test_select") + select_related = self.registry.get_select_related() + self.assertIn("test_select", select_related) + + def test_get_key_string(self): + key = self.registry._get_key("test_key") + self.assertEqual(key, "test_key") + + def test_get_key_prefetch_object(self): + prefetch = Prefetch("test_prefetch") + key = self.registry._get_key(prefetch) + self.assertEqual(key, "test_prefetch") + + def test_get_key_callable(self): + def test_callable(): + pass + + key = self.registry._get_key(test_callable) + self.assertEqual(key, "test_callable") + + def test_get_key_unsupported_type(self): + with self.assertRaises(ValueError): + self.registry._get_key(123) + + @patch("oscar_odin.mappings.prefetching.prefetch.prefetch_registry") + def test_prefetch_product_queryset_basic(self, mock_registry): + mock_registry.get_select_related.return_value = ["product_class", "parent"] + mock_registry.get_prefetches.return_value = { + "images": "images", + "stockrecords": "stockrecords", + } + mock_registry.get_children_prefetches.return_value = {} + + result = prefetch_product_queryset(self.mock_queryset) + + self.mock_queryset.select_related.assert_called_once_with( + "product_class", "parent" + ) + self.mock_queryset.prefetch_related.assert_has_calls( + [call("images"), call("stockrecords")], any_order=True + ) + + @patch("oscar_odin.mappings.prefetching.prefetch.prefetch_registry") + def test_prefetch_product_queryset_with_callable(self, mock_registry): + def mock_callable(qs, **kwargs): + return qs.prefetch_related("callable_prefetch") + + mock_registry.get_select_related.return_value = [] + mock_registry.get_prefetches.return_value = {"callable": mock_callable} + mock_registry.get_children_prefetches.return_value = {} + + result = prefetch_product_queryset(self.mock_queryset) + + self.mock_queryset.prefetch_related.assert_called_once_with("callable_prefetch") + + @patch("oscar_odin.mappings.prefetching.prefetch.prefetch_registry") + def test_prefetch_product_queryset_with_children(self, mock_registry): + mock_registry.get_select_related.return_value = [] + mock_registry.get_prefetches.return_value = {} + mock_registry.get_children_prefetches.return_value = { + "children_images": "children__images", + "children_stockrecords": "children__stockrecords", + } + + result = prefetch_product_queryset(self.mock_queryset, include_children=True) + + self.mock_queryset.prefetch_related.assert_has_calls( + [call("children__images"), call("children__stockrecords")], any_order=True + ) + + @patch("oscar_odin.mappings.prefetching.prefetch.prefetch_registry") + def test_prefetch_product_queryset_with_prefetch_object(self, mock_registry): + mock_queryset = Mock(spec=ProductQuerySet) + prefetch_obj = Prefetch("custom_prefetch", queryset=mock_queryset) + mock_registry.get_select_related.return_value = [] + mock_registry.get_prefetches.return_value = {"custom": prefetch_obj} + mock_registry.get_children_prefetches.return_value = {} + + result = prefetch_product_queryset(self.mock_queryset) + + self.mock_queryset.prefetch_related.assert_called_once_with(prefetch_obj) + + @patch("oscar_odin.mappings.prefetching.prefetch.prefetch_registry") + def test_prefetch_product_queryset_with_additional_kwargs(self, mock_registry): + def mock_callable(qs, **kwargs): + if kwargs.get("custom_arg"): + return qs.prefetch_related("custom_prefetch") + return qs + + mock_registry.get_select_related.return_value = [] + mock_registry.get_prefetches.return_value = {"callable": mock_callable} + mock_registry.get_children_prefetches.return_value = {} + + result = prefetch_product_queryset(self.mock_queryset, custom_arg=True) + + self.mock_queryset.prefetch_related.assert_called_once_with("custom_prefetch")