Просмотр исходного кода

Fixes #21975: Prefetch all related data during CSV bulk export (#21976)

Jeremy Stretch 1 месяц назад
Родитель
Сommit
789085cc33

+ 11 - 4
netbox/netbox/tables/tables.py

@@ -116,16 +116,23 @@ class BaseTable(tables.Table):
             self.sequence.remove('actions')
             self.sequence.append('actions')
 
-    def _apply_prefetching(self):
+    def _apply_prefetching(self, columns=None):
         """
-        Dynamically update the table's QuerySet to ensure related fields are pre-fetched
+        Dynamically update the table's QuerySet to ensure related fields are pre-fetched.
+
+        Args:
+            columns: An optional iterable of column names for which to apply prefetching,
+                regardless of visibility. If None, only currently visible columns are used.
         """
         if not isinstance(self.data, TableQuerysetData):
             return
 
         prefetch_fields = []
-        for column in self.columns:
-            if not column.visible:
+        for column in self.columns.iterall():
+            if columns is not None:
+                if column.name not in columns:
+                    continue
+            elif not column.visible:
                 # Skip hidden columns
                 continue
             model = getattr(self.Meta, 'model')  # Must be called *after* resolving columns

+ 32 - 0
netbox/netbox/tests/test_tables.py

@@ -47,6 +47,38 @@ class BaseTableTest(TestCase):
         prefetch_lookups = table.data.data._prefetch_related_lookups
         self.assertEqual(prefetch_lookups, tuple())
 
+    def test_prefetch_all_columns_for_export(self):
+        """
+        Verify that related fields for *all* table columns are prefetched when preparing a CSV
+        export, including columns which are not currently visible in the user's configured view.
+        """
+        request = RequestFactory().get('/')
+        request.user = self.user
+
+        # Configure the table with only local-field columns visible. Related columns like 'site',
+        # 'rack', and 'region' are hidden in the user's view.
+        self.user.config.set(
+            'tables.DeviceTable.columns',
+            ['name', 'status'],
+            commit=True,
+        )
+        table = DeviceTable(Device.objects.all())
+        table.configure(request)
+
+        # With only local-field columns visible, no relations should be prefetched yet.
+        self.assertEqual(table.data.data._prefetch_related_lookups, tuple())
+
+        # Simulate the CSV "All data" export path: re-apply prefetching for every column that
+        # will be included in the export, regardless of visibility.
+        export_columns = [
+            col_name for col_name, _ in table.selected_columns + table.available_columns
+        ]
+        table._apply_prefetching(columns=export_columns)
+
+        prefetch_lookups = table.data.data._prefetch_related_lookups
+        self.assertIn('rack', prefetch_lookups)
+        self.assertIn('site__region', prefetch_lookups)
+
     def test_configure_anonymous_user_with_ordering(self):
         """
         Verify that table.configure() does not raise an error when an anonymous

+ 6 - 1
netbox/netbox/views/generic/bulk_views.py

@@ -93,11 +93,16 @@ class ObjectListView(BaseMultiObjectView, ActionsMixin, TableMixin):
             delimiter: The character used to separate columns (a comma is used by default)
         """
         exclude_columns = {'pk', 'actions'}
+        all_columns = [col_name for col_name, _ in table.selected_columns + table.available_columns]
         if columns:
-            all_columns = [col_name for col_name, _ in table.selected_columns + table.available_columns]
             exclude_columns.update({
                 col for col in all_columns if col not in columns
             })
+
+        # Ensure related objects are prefetched for every column that will be exported, not just
+        # those currently visible in the configured table view.
+        table._apply_prefetching(columns=[c for c in all_columns if c not in exclude_columns])
+
         exporter = TableExport(
             export_format=TableExport.CSV,
             table=table,