custom attributes prefetching (#29531)

src/authentic2/manager/user_views.py

@@ -11,6 +11,7 @@ from django.core.mail import EmailMultiAlternatives
 from django.template import loader
 from django.core.urlresolvers import reverse
 from django.contrib.auth import get_user_model
+from django.contrib.contenttypes.models import ContentType
 from django.contrib import messages
 from django.http import HttpResponseRedirect, QueryDict
 from django.views.generic.detail import SingleObjectMixin
@@ -20,7 +21,7 @@ from import_export.fields import Field
 import tablib
 
 from authentic2.constants import SWITCH_USER_SESSION_KEY
-from authentic2.models import Attribute, PasswordReset
+from authentic2.models import Attribute, AttributeValue, PasswordReset
 from authentic2.utils import switch_user, send_password_reset_mail, redirect, send_email_change_email
 from authentic2.a2_rbac.utils import get_default_ou
 from authentic2 import hooks
@@ -323,11 +324,6 @@ class UsersExportView(ExportMixin, UsersView):
     resource_class = UserResource
     export_prefix = 'users-'
 
-    def get_queryset(self):
-        '''Prefetch attribute values.'''
-        qs = super(UsersExportView, self).get_queryset()
-        return qs.prefetch_related('attribute_values', 'attribute_values__attribute')
-
     @property
     def csv(self):
         return self._dataset.export('csv')
@@ -338,7 +334,16 @@ class UsersExportView(ExportMixin, UsersView):
         attributes = [attr.name for attr in Attribute.objects.all()]
         headers = fields + tuple(['attribute_%s' % attr for attr in attributes])
 
+        at_mapping = {a.id: a for a in Attribute.objects.all()}
+        avs = AttributeValue.objects.filter(
+            content_type=ContentType.objects.get_for_model(get_user_model())).values()
+        user_attrs = collections.defaultdict(dict)
+        for av in avs:
+            user_attrs[av['object_id']][at_mapping[av['attribute_id']].name] = av['content']
+
         def iso(rec):
+            if rec is None or rec == {}:
+                return ''
             if hasattr(rec, 'strftime'):
                 if isinstance(rec, datetime.datetime):
                     _format = "%Y-%m-%d %H:%M:%S"
@@ -356,9 +361,9 @@ class UsersExportView(ExportMixin, UsersView):
                     value = getattr(user, field)
                 record.append(value)
 
-            attr_d = user.attributes.values
+            attr_d = user_attrs[user.pk]
             for attr in attributes:
-                record.append(attr_d.get(attr, ''))
+                record.append(attr_d.get(attr))
 
             return map(iso, record)
 

tests/test_user_manager.py

@@ -124,7 +124,7 @@ def test_export_csv(settings, app, superuser, django_assert_num_queries):
     user_count = User.objects.count()
     # queries should be batched to keep prefetching working without
     # overspending memory for the queryset cache, 4 queries by batches
-    num_queries = 9 + 4 * (user_count / DEFAULT_BATCH_SIZE + bool(user_count % DEFAULT_BATCH_SIZE))
+    num_queries = 4 + 4 * (user_count / DEFAULT_BATCH_SIZE + bool(user_count % DEFAULT_BATCH_SIZE))
     with django_assert_num_queries(num_queries):
          response = response.click('CSV')
     table = list(csv.reader(response.content.splitlines()))