Closes #23563: Replace display attribute prefetching in global search with FETCH_PEERS

netbox/netbox/search/backends.py

@@ -5,8 +5,7 @@ import netaddr
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
 from django.db import DatabaseError, ProgrammingError, transaction
-from django.db.models import F, Q, Window, prefetch_related_objects
-from django.db.models.fields.related import ForeignKey
+from django.db.models import FETCH_PEERS, F, Q, Window
 from django.db.models.functions import window
 from django.utils.module_loading import import_string
 from django.utils.translation import gettext_lazy as _
@@ -15,7 +14,6 @@ from netaddr.core import AddrFormatError
 from core.models import ObjectType
 from extras.models import CachedValue, CustomField
 from netbox.registry import registry
-from utilities.object_types import object_type_identifier
 from utilities.querysets import RestrictedPrefetch
 from utilities.string import title
 
@@ -203,12 +201,6 @@ class CachedValueSearchBackend(SearchBackend):
             )
         )[:MAX_RESULTS]
 
-        # Gather all ObjectTypes present in the search results (used for prefetching related
-        # objects). This must be done before generating the final results list, which returns
-        # a RawQuerySet.
-        object_type_ids = set(queryset.values_list('object_type', flat=True))
-        object_types = ObjectType.objects.filter(pk__in=object_type_ids)
-
         # Construct a Prefetch to pre-fetch only those related objects for which the
         # user has permission to view.
         if user:
@@ -218,33 +210,13 @@ class CachedValueSearchBackend(SearchBackend):
 
         # Wrap the base query to return only the lowest-weight result for each object
         # Hat-tip to https://blog.oyam.dev/django-filter-by-window-function/ for the solution
+        # FETCH_PEERS passes to the prefetched objects, batching the relations read when rendering results
         sql, params = queryset.query.sql_with_params()
-        results = CachedValue.objects.prefetch_related(*prefetch).raw(
+        results = CachedValue.objects.fetch_mode(FETCH_PEERS).prefetch_related(*prefetch).raw(
             f"SELECT * FROM ({sql}) t WHERE row_number = 1",
             params
         )
 
-        # Iterate through each ObjectType represented in the search results and prefetch any
-        # related objects necessary to render the prescribed display attributes (display_attrs).
-        for object_type in object_types:
-            model = object_type.model_class()
-            indexer = registry['search'].get(object_type_identifier(object_type))
-            if not (display_attrs := getattr(indexer, 'display_attrs', None)):
-                continue
-
-            # Add ForeignKey fields to prefetch list
-            prefetch_fields = []
-            for attr in display_attrs:
-                field = model._meta.get_field(attr)
-                if type(field) is ForeignKey:
-                    prefetch_fields.append(f'object__{attr}')
-
-            # Compile a list of all CachedValues referencing this object type, and prefetch
-            # any related objects
-            if prefetch_fields:
-                objects = [r for r in results if r.object_type_id == object_type.pk]
-                prefetch_related_objects(objects, *prefetch_fields)
-
         # Omit any results pertaining to an object the user does not have permission to view
         ret = []
         for r in results:

netbox/netbox/tests/test_search.py

@@ -9,11 +9,13 @@ from django.contrib.contenttypes.models import ContentType
 from django.db import DEFAULT_DB_ALIAS, ProgrammingError, connection, transaction
 from django.db.models.signals import post_delete, post_save
 from django.test import TestCase, TransactionTestCase
+from django.test.utils import CaptureQueriesContext
 from redis.exceptions import ConnectionError as RedisConnectionError
 
-from dcim.models import Manufacturer, Site
+from dcim.models import Device, DeviceRole, DeviceType, Manufacturer, Site
 from dcim.search import SiteIndex
 from extras.models import CachedValue
+from ipam.models import Prefix
 from netbox.search import deferred
 from netbox.search.backends import SearchBackend, search_backend
 from netbox.search.jobs import SearchCacheJob
@@ -193,9 +195,9 @@ class SearchBackendTestCase(TestCase):
         results = search_backend.search('xxxxx')
         self.assertEqual(len(results), 0)
 
-    def test_search_prefetch_display_attrs(self):
+    def test_search_batches_display_attrs(self):
         """
-        Check that ForeignKey fields listed in an indexer's display_attrs are prefetched.
+        Check that a ForeignKey listed in an indexer's display_attrs is fetched for all results with one query.
         """
         tenant = Tenant.objects.create(name='Tenant 1', slug='tenant-1')
         Site.objects.update(tenant=tenant)
@@ -203,8 +205,42 @@ class SearchBackendTestCase(TestCase):
 
         results = search_backend.search('site')
         self.assertEqual(len(results), 3)
-        for result in results:
-            self.assertTrue(Site.tenant.field.is_cached(result.object))
+        with self.assertNumQueries(1):
+            for result in results:
+                self.assertIn(tenant, result.display_attrs.values())
+
+    def test_search_batches_related_objects(self):
+        """
+        Check that the number of queries doesn't grow with the number of results when rendering them reads
+        related objects through __str__() or a GenericForeignKey display attribute.
+        """
+        manufacturer = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
+        device_type = DeviceType.objects.create(manufacturer=manufacturer, model='Device Type 1', slug='device-type-1')
+        role = DeviceRole.objects.create(name='Device Role 1', slug='device-role-1')
+
+        def render_results(sites):
+            # Unnamed devices render as "<manufacturer> <model> (<pk>)"; prefixes display their scope (a site)
+            for site in sites:
+                Device.objects.create(site=site, device_type=device_type, role=role, description='Batched')
+                Prefix.objects.create(prefix=f'10.0.{site.pk}.0/24', scope=site, description='Batched')
+            search_backend.cache(Device.objects.all())
+            search_backend.cache(Prefix.objects.all())
+
+            ContentType.objects.clear_cache()
+            with CaptureQueriesContext(connection) as queries:
+                results = search_backend.search('batched')
+                rendered = [(result.name, list(result.display_attrs.values())) for result in results]
+            return rendered, len(queries)
+
+        sites = list(Site.objects.order_by('name'))
+        rendered, query_count = render_results(sites[:1])
+        self.assertEqual(len(rendered), 2)
+        rendered, more_query_count = render_results(sites[1:])
+        self.assertEqual(len(rendered), 6)
+        self.assertEqual(more_query_count, query_count)
+
+        self.assertEqual(sum(name.startswith('Manufacturer 1 Device Type 1') for name, _ in rendered), 3)
+        self.assertCountEqual([attrs[0] for name, attrs in rendered if name.startswith('10.0.')], sites)
 
 
 class DeferredCachingTestCase(RQQueueTestMixin, TestCase):