diff --git a/src/authentic2/manager/user_views.py b/src/authentic2/manager/user_views.py index 3e03e5917..3d76da1cb 100644 --- a/src/authentic2/manager/user_views.py +++ b/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) diff --git a/tests/test_user_manager.py b/tests/test_user_manager.py index f5ca8d1b1..a23fc1b87 100644 --- a/tests/test_user_manager.py +++ b/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()))