Use a middleware in HK to avoid connection leaks

The middleware also handle thread-safety issues

2 files changed, 22 insertions, 12 deletions

diff --git a/kittystore/storm/__init__.py b/kittystore/storm/__init__.py
index 8174e55..da1a56d 100644
--- a/kittystore/storm/__init__.py
+++ b/kittystore/storm/__init__.py
@@ -31,20 +31,23 @@ class ThreadSafeStorePool(object):
         try:
             return self._local.store
         except AttributeError:
-            self._local.store = self.create_store()
+            self._local.store = create_store(self.url, self.debug)
             return self._local.store
 
-    def create_store(self):
-        if self.debug:
-            storm.tracer.debug(True, stream=sys.stdout)
-        database = create_database(self.url)
-        store = Store(database)
-        dbtype = self.url.partition(":")[0]
-        dbschema = Schema(schema.CREATES[dbtype], [], [], schema)
-        dbschema.upgrade(store)
-        return StormStore(store, self.debug)
+
+def create_store(url, debug):
+    if debug:
+        storm.tracer.debug(True, stream=sys.stdout)
+    database = create_database(url)
+    store = Store(database)
+    dbtype = url.partition(":")[0]
+    dbschema = Schema(schema.CREATES[dbtype], [], [], schema)
+    dbschema.upgrade(store)
+    return StormStore(store, debug)
 
 
 def get_storm_store(url, debug=False):
-    store_pool = ThreadSafeStorePool(url, debug)
-    return store_pool.get()
+    # Thread safety is managed by the middleware
+    #store_pool = ThreadSafeStorePool(url, debug)
+    #return store_pool.get()
+    return create_store(url, debug)
diff --git a/kittystore/storm/store.py b/kittystore/storm/store.py
index 0c5b91e..8bb3872 100644
--- a/kittystore/storm/store.py
+++ b/kittystore/storm/store.py
@@ -366,3 +366,10 @@ class StormStore(object):
     def commit(self):
         """Commit transaction to the database."""
         self.db.commit()
+
+    def close(self):
+        """Close the connection."""
+        self.db.close()
+
+    def rollback(self):
+        self.db.rollback()