From b09ad859788ea70e46afe22f20c79c80a8de377f Mon Sep 17 00:00:00 2001
From: Aaron Tan <jangsutsr@gmail.com>
Date: Fri, 9 Jun 2017 13:47:05 -0400
Subject: [PATCH] Add count to host group summary field.

---
 awx/api/serializers.py     | 13 +++++++++----
 awx/main/utils/__init__.py |  4 ++--
 awx/main/utils/common.py   |  6 ++++++
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/awx/api/serializers.py b/awx/api/serializers.py
index 13efa6a34b..a36e4c2b86 100644
--- a/awx/api/serializers.py
+++ b/awx/api/serializers.py
@@ -44,7 +44,8 @@ from awx.main.access import get_user_capabilities
 from awx.main.fields import ImplicitRoleField
 from awx.main.utils import (
     get_type_for_model, get_model_for_type, timestamp_apiformat,
-    camelcase_to_underscore, getattrd, parse_yaml_or_json)
+    camelcase_to_underscore, getattrd, parse_yaml_or_json,
+    has_model_field_prefetched)
 from awx.main.utils.filters import SmartFilter
 
 from awx.main.validators import vars_validate_or_raise
@@ -1247,8 +1248,12 @@ class HostSerializer(BaseSerializerWithVariables):
             d['last_job']['job_template_name'] = obj.last_job.job_template.name
         except (KeyError, AttributeError):
             pass
-        group_list = [{'id': g.id, 'name': g.name} for g in obj.groups.order_by('name')[:5]]
-        d.setdefault('groups', group_list)
+        if has_model_field_prefetched(obj, 'groups'):
+            group_list = sorted([{'id': g.id, 'name': g.name} for g in obj.groups.all()], key=lambda x: x['name'])[:5]
+        else:
+            group_list = [{'id': g.id, 'name': g.name} for g in obj.groups.all().order_by('name')[:5]]
+        group_cnt = obj.groups.count()
+        d.setdefault('groups', {'count': group_cnt, 'results': group_list})
         d.setdefault('recent_jobs', [{
             'id': j.job.id,
             'name': j.job.job_template.name if j.job.job_template is not None else "",
@@ -2179,7 +2184,7 @@ class OrganizationCredentialSerializerCreate(CredentialSerializerCreate):
 class LabelsListMixin(object):
 
     def _summary_field_labels(self, obj):
-        if hasattr(obj, '_prefetched_objects_cache') and obj.labels.prefetch_cache_name in obj._prefetched_objects_cache:
+        if has_model_field_prefetched(obj, 'labels'):
             label_list = [{'id': x.id, 'name': x.name} for x in obj.labels.all()[:10]]
             label_ct = len(obj.labels.all())
         else:
diff --git a/awx/main/utils/__init__.py b/awx/main/utils/__init__.py
index 360f5a65f2..fbdb9d11e0 100644
--- a/awx/main/utils/__init__.py
+++ b/awx/main/utils/__init__.py
@@ -19,6 +19,6 @@ from awx.main.utils.common import ( # noqa
     build_proot_temp_dir,
     wrap_args_with_proot,
     get_system_task_capacity,
-    decrypt_field_value
+    decrypt_field_value,
+    has_model_field_prefetched
 )
-
diff --git a/awx/main/utils/common.py b/awx/main/utils/common.py
index 20b49fcdcf..dd42078d43 100644
--- a/awx/main/utils/common.py
+++ b/awx/main/utils/common.py
@@ -882,3 +882,9 @@ def get_search_fields(model):
                           'name', 'description'):
             fields.append(field.name)
     return fields
+
+
+def has_model_field_prefetched(model_obj, field_name):
+    # NOTE: Update this function if django internal implementation changes.
+    return getattr(getattr(model_obj, field_name, None),
+                   'prefetch_cache_name', '') in getattr(model_obj, '_prefetched_objects_cache', {})