Use a middleware to avoid KS connection leaks

8 files changed, 85 insertions, 18 deletions

diff --git a/hyperkitty/api.py b/hyperkitty/api.py
index e085139..f28a1af 100644
--- a/hyperkitty/api.py
+++ b/hyperkitty/api.py
@@ -7,8 +7,7 @@ from django.http import HttpResponseNotModified, HttpResponse
 import json
 import re
 
-from kittystore import get_store
-
+from hyperkitty.lib import get_store
 from hyperkitty.utils import log
 
 
@@ -19,7 +18,7 @@ class EmailResource(View):
 
     def get(self, request, mlist_fqdn, messageid):
         list_name = mlist_fqdn.split('@')[0]
-        store = get_store(settings.KITTYSTORE_URL)
+        store = get_store(request)
         email = store.get_message_by_hash_from_list(list_name, messageid)
         if not email:
             return HttpResponse(status=404)
@@ -34,7 +33,7 @@ class ThreadResource(View):
 
     def get(self, request, mlist_fqdn, threadid):
         list_name = mlist_fqdn.split('@')[0]
-        store = get_store(settings.KITTYSTORE_URL)
+        store = get_store(request)
         thread = store.get_messages_in_thread(list_name, threadid)
         if not thread:
             return HttpResponse(status=404)
@@ -63,7 +62,7 @@ class SearchResource(View):
                 re.compile(regex, re.IGNORECASE)}
 
         #print query_string, field, keyword
-        store = get_store(settings.KITTYSTORE_URL)
+        store = get_store(request)
         threads = store.search_archives(list_name, query_string)
         if not threads:
             return HttpResponse(status=404)
diff --git a/hyperkitty/lib/__init__.py b/hyperkitty/lib/__init__.py
index be8e27d..710011a 100644
--- a/hyperkitty/lib/__init__.py
+++ b/hyperkitty/lib/__init__.py
@@ -5,6 +5,7 @@ from hashlib import md5
 import datetime
 
 from django.conf import settings
+import kittystore
 
 from hyperkitty.utils import log
 
@@ -41,3 +42,6 @@ def get_months(store, list_name):
     archives[now.year] = range(1, 13)[:now.month]
     return archives
 
+
+def get_store(request):
+    return request.environ["kittystore.store"]
diff --git a/hyperkitty/lib/store.py b/hyperkitty/lib/store.py
new file mode 100644
index 0000000..c20c039
--- /dev/null
+++ b/hyperkitty/lib/store.py
@@ -0,0 +1,64 @@
+"""
+WSGI and Django middlewares for KittyStore
+
+Inspired by http://pypi.python.org/pypi/middlestorm
+"""
+
+from threading import local
+
+from django.conf import settings
+import kittystore
+
+
+class KittyStoreWSGIMiddleware(object):
+    """WSGI middleware.
+    Add KittyStore object in environ['kittystore.store']. Each thread contains
+    own store.
+    """
+
+    def __init__(self, app):
+        """Create WSGI middleware.
+        :param app: top level application or middleware.
+        :param database: instance of Database returned create_database.
+        """
+        self._app = app
+        self._local = local()
+
+    def __call__(self, environ, start_response):
+        try:
+            environ['kittystore.store']  = self._local.store
+        except AttributeError:
+            environ['kittystore.store']  = \
+                    self._local.__dict__.setdefault('store',
+                        kittystore.get_store(settings.KITTYSTORE_URL))
+        try:
+            return self._app(environ, start_response)
+        finally:
+            environ['kittystore.store'].rollback()
+            #environ['kittystore.store'].close()
+
+
+class KittyStoreDjangoMiddleware(object):
+    """Django middleware.
+    Add KittyStore object in environ['kittystore.store']. Each thread contains
+    own store.
+    """
+
+    def __init__(self):
+        """Create Django middleware."""
+        self._local = local()
+
+    def process_request(self, request):
+        try:
+            request.environ['kittystore.store']  = self._local.store
+        except AttributeError:
+            request.environ['kittystore.store']  = \
+                    self._local.__dict__.setdefault('store',
+                        kittystore.get_store(settings.KITTYSTORE_URL))
+
+    #def process_response(self, request, response):
+    #    request.environ['kittystore.store'].close()
+    #    return response
+
+    def process_exception(self, request, exception):
+        request.environ['kittystore.store'].rollback()
diff --git a/hyperkitty/models.py b/hyperkitty/models.py
index dae055a..6a8bd63 100644
--- a/hyperkitty/models.py
+++ b/hyperkitty/models.py
@@ -61,6 +61,8 @@ class UserProfile(models.Model):
         except Rating.DoesNotExist:
             votes = {}
 
+        # TODO: warning, not thread-safe, should get the cached connection from
+        # the WSGI environment
         store = get_store(settings.KITTYSTORE_URL)
         for vote in votes:
             list_name = vote.list_address.split('@')[0]
diff --git a/hyperkitty/views/list.py b/hyperkitty/views/list.py
index 8825511..a8cae9b 100644
--- a/hyperkitty/views/list.py
+++ b/hyperkitty/views/list.py
@@ -16,11 +16,10 @@ from django.core.paginator import Paginator, EmptyPage, PageNotAnInteger, Invali
 from django.contrib.auth.decorators import (login_required,
                                             permission_required,
                                             user_passes_test)
-from kittystore import get_store
 
 from hyperkitty.models import Rating, Tag
 #from hyperkitty.lib.mockup import *
-from hyperkitty.lib import get_months
+from hyperkitty.lib import get_months, get_store
 from forms import *
 from hyperkitty.utils import log
 
@@ -69,7 +68,7 @@ def archives(request, mlist_fqdn, year=None, month=None, day=None):
 
     search_form = SearchForm(auto_id=False)
     t = loader.get_template('month_view.html')
-    store = get_store(settings.KITTYSTORE_URL)
+    store = get_store(request)
     threads = store.get_threads(mlist_fqdn, start=begin_date,
         end=end_date)
 
@@ -174,7 +173,7 @@ def list(request, mlist_fqdn=None):
     end_date = datetime(today.year, today.month, today.day)
     begin_date = end_date - timedelta(days=32)
 
-    store = get_store(settings.KITTYSTORE_URL)
+    store = get_store(request)
     threads = store.get_threads(list_name=mlist_fqdn, start=begin_date,
         end=end_date)
 
@@ -272,7 +271,7 @@ def _search_results_page(request, mlist_fqdn, threads, search_type,
     except (EmptyPage, InvalidPage):
         threads = paginator.page(paginator.num_pages)
 
-    store = get_store(settings.KITTYSTORE_URL)
+    store = get_store(request)
     cnt = 0
     for msg in threads.object_list:
         msg.email = msg.sender_email.strip()
@@ -318,7 +317,7 @@ def search(request, mlist_fqdn):
 def search_keyword(request, mlist_fqdn, target, keyword, page=1):
     ## Should we remove the code below? 
     ## If urls.py does it job we should never need it
-    store = get_store(settings.KITTYSTORE_URL)
+    store = get_store(request)
     if not keyword:
         keyword = request.GET.get('keyword')
     if not target:
diff --git a/hyperkitty/views/message.py b/hyperkitty/views/message.py
index 86ae139..dd80709 100644
--- a/hyperkitty/views/message.py
+++ b/hyperkitty/views/message.py
@@ -9,10 +9,10 @@ from django.core.paginator import Paginator, EmptyPage, PageNotAnInteger, Invali
 from django.contrib.auth.decorators import (login_required,
                                             permission_required,
                                             user_passes_test)
-from kittystore import get_store
 
 from hyperkitty.models import Rating
 #from hyperkitty.lib.mockup import *
+from hyperkitty.lib import get_store
 from forms import *
 from hyperkitty.utils import log
 
@@ -26,7 +26,7 @@ def index (request, mlist_fqdn, hashid):
 
     search_form = SearchForm(auto_id=False)
     t = loader.get_template('message.html')
-    store = get_store(settings.KITTYSTORE_URL)
+    store = get_store(request)
     message = store.get_message_by_hash_from_list(mlist_fqdn, hashid)
     message.sender_email = message.sender_email.strip()
     # Extract all the votes for this message
diff --git a/hyperkitty/views/pages.py b/hyperkitty/views/pages.py
index a3de806..c9d6e19 100644
--- a/hyperkitty/views/pages.py
+++ b/hyperkitty/views/pages.py
@@ -17,10 +17,10 @@ from django.core.paginator import Paginator, EmptyPage, PageNotAnInteger, Invali
 from django.contrib.auth.decorators import (login_required,
                                             permission_required,
                                             user_passes_test)
-from kittystore import get_store
 
 from hyperkitty.models import Rating
 #from hyperkitty.lib.mockup import *
+from hyperkitty.lib import get_store
 from forms import *
 from hyperkitty.utils import log
 
@@ -31,7 +31,7 @@ def index(request):
     base_url = settings.MAILMAN_API_URL % {
         'username': settings.MAILMAN_USER, 'password': settings.MAILMAN_PASS}
 
-    store = get_store(settings.KITTYSTORE_URL)
+    store = get_store(request)
     list_data = store.get_list_names()
     log("warn", repr(list_data))
 
diff --git a/hyperkitty/views/thread.py b/hyperkitty/views/thread.py
index bef9c40..b5bf80a 100644
--- a/hyperkitty/views/thread.py
+++ b/hyperkitty/views/thread.py
@@ -9,13 +9,12 @@ from django.core.paginator import Paginator, EmptyPage, PageNotAnInteger, Invali
 from django.contrib.auth.decorators import (login_required,
                                             permission_required,
                                             user_passes_test)
-from kittystore import get_store
 
 from hyperkitty.models import Rating, Tag
 #from hyperkitty.lib.mockup import *
 from forms import *
 from hyperkitty.utils import log
-from hyperkitty.lib import get_months
+from hyperkitty.lib import get_months, get_store
 
 
 
@@ -25,7 +24,7 @@ def thread_index (request, mlist_fqdn, threadid):
 
     search_form = SearchForm(auto_id=False)
     t = loader.get_template('thread.html')
-    store = get_store(settings.KITTYSTORE_URL)
+    store = get_store(request)
     messages = store.get_messages_in_thread(mlist_fqdn, threadid)
     #prev_thread = mongo.get_thread_name(list_name, int(threadid) - 1)
     prev_thread = []