Fix related hosts

allegro · Nov 26, 2024 · 68030f3 · 68030f3
1 parent 6271315
commit 68030f3
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 26 deletions.
diff --git a/src/ralph/data_center/api/serializers.py b/src/ralph/data_center/api/serializers.py
@@ -1,6 +1,5 @@
 # -*- coding: utf-8 -*-
 
-from django.contrib.contenttypes.models import ContentType
 from rest_framework import serializers
 
 from ralph.api import RalphAPISerializer
@@ -27,7 +26,6 @@
     VIP
 )
 from ralph.security.api import SecurityScanSerializer
-from ralph.virtual.models import VirtualServer
 
 
 class ClusterTypeSerializer(RalphAPISerializer):
@@ -132,40 +130,32 @@ class DataCenterAssetSerializer(ComponentSerializerMixin, AssetSerializer):
     securityscan = SecurityScanSerializer()
     related_hosts = serializers.SerializerMethodField()
 
-    def _get_serialized_sublist(self, full_list, serializer_class, cond):
+    def _get_serialized_sublist(self, full_list, serializer_class):
         return serializer_class(
-            [elem for elem in full_list if cond(elem)],
+            full_list,
             many=True, context=self.context
         ).data
 
     def _get_physical_servers(self, obj):
-        dca = ContentType.objects.get_for_model(DataCenterAsset)
         return self._get_serialized_sublist(
-            obj.children.all(),
-            DataCenterAssetSimpleSerializer,
-            lambda child: child.content_type == dca
+            obj.physical_servers,
+            DataCenterAssetSimpleSerializer
         )
 
     def _get_virtual_servers(self, obj):
         from ralph.virtual.api import VirtualServerSimpleSerializer
-        vs = ContentType.objects.get_for_model(VirtualServer)
         return self._get_serialized_sublist(
-            obj.children.all(),
-            VirtualServerSimpleSerializer,
-            lambda child: child.content_type == vs
+            obj.virtual_servers,
+            VirtualServerSimpleSerializer
         )
 
     def _get_cloud_hosts(self, obj):
         from ralph.virtual.api import CloudHostSimpleSerializer
         return self._get_serialized_sublist(
-            obj.cloudhost_set.all(),
-            CloudHostSimpleSerializer,
-            lambda host: True
-        )
+            obj.cloud_hosts,
+            CloudHostSimpleSerializer)
 
     def get_related_hosts(self, obj):
-        if hasattr(obj, 'content_type'):
-            obj = obj.content_type.get_object_for_this_type(pk=obj.pk)
         return {
             "virtual_servers": self._get_virtual_servers(obj),
             "physical_servers": self._get_physical_servers(obj),

diff --git a/src/ralph/data_center/api/views.py b/src/ralph/data_center/api/views.py
@@ -42,6 +42,7 @@
     ServerRoom,
     VIP
 )
+from ralph.virtual.models import CloudHost, VirtualServer
 
 
 class DataCenterAssetFilterSet(NetworkableObjectFilters):
@@ -50,7 +51,7 @@ class Meta(NetworkableObjectFilters.Meta):
 
 
 class DataCenterAssetViewSet(BaseObjectViewSetMixin, RalphAPIViewSet):
-    queryset = DataCenterAsset.polymorphic_objects.all()
+    queryset = DataCenterAsset.objects.all()
     serializer_class = DataCenterAssetSerializer
     save_serializer_class = DataCenterAssetSaveSerializer
     select_related = DataCenterAssetAdmin.list_select_related + [
@@ -62,12 +63,25 @@ class DataCenterAssetViewSet(BaseObjectViewSetMixin, RalphAPIViewSet):
         'asset_ptr',
     ]
     prefetch_related = base_object_descendant_prefetch_related + [
-        'children',
+        Prefetch(
+            'children',
+            queryset=VirtualServer.objects.select_related('parent'),
+            to_attr='virtual_servers'
+        ),
+        Prefetch(
+            'children',
+            queryset=DataCenterAsset.objects.select_related('parent'),
+            to_attr='physical_servers'
+        ),
         'rack__server_room__data_center',
         'connections',
         'tags',
         'memory_set',
-        'cloudhost_set',
+        Prefetch(
+            'cloudhost_set',
+            queryset=CloudHost.objects.select_related('parent'),
+            to_attr='cloud_hosts'
+        ),
         Prefetch(
             'ethernet_set',
             queryset=Ethernet.objects.select_related('ipaddress')
@@ -95,11 +109,8 @@ def get_queryset(self):
             ConfigurationModule,
             ServiceEnvironment
         )
-        return (
-            DataCenterAsset.polymorphic_objects
-            .select_related(*self.select_related)
-            .polymorphic_prefetch_related(DataCenterAsset=self.prefetch_related)
-        )
+        qs = super().get_queryset()
+        return qs
 
 
 class AccessoryViewSet(RalphAPIViewSet):