Move YAML export prefetching to the models and drop get_export_queryset()

YAMLExporter now applies the prefetching named by a model's yaml_prefetch
attribute, along with the FETCH_PEERS fetch mode, so ObjectListView no
longer needs an export queryset hook.

docs/plugins/development/views.md

@@ -154,7 +154,6 @@ Below are the class definitions for NetBox's multi-object views. These views han
     options:
       members:
         - get_table
-        - get_export_queryset
         - export_table
         - export_template
 

netbox/dcim/models/devices.py

@@ -220,6 +220,8 @@ class DeviceType(ImageAttachmentsMixin, PrimaryModel, WeightMixin):
     prerequisite_models = (
         'dcim.Manufacturer',
     )
+    # Avoid one module_bay_types query per module bay template when exporting to YAML
+    yaml_prefetch = ('modulebaytemplates__module_bay_types',)
 
     class Meta:
         ordering = ['manufacturer', 'model']

netbox/dcim/models/modules.py

@@ -227,6 +227,8 @@ class ModuleType(ImageAttachmentsMixin, PrimaryModel, WeightMixin):
     prerequisite_models = (
         'dcim.Manufacturer',
     )
+    # Avoid one module_bay_types query per module type when exporting to YAML
+    yaml_prefetch = ('module_bay_types',)
 
     class Meta:
         ordering = ('profile', 'manufacturer', 'model')

netbox/dcim/tests/test_views.py

@@ -11,6 +11,7 @@ import yaml
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.messages import get_messages
 from django.db import connection
+from django.db.models import FETCH_PEERS
 from django.http import QueryDict, StreamingHttpResponse
 from django.test import override_settings, tag
 from django.test.utils import CaptureQueriesContext
@@ -22,7 +23,6 @@ from core.models import ObjectChange, ObjectType
 from dcim.choices import *
 from dcim.constants import *
 from dcim.models import *
-from dcim.views import DeviceTypeListView, ModuleTypeListView
 from extras.models import ConfigContext, ConfigTemplate
 from ipam.models import ASN, RIR, VLAN, VRF, IPAddress, VLANGroup
 from netbox.choices import (
@@ -1288,15 +1288,12 @@ inventory-items:
         one_bay_device_type = make_device_type('Export Query DT One Bay', 1)
         five_bay_device_type = make_device_type('Export Query DT Five Bays', 5)
 
-        view = DeviceTypeListView()
         exporter = get_exporter('yaml')
-        view.queryset = DeviceType.objects.filter(pk=one_bay_device_type.pk)
         with CaptureQueriesContext(connection) as one_bay_queries:
-            exporter.export(request=None, queryset=view.get_export_queryset(exporter))
+            exporter.export(request=None, queryset=DeviceType.objects.filter(pk=one_bay_device_type.pk))
 
-        view.queryset = DeviceType.objects.filter(pk=five_bay_device_type.pk)
         with CaptureQueriesContext(connection) as five_bay_queries:
-            exporter.export(request=None, queryset=view.get_export_queryset(exporter))
+            exporter.export(request=None, queryset=DeviceType.objects.filter(pk=five_bay_device_type.pk))
 
         self.assertEqual(len(one_bay_queries), len(five_bay_queries))
 
@@ -2187,10 +2184,10 @@ interfaces:
             self.assertEqual(channel.parent, imported_parent)
 
     def test_bulk_yaml_export_prefetches_module_bay_types_on_the_module_type_itself(self):
-        """Compares an unprefetched to_yaml() call per instance against a YAML export (which
-        prefetches and issues no queries of its own beyond that), rather than a row-count
-        comparison, which other per-instance relations that legitimately scale with it would
-        swamp."""
+        """Compares an unprefetched to_yaml() call per instance (using the same fetch mode) against
+        a YAML export (which prefetches and issues no queries of its own beyond that), rather than
+        a row-count comparison, which other per-instance relations that legitimately scale with it
+        would swamp."""
         manufacturer = Manufacturer.objects.create(name='Export Query MT Manufacturer', slug='export-query-mt-mfr')
         bay_type = ModuleBayType.objects.create(name='Export Query MT SFP28', slug='export-query-mt-sfp28')
 
@@ -2202,13 +2199,10 @@ interfaces:
         pks = [mt.pk for mt in module_types]
 
         with CaptureQueriesContext(connection) as unprefetched:
-            [obj.to_yaml() for obj in ModuleType.objects.filter(pk__in=pks)]
+            [obj.to_yaml() for obj in ModuleType.objects.filter(pk__in=pks).fetch_mode(FETCH_PEERS)]
 
-        view = ModuleTypeListView()
-        exporter = get_exporter('yaml')
-        view.queryset = ModuleType.objects.filter(pk__in=pks)
         with CaptureQueriesContext(connection) as prefetched:
-            exporter.export(request=None, queryset=view.get_export_queryset(exporter))
+            get_exporter('yaml').export(request=None, queryset=ModuleType.objects.filter(pk__in=pks))
 
         # Without the prefetch, each of the 5 module types issues its own module_bay_types
         # query; with it, exactly one query serves all 5.

netbox/dcim/views.py

@@ -20,7 +20,6 @@ from extras.views import ObjectConfigContextView, ObjectRenderConfigView
 from ipam.models import ASN, VLAN, IPAddress, Prefix, VLANGroup
 from ipam.tables import VLANTranslationRuleTable
 from ipam.ui.panels import FHRPGroupAssignmentsPanel
-from netbox.exporters import YAMLExporter
 from netbox.object_actions import *
 from netbox.ui import actions, layout
 from netbox.ui.breadcrumbs import Breadcrumb, filtered_list_url, object_view_url
@@ -1450,13 +1449,6 @@ class DeviceTypeListView(generic.ObjectListView):
     filterset_form = forms.DeviceTypeFilterForm
     table = tables.DeviceTypeTable
 
-    def get_export_queryset(self, exporter):
-        queryset = super().get_export_queryset(exporter)
-        if isinstance(exporter, YAMLExporter):
-            # Avoid one module_bay_types query per module bay template across the export.
-            queryset = queryset.prefetch_related('modulebaytemplates__module_bay_types')
-        return queryset
-
 
 @register_model_view(DeviceType)
 class DeviceTypeView(GetRelatedModelsMixin, generic.ObjectView):
@@ -1917,15 +1909,6 @@ class ModuleTypeListView(generic.ObjectListView):
     filterset_form = forms.ModuleTypeFilterForm
     table = tables.ModuleTypeTable
 
-    def get_export_queryset(self, exporter):
-        queryset = super().get_export_queryset(exporter)
-        if isinstance(exporter, YAMLExporter):
-            # Avoid one module_bay_types query per module type across the export. (Unlike
-            # DeviceType.to_yaml(), ModuleType.to_yaml() doesn't export module bay templates at
-            # all, so there's nothing to prefetch alongside it.)
-            queryset = queryset.prefetch_related('module_bay_types')
-        return queryset
-
 
 @register_model_view(ModuleType)
 class ModuleTypeView(GetRelatedModelsMixin, generic.ObjectView):

netbox/netbox/exporters.py

@@ -4,6 +4,7 @@ from itertools import chain
 
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
+from django.db.models import FETCH_PEERS
 from django.utils.translation import gettext as _
 
 from netbox.registry import registry
@@ -125,6 +126,8 @@ class YAMLExporter(Exporter):
         return hasattr(model, 'to_yaml')
 
     def export(self, *, request, queryset):
+        # Prefetch the relations named by the model, and batch any other foreign key lookups across all objects
+        queryset = queryset.prefetch_related(*getattr(queryset.model, 'yaml_prefetch', ())).fetch_mode(FETCH_PEERS)
         return '---\n'.join(obj.to_yaml() for obj in queryset)
 
 

netbox/netbox/tests/test_exporters.py

@@ -4,8 +4,10 @@ from unittest import skipIf
 
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
+from django.db import connection
 from django.http import StreamingHttpResponse
 from django.test import RequestFactory, SimpleTestCase, override_settings
+from django.test.utils import CaptureQueriesContext
 from django.urls import reverse
 
 from dcim.models import DeviceType, Manufacturer, Site
@@ -110,6 +112,22 @@ class CSVExporterTestCase(SimpleTestCase):
         self.assertEqual(list(data), ['Name,Count\r\n', 'Site 1,1\r\n', '"Site, 2",\r\n'])
 
 
+class YAMLExporterTestCase(TestCase):
+
+    def test_export_fetches_peers(self):
+        """
+        Related objects are fetched for all exported objects at once, rather than one by one.
+        """
+        for i in range(1, 4):
+            manufacturer = Manufacturer.objects.create(name=f'Manufacturer {i}', slug=f'manufacturer-{i}')
+            DeviceType.objects.create(manufacturer=manufacturer, model=f'Device Type {i}', slug=f'device-type-{i}')
+
+        with CaptureQueriesContext(connection) as queries:
+            get_exporter('yaml').export(request=None, queryset=DeviceType.objects.all())
+        manufacturer_queries = [q for q in queries.captured_queries if 'FROM "dcim_manufacturer"' in q['sql']]
+        self.assertEqual(len(manufacturer_queries), 1)
+
+
 class ExportViewTestCase(TestCase):
 
     @classmethod

netbox/netbox/views/generic/bulk_views.py

@@ -112,15 +112,6 @@ class ObjectListView(BaseMultiObjectView, ActionsMixin, TableMixin):
     # Export methods
     #
 
-    def get_export_queryset(self, exporter):
-        """
-        Return the queryset of objects to be exported using the given exporter.
-
-        Args:
-            exporter: The Exporter instance
-        """
-        return self.queryset
-
     def get_export_rows(self, table, columns=None):
         """
         Return the column headers and an iterator of cell values for each row to be exported from the table.
@@ -180,8 +171,7 @@ class ObjectListView(BaseMultiObjectView, ActionsMixin, TableMixin):
         """
         Deprecated: Export the queryset of objects as concatenated YAML documents.
         """
-        exporter = get_exporter('yaml')
-        return exporter.export(request=self.request, queryset=self.get_export_queryset(exporter))
+        return get_exporter('yaml').export(request=self.request, queryset=self.queryset)
 
     # TODO: Remove in NetBox v5.0
     def export_table(self, table, columns=None, filename=None, delimiter=None):
@@ -273,12 +263,11 @@ class ObjectListView(BaseMultiObjectView, ActionsMixin, TableMixin):
                 exporter = get_exporter('yaml')
             else:
                 exporter = get_exporter('csv')
-            queryset = self.get_export_queryset(exporter)
             if not isinstance(exporter, TableExporter):
-                data = exporter.export(request=request, queryset=queryset)
+                data = exporter.export(request=request, queryset=self.queryset)
                 return self.get_export_response(exporter, data)
 
-            table = self.get_table(queryset, request, has_table_actions)
+            table = self.get_table(self.queryset, request, has_table_actions)
             columns = [name for name, _ in table.selected_columns] if current_view else None
 
             # TODO: Remove in NetBox v5.0 along with export_table()