custom attributes prefetching (#29531)

This commit is contained in:
Emmanuel Cazenave 2019-01-15 14:52:06 +01:00
parent 56c72c2cf2
commit 93c52a940c
2 changed files with 14 additions and 9 deletions

View File

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

View File

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