Merge pull request #3869 from allegro/fix_related_hosts

Fix related hosts

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):
@@ -122,7 +120,7 @@ class Meta(SimpleRackSerializer.Meta):
 class DataCenterAssetSimpleSerializer(RalphAPISerializer):
     class Meta:
         model = DataCenterAsset
-        fields = ['hostname', 'url']
+        fields = ['id', 'hostname', 'url']
         _skip_tags_field = True
 
 
@@ -132,42 +130,21 @@ class DataCenterAssetSerializer(ComponentSerializerMixin, AssetSerializer):
     securityscan = SecurityScanSerializer()
     related_hosts = serializers.SerializerMethodField()
 
-    def _get_serialized_sublist(self, full_list, serializer_class, cond):
-        return serializer_class(
-            [elem for elem in full_list if cond(elem)],
-            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
-        )
-
-    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
-        )
-
-    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
-        )
-
     def get_related_hosts(self, obj):
+        from ralph.virtual.api import CloudHostSimpleSerializer
+        from ralph.virtual.api import VirtualServerSimpleSerializer
+        # attributes "virtual_servers", "physical_servers" and "cloud_hosts"
+        # are custom prefetches, see DataCenterAssetViewSet
         return {
-            "virtual_servers": self._get_virtual_servers(obj),
-            "physical_servers": self._get_physical_servers(obj),
-            "cloud_hosts": self._get_cloud_hosts(obj)
+            "virtual_servers": VirtualServerSimpleSerializer(
+                obj.virtual_servers, many=True, context=self.context
+            ).data,
+            "physical_servers": DataCenterAssetSimpleSerializer(
+                obj.physical_servers, many=True, context=self.context
+            ).data,
+            "cloud_hosts": CloudHostSimpleSerializer(
+                obj.cloud_hosts, many=True, context=self.context
+            ).data,
         }
 
     class Meta(AssetSerializer.Meta):

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):

src/ralph/data_center/tests/test_api.py

@@ -132,6 +132,51 @@ def test_get_data_center_asset_details_related_hosts(self):
             response.data['related_hosts']['physical_servers'][0]['hostname'],
             dc_asset_4.hostname
         )
+    def test_get_data_center_asset_list_related_hosts(self):
+        dc_asset_5 = DataCenterAssetFullFactory()
+        cloud_host = CloudHostFactory(
+            hypervisor=dc_asset_5
+        )
+        virtual_server = VirtualServerFactory(
+            parent=dc_asset_5
+        )
+        virtual_server_2 = VirtualServerFactory(
+            parent=dc_asset_5,
+            hostname='random_test_hostname'
+        )
+        dc_asset_6 = DataCenterAssetFullFactory(
+            parent=dc_asset_5
+        )
+        url = f"{reverse('datacenterasset-list')}?hostname={dc_asset_5.hostname}"
+        response = self.client.get(url, format='json')
+        self.assertEqual(
+            len(response.data['results']), 1
+        )
+        self.assertEqual(
+            len(response.data['results'][0]['related_hosts']['cloud_hosts']), 1
+        )
+        self.assertEqual(
+            len(response.data['results'][0]['related_hosts']['virtual_servers']), 2
+        )
+        self.assertIn(
+            response.data['results'][0]['related_hosts']['virtual_servers'][0]['hostname'],
+            (virtual_server.hostname, virtual_server_2.hostname)
+        )
+        self.assertIn(
+            response.data['results'][0]['related_hosts']['virtual_servers'][1]['hostname'],
+            (virtual_server.hostname, virtual_server_2.hostname)
+        )
+        self.assertEqual(
+            response.data['results'][0]['related_hosts']['cloud_hosts'][0]['hostname'],
+            cloud_host.hostname
+        )
+        self.assertEqual(
+            len(response.data['results'][0]['related_hosts']['physical_servers']), 1
+        )
+        self.assertEqual(
+            response.data['results'][0]['related_hosts']['physical_servers'][0]['hostname'],
+            dc_asset_6.hostname
+        )
 
     def test_create_data_center_asset(self):
         url = reverse('datacenterasset-list')

src/ralph/virtual/api.py

@@ -49,7 +49,7 @@ class CloudHostSimpleSerializer(BaseObjectSerializer):
 
     class Meta:
         model = CloudHost
-        fields = ['hostname', 'ip_addresses', 'url']
+        fields = ['id', 'hostname', 'ip_addresses', 'url']
 
 
 class CloudProjectSimpleSerializer(BaseObjectSerializer):
@@ -152,7 +152,7 @@ class Meta:
 class VirtualServerSimpleSerializer(BaseObjectSerializer):
     class Meta(BaseObjectSerializer.Meta):
         model = VirtualServer
-        fields = ['hostname', 'url']
+        fields = ['id', 'hostname', 'url']
         _skip_tags_field = True
         exclude = None