Merge branch 'feature/issue-1303-clean-up-riak-clients' into develop

go/account/tasks.py

@@ -1,8 +1,10 @@
-from celery.task import task
+from contextlib import closing
 
+from celery.task import task
 from django.contrib.auth import get_user_model
 
 from go.account.utils import send_user_account_summary
+from go.base.utils import vumi_api, vumi_api_for_user
 
 
 @task(ignore_result=True)
@@ -11,26 +13,27 @@ def update_account_details(user_id, first_name=None, last_name=None,
                            confirm_start_conversation=None,
                            email_summary=None):
     user = get_user_model().objects.get(pk=user_id)
-    profile = user.get_profile()
-    account = profile.get_user_account()
+    with closing(vumi_api_for_user(user)) as user_api:
+        account = user_api.get_user_account()
 
-    if new_password:
-        user.set_password(new_password)
-    user.first_name = first_name
-    user.last_name = last_name
-    user.email = email_address
-    user.save()
+        if new_password:
+            user.set_password(new_password)
+        user.first_name = first_name
+        user.last_name = last_name
+        user.email = email_address
+        user.save()
 
-    account.msisdn = unicode(msisdn)
-    account.confirm_start_conversation = confirm_start_conversation
-    account.email_summary = email_summary
-    account.save()
+        account.msisdn = unicode(msisdn)
+        account.confirm_start_conversation = confirm_start_conversation
+        account.email_summary = email_summary
+        account.save()
 
 
 @task(ignore_result=True)
 def send_scheduled_account_summary(interval):
     users = get_user_model().objects.all()
-    for user in users:
-        user_account = user.get_profile().get_user_account()
-        if user_account.email_summary == interval:
-            send_user_account_summary(user)
+    with closing(vumi_api()) as api:
+        for user in users:
+            user_account = vumi_api_for_user(user, api).get_user_account()
+            if user_account.email_summary == interval:
+                send_user_account_summary(user)

go/account/utils.py

@@ -1,3 +1,4 @@
+from contextlib import closing
 from operator import attrgetter
 
 from go.base.utils import vumi_api_for_user
@@ -32,21 +33,22 @@ def get_messages_count(conversations):
 
 
 def send_user_account_summary(user):
-    user_api = vumi_api_for_user(user)
-    contact_store = user_api.contact_store
-    conv_store = user_api.conversation_store
+    with closing(vumi_api_for_user(user)) as user_api:
+        contact_store = user_api.contact_store
+        conv_store = user_api.conversation_store
 
-    contact_keys = contact_store.list_contacts()
-    uniques = get_uniques(contact_store, contact_keys=contact_keys,
-                            plucker=attrgetter('msisdn'))
-    conversation_keys = conv_store.list_conversations()
+        contact_keys = contact_store.list_contacts()
+        uniques = get_uniques(
+            contact_store, contact_keys=contact_keys,
+            plucker=attrgetter('msisdn'))
+        conversation_keys = conv_store.list_conversations()
 
-    all_conversations = []
-    bunches = conv_store.conversations.load_all_bunches(conversation_keys)
-    for bunch in bunches:
-        all_conversations.extend([user_api.wrap_conversation(conv)
-                                    for conv in bunch])
-    all_conversations.sort(key=(lambda conv: conv.created_at), reverse=True)
+        all_conversations = []
+        bunches = conv_store.conversations.load_all_bunches(conversation_keys)
+        for bunch in bunches:
+            all_conversations.extend(
+                [user_api.wrap_conversation(conv) for conv in bunch])
+        all_conversations.sort(key=attrgetter('created_at'), reverse=True)
 
     active_conversations = {}
     known_types = configured_conversation_types()

go/account/views.py

@@ -29,9 +29,8 @@ class GoRegistrationView(RegistrationView):
 
 @login_required
 def details(request):
-    profile = request.user.get_profile()
     token_manager = DjangoTokenManager(request.user_api.api.token_manager)
-    account = profile.get_user_account()
+    account = request.user_api.get_user_account()
     account_form = AccountForm(request.user, initial={
         'name': request.user.first_name,
         'surname': request.user.last_name,

go/api/conversation_api/conversation_api.py

@@ -157,6 +157,7 @@ def setup_worker(self):
     @inlineCallbacks
     def teardown_worker(self):
         yield self.webserver.loseConnection()
+        yield self.vumi_api.close()
 
     def setup_connectors(self):
         pass

go/api/go_api/go_api.py

@@ -289,6 +289,7 @@ def setup_worker(self):
     def teardown_worker(self):
         if self._web_service is not None:
             yield self._web_service.stopService()
+        yield self.vumi_api.close()
 
     def setup_connectors(self):
         pass

go/apps/surveys/tasks.py

@@ -1,14 +1,13 @@
+from contextlib import closing
 from StringIO import StringIO
 from zipfile import ZipFile, ZIP_DEFLATED
 
 from celery.task import task
-
 from django.conf import settings
 from django.core.mail import EmailMessage
 
-from go.vumitools.api import VumiUserApi
 from go.base.models import UserProfile
-
+from go.base.utils import vumi_api
 from go.apps.surveys.view_definition import get_poll_config
 
 
@@ -18,26 +17,27 @@ def export_vxpolls_data(account_key, conversation_key, include_old_questions):
     Export the data from a vxpoll and send it as a zipped attachment
     via email.
     """
-    api = VumiUserApi.from_config_sync(account_key, settings.VUMI_API_CONFIG)
-    user_profile = UserProfile.objects.get(user_account=account_key)
-    conversation = api.get_wrapped_conversation(conversation_key)
-
-    poll_id = 'poll-%s' % (conversation.key,)
-    pm, poll_data = get_poll_config(poll_id)
-    poll = pm.get(poll_id)
-    csv_data = pm.export_user_data_as_csv(
-        poll, include_old_questions=include_old_questions)
-    email = EmailMessage(
-        'Survey export for: %s' % (conversation.name,),
-        'Please find the data for the survey %s attached.\n' % (
-            conversation.name),
-        settings.DEFAULT_FROM_EMAIL, [user_profile.user.email])
-
-    zipio = StringIO()
-    zf = ZipFile(zipio, "a", ZIP_DEFLATED)
-    zf.writestr("survey-data-export.csv", csv_data)
-    zf.close()
-
-    email.attach('survey-data-export.zip', zipio.getvalue(),
-                 'application/zip')
-    email.send()
+    with closing(vumi_api()) as api:
+        user_api = api.get_user_api(account_key)
+        user_profile = UserProfile.objects.get(user_account=account_key)
+        conversation = user_api.get_wrapped_conversation(conversation_key)
+
+        poll_id = 'poll-%s' % (conversation.key,)
+        pm, poll_data = get_poll_config(poll_id)
+        poll = pm.get(poll_id)
+        csv_data = pm.export_user_data_as_csv(
+            poll, include_old_questions=include_old_questions)
+        email = EmailMessage(
+            'Survey export for: %s' % (conversation.name,),
+            'Please find the data for the survey %s attached.\n' % (
+                conversation.name),
+            settings.DEFAULT_FROM_EMAIL, [user_profile.user.email])
+
+        zipio = StringIO()
+        zf = ZipFile(zipio, "a", ZIP_DEFLATED)
+        zf.writestr("survey-data-export.csv", csv_data)
+        zf.close()
+
+        email.attach('survey-data-export.zip', zipio.getvalue(),
+                     'application/zip')
+        email.send()

go/base/command_utils.py

@@ -56,12 +56,20 @@ def make_email_option():
 
 
 class BaseGoCommand(BaseCommand):
+    _vumi_api = None
+
+    @property
+    def vumi_api(self):
+        if self._vumi_api is None:
+            self._vumi_api = vumi_api()
+        return self._vumi_api
+
     def list_commands(self):
         return [opt.const for opt in self.option_list
                 if opt.action == 'append_const' and opt.dest == 'command']
 
     def user_api_for_user(self, user):
-        return vumi_api_for_user(user)
+        return vumi_api_for_user(user, self.vumi_api)
 
     def mk_all_user_apis(self):
         apis = [(user, self.user_api_for_user(user)) for user in get_users()]
@@ -79,13 +87,14 @@ def mk_user_api(self, email_address=None, options=None):
         user_api = self.user_api_for_user(user)
         return user, user_api
 
-    def mk_vumi_api(self):
-        return vumi_api()
-
     def handle(self, *args, **options):
         self.options = options
-        self.vumi_api = self.mk_vumi_api()
-        return self.dispatch_command(*args, **options)
+        try:
+            return self.dispatch_command(*args, **options)
+        finally:
+            if self._vumi_api is not None:
+                self._vumi_api.close()
+                self._vumi_api = None
 
     def dispatch_command(self, *args, **options):
         commands = options.get('command', [])

go/base/management/commands/go_setup_env.py

@@ -157,6 +157,10 @@ def handle(self, *apps, **options):
 
         self.write_startup_script()
 
+    def cleanup(self):
+        self.redis.close_manager()
+        self.riak.close_manager()
+
     def setup_backend(self, config):
         self.redis = RedisManager.from_config(config['redis_manager'])
         self.riak = RiakManager.from_config(config['riak_manager'])
@@ -165,6 +169,9 @@ def setup_backend(self, config):
             self.redis.sub_manager('tagpool_store'))
         self.api = VumiApi(self.riak, self.redis)
 
+    def get_user_api(self, user):
+        return vumi_api_for_user(user, self.api)
+
     def read_yaml(self, file_path):
         # We remove a top-level '__ignore__' key which can contain blocks
         # referenced elsewhere.
@@ -243,7 +250,7 @@ def setup_account(self, user_info):
         user.save()
 
         profile = user.get_profile()
-        account = profile.get_user_account()
+        account = profile.get_user_account(self.api)
 
         for pool_name, max_keys in user_info['tagpools']:
             self.assign_tagpool(account, pool_name, max_keys)
@@ -276,14 +283,14 @@ def assign_application(self, account, application_module):
         return app_permission
 
     def setup_channels(self, user, channels):
-        user_api = vumi_api_for_user(user)
+        user_api = self.get_user_api(user)
         for channel in channels:
             tag = tuple(channel.split(':'))
             user_api.acquire_specific_tag(tag)
             self.stdout.write('Tag %s acquired\n' % (tag,))
 
     def setup_routers(self, user, routers):
-        user_api = vumi_api_for_user(user)
+        user_api = self.get_user_api(user)
         for router_info in routers:
             router_info = router_info.copy()  # So we can modify it.
             self.router_info.append({
@@ -316,7 +323,7 @@ def setup_routers(self, user, routers):
             self.stdout.write('Router %s created\n' % (router.key,))
 
     def setup_conversations(self, user, conversations):
-        user_api = vumi_api_for_user(user)
+        user_api = self.get_user_api(user)
         for conv_info in conversations:
             conv_info = conv_info.copy()  # So we can modify it.
             self.conversation_info.append({
@@ -362,7 +369,7 @@ def setup_routing(self, user, account_objects):
             rt.add_entry(
                 str(connectors[src]), src_ep, str(connectors[dst]), dst_ep)
 
-        user_account = vumi_api_for_user(user).get_user_account()
+        user_account = self.get_user_api(user).get_user_account()
         user_account.routing_table = rt
         user_account.save()
 
@@ -584,7 +591,7 @@ def create_billing_api_supervisord_conf(self):
         self.stdout.write('Wrote %s.\n' % (fn,))
 
     def setup_contact_groups(self, user, contact_groups):
-        user_api = vumi_api_for_user(user)
+        user_api = self.get_user_api(user)
         for group_info in contact_groups:
             self.contact_group_info.append({
                 'account': user.email,

go/base/middleware.py

@@ -1,5 +1,6 @@
 import time
 import logging
+from contextlib import closing
 
 from django.core.urlresolvers import resolve, Resolver404
 from vumi.blinkenlights.metrics import Metric
@@ -20,6 +21,23 @@ def process_request(self, request):
             SessionManager.set_user_account_key(
                 request.session, user_api.user_account_key)
 
+    def _cleanup_user_api(self, request):
+        try:
+            user_api = request.user_api
+        except AttributeError:
+            # Ignoring AttributeError as that just means this request wasn't
+            # seen by process_request (this is normal when other middleware
+            # returns a response).
+            return
+        user_api.close()
+
+    def process_response(self, request, response):
+        self._cleanup_user_api(request)
+        return response
+
+    def process_exception(self, request, exception):
+        self._cleanup_user_api(request)
+
 
 class ResponseTimeMiddleware(object):
     """
@@ -62,7 +80,8 @@ def process_response(self, request, response):
         response_time = start_time - time.time()
         response['X-Response-Time'] = response_time
         # TODO: Better way to fire these metrics.
-        metrics = vumi_api().get_metric_manager(get_django_metric_prefix())
-        metrics.oneshot(metric, response_time)
-        metrics.publish_metrics()
+        with closing(vumi_api()) as api:
+            metrics = api.get_metric_manager(get_django_metric_prefix())
+            metrics.oneshot(metric, response_time)
+            metrics.publish_metrics()
         return response