From a1238b3cbe0227d966e788bc59e56e2b4b1a2a05 Mon Sep 17 00:00:00 2001
From: Oliver Walters <oliver.henry.walters@gmail.com>
Date: Sun, 30 Nov 2025 02:06:55 +0000
Subject: [PATCH] Validate queryset annotation

- Add unit test with large dataset
- Ensure number of queries is fixed
- Fix for prefetching check
---
 src/backend/InvenTree/InvenTree/models.py | 20 +++---
 src/backend/InvenTree/common/test_api.py  | 78 +++++++++++++++++++++++
 2 files changed, 87 insertions(+), 11 deletions(-)

diff --git a/src/backend/InvenTree/InvenTree/models.py b/src/backend/InvenTree/InvenTree/models.py
index c8e8eb4ecf..fdb22675e1 100644
--- a/src/backend/InvenTree/InvenTree/models.py
+++ b/src/backend/InvenTree/InvenTree/models.py
@@ -454,11 +454,12 @@ class ReferenceIndexingMixin(models.Model):
 class ContentTypeMixin:
     """Mixin class which supports retrieval of the ContentType for a model instance."""
 
-    def get_content_type(self):
+    @classmethod
+    def get_content_type(cls):
         """Return the ContentType object associated with this model."""
         from django.contrib.contenttypes.models import ContentType
 
-        return ContentType.objects.get_for_model(self.__class__)
+        return ContentType.objects.get_for_model(cls)
 
 
 class InvenTreeModel(ContentTypeMixin, PluginValidationMixin, models.Model):
@@ -532,14 +533,10 @@ class InvenTreeParameterMixin(InvenTreePermissionCheckMixin, models.Model):
         Returns:
             Annotated queryset
         """
-        from common.models import Parameter
-
         return queryset.prefetch_related(
-            models.Prefetch(
-                'parameters_list',
-                queryset=Parameter.objects.all().select_related('template'),
-                to_attr='parameters_list_prefetched',
-            )
+            'parameters_list',
+            'parameters_list__model_type',
+            'parameters_list__template',
         )
 
     @property
@@ -548,8 +545,9 @@ class InvenTreeParameterMixin(InvenTreePermissionCheckMixin, models.Model):
 
         This will return pre-fetched data if available (i.e. in a serializer context).
         """
-        if hasattr(self, 'parameters_list_prefetched'):
-            return self.parameters_list_prefetched
+        # Check the query cache for pre-fetched parameters
+        if 'parameters_list' in getattr(self, '_prefetched_objects_cache', {}):
+            return self._prefetched_objects_cache['parameters_list']
 
         return self.parameters_list.all()
 
diff --git a/src/backend/InvenTree/common/test_api.py b/src/backend/InvenTree/common/test_api.py
index d5a465f8c8..048ae4c92c 100644
--- a/src/backend/InvenTree/common/test_api.py
+++ b/src/backend/InvenTree/common/test_api.py
@@ -192,3 +192,81 @@ class ParameterAPITests(InvenTreeAPITestCase):
         self.assertFalse(
             common.models.Parameter.objects.filter(pk=parameter.pk).exists()
         )
+
+    def test_parameter_annotation(self):
+        """Test that we can annotate parameters against a queryset."""
+        from company.models import Company
+
+        templates = []
+        parameters = []
+        companies = []
+
+        for ii in range(100):
+            company = Company(
+                name=f'Test Company {ii}',
+                description='A company for testing parameter annotations',
+            )
+            companies.append(company)
+
+        Company.objects.bulk_create(companies)
+
+        # Let's create a large number of parameters
+        for ii in range(25):
+            templates.append(
+                common.models.ParameterTemplate(
+                    name=f'Test Parameter {ii}',
+                    units='',
+                    description='A parameter for testing annotations',
+                    model_type=Company.get_content_type(),
+                    enabled=True,
+                )
+            )
+
+        common.models.ParameterTemplate.objects.bulk_create(templates)
+
+        # Create a parameter for every company against every template
+        for company in companies:
+            for template in templates:
+                parameters.append(
+                    common.models.Parameter(
+                        template=template,
+                        model_type=company.get_content_type(),
+                        model_id=company.pk,
+                        data=f'Test data for {company.name} - {template.name}',
+                    )
+                )
+
+        common.models.Parameter.objects.bulk_create(parameters)
+
+        self.assertEqual(
+            common.models.Parameter.objects.count(), len(companies) * len(templates)
+        )
+
+        # We will fetch the companies, annotated with all parameters
+        url = reverse('api-company-list')
+
+        # By default, we do not expect any parameter annotations
+        response = self.get(url, data={'limit': 5})
+
+        self.assertEqual(response.data['count'], len(companies))
+        for company in response.data['results']:
+            self.assertNotIn('parameters', company)
+
+        # Fetch all companies, explicitly without parameters
+        with self.assertNumQueriesLessThan(20):
+            response = self.get(url, data={'parameters': False})
+
+        # Now, annotate with parameters
+        # This must be done efficiently, without an 1 + N query pattern
+        with self.assertNumQueriesLessThan(45):
+            response = self.get(url, data={'parameters': True})
+
+        self.assertEqual(len(response.data), len(companies))
+
+        for company in response.data:
+            self.assertIn('parameters', company)
+            self.assertEqual(
+                len(company['parameters']),
+                len(templates),
+                'Incorrect number of parameter annotations found',
+            )