Make prefetches work for filtered and not filtered export

allegro · Oct 31, 2024 · 04002d9 · 04002d9
1 parent b198817
commit 04002d9
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 7 deletions.
diff --git a/src/ralph/data_center/admin.py b/src/ralph/data_center/admin.py
@@ -35,7 +35,7 @@
 from ralph.admin.views.main import RalphChangeList
 from ralph.admin.views.multiadd import MulitiAddAdminMixin
 from ralph.assets.invoice_report import AssetInvoiceReportMixin
-from ralph.assets.models.base import BaseObject
+from ralph.assets.models.base import BaseObject, BaseObjectPolymorphicQuerySet
 from ralph.assets.models.components import Ethernet
 from ralph.assets.views import ComponentsAdminView
 from ralph.attachments.admin import AttachmentsMixin
@@ -492,13 +492,25 @@ class DataCenterAssetAdmin(
     )
 
     def get_export_queryset(self, request):
-        # we skip method in RalphAdminImportExportMixin because it doesn't return a queryset
-        # sad :(
-        return super(RalphAdminImportExportMixin, self).get_export_queryset(request).select_related(
-            *self.list_select_related
-        ).polymorphic_prefetch_related(
-            DataCenterAsset=['tags', 'ethernet_set__ipaddress', 'parent__ethernet_set__ipaddress'],
+        qs = (
+            super(RalphAdminImportExportMixin, self)
+            .get_export_queryset(request)
+            .select_related(
+                *self.list_select_related
+            )
         )
+        if isinstance(qs, BaseObjectPolymorphicQuerySet):
+            return qs.polymorphic_prefetch_related(
+                DataCenterAsset=[
+                    'tags',
+                    'ethernet_set__ipaddress',
+                    'parent__ethernet_set__ipaddress'
+                ]
+            )
+        else:
+            return qs.prefetch_related(
+                'tags', 'ethernet_set__ipaddress', 'parent__ethernet_set__ipaddress'
+            )
 
     def get_multiadd_fields(self, obj=None):
         multiadd_fields = [

diff --git a/src/ralph/data_importer/tests/test_export.py b/src/ralph/data_importer/tests/test_export.py
@@ -114,6 +114,16 @@ def test_data_center_asset_export_queries_count(self):
             DataCenterAsset
         ), max_queries=12)
 
+    def test_data_center_asset_export_filtered(self):
+        self._init(10)
+        first_id = next(iter(self.data_center_assets_map.keys()))
+        with CaptureQueriesContext(connections['default']) as cqc:
+            export_data = self._export(
+                DataCenterAsset, filters={'id': first_id}
+            )
+            queries = len(cqc)
+        self.assertEqual(len(export_data.dict), 1)
+        self.assertLessEqual(queries, 10)
 
 class DataCenterAssetExporterTestCaseWithParent(DataCenterAssetExporterTestCase):
     def _init(self, num=10):