diff --git a/go/contacts/tasks.py b/go/contacts/tasks.py index 9b1b90c47..56ba1da34 100644 --- a/go/contacts/tasks.py +++ b/go/contacts/tasks.py @@ -149,6 +149,16 @@ def export_contacts(account_key, contact_keys, include_extra=True): api = VumiUserApi.from_config_sync(account_key, settings.VUMI_API_CONFIG) contact_store = api.contact_store + all_key_count = len(contact_keys) + + message_content_template = 'Please find the CSV data for %s contact(s)' + + if all_key_count > settings.CONTACT_EXPORT_TASK_LIMIT: + contact_keys = contact_keys[:settings.CONTACT_EXPORT_TASK_LIMIT] + message_content_template = '\n'.join([ + 'NOTE: There are too many contacts to export.', + 'Please find the CSV data for %%s (out of %s) contacts.' % ( + all_key_count,)]) contacts = contacts_by_key(contact_store, *contact_keys) data = contacts_to_csv(contacts, include_extra) @@ -159,8 +169,7 @@ def export_contacts(account_key, contact_keys, include_extra=True): user_profile = UserProfile.objects.get(user_account=account_key) email = EmailMessage( - 'Contacts export', - 'Please find the CSV data for %s contact(s)' % len(contacts), + 'Contacts export', message_content_template % len(contacts), settings.DEFAULT_FROM_EMAIL, [user_profile.user.email]) email.attach('contacts-export.zip', file, 'application/zip') diff --git a/go/contacts/tests.py b/go/contacts/tests.py index 54652098c..f0fd6456d 100644 --- a/go/contacts/tests.py +++ b/go/contacts/tests.py @@ -221,6 +221,64 @@ def test_contact_exporting(self): self.assertTrue(contents) self.assertEqual(mime_type, 'application/zip') + def test_contact_exporting_too_many(self): + """ + If we have too many contacts to export, we only export a subset. + """ + self.vumi_helper.patch_settings(CONTACT_EXPORT_TASK_LIMIT=2) + + c1 = self.mkcontact() + c1.extra['foo'] = u'bar' + c1.save() + + c2 = self.mkcontact() + c2.extra['foo'] = u'baz' + c2.save() + + c3 = self.mkcontact() + c3.extra['foo'] = u'quux' + c3.save() + + response = self.client.post(reverse('contacts:people'), { + '_export': True, + 'contact': [c1.key, c2.key, c3.key], + }) + + self.assertContains( + response, + "The export is scheduled and should complete within a few" + " minutes.") + + self.assertEqual(len(mail.outbox), 1) + [email] = mail.outbox + [(file_name, contents, mime_type)] = email.attachments + + self.assertEqual(email.recipients(), [self.user_email]) + self.assertTrue('Contacts export' in email.subject) + self.assertTrue( + 'NOTE: There are too many contacts to export.' in email.body) + self.assertTrue('2 (out of 3) contacts' in email.body) + self.assertEqual(file_name, 'contacts-export.zip') + + zipfile = ZipFile(StringIO(contents), 'r') + csv_contents = zipfile.open('contacts-export.csv', 'r').read() + + [header, c1_data, c2_data, _] = csv_contents.split('\r\n') + + self.assertEqual( + header, + ','.join([ + 'key', 'name', 'surname', 'email_address', 'msisdn', 'dob', + 'twitter_handle', 'facebook_id', 'bbm_pin', 'gtalk_id', + 'mxit_id', 'wechat_id', 'created_at', 'foo'])) + + self.assertTrue(c1_data.startswith(c1.key)) + self.assertTrue(c1_data.endswith(',bar')) + self.assertTrue(c2_data.startswith(c2.key)) + self.assertTrue(c2_data.endswith(',baz')) + self.assertTrue(contents) + self.assertEqual(mime_type, 'application/zip') + def test_exporting_all_contacts(self): c1 = self.mkcontact() c1.extra['foo'] = u'bar' diff --git a/go/settings.py b/go/settings.py index e2cf00f88..db2d31e85 100644 --- a/go/settings.py +++ b/go/settings.py @@ -309,6 +309,10 @@ def static_paths(paths): # }, } + +# Exporting hundreds of thousands of contacts makes celery use all the memory. +CONTACT_EXPORT_TASK_LIMIT = 100000 + try: from production_settings import * except ImportError as err: