Move the store pool from HyperKitty to the Storm module

1 files changed, 34 insertions, 8 deletions

diff --git a/kittystore/storm/__init__.py b/kittystore/storm/__init__.py
index e4d716d..8174e55 100644
--- a/kittystore/storm/__init__.py
+++ b/kittystore/storm/__init__.py
@@ -3,6 +3,7 @@
 from __future__ import absolute_import, with_statement
 
 import sys
+import threading
 
 import storm.tracer
 from storm.locals import create_database, Store
@@ -13,12 +14,37 @@ from . import schema
 from .store import StormStore
 
 
+class ThreadSafeStorePool(object):
+    """
+    Storm does not have a thread pool, like SQLAlchemy. Solve the threading
+    problem by keeping the store in a thread-local object.
+
+    http://unpythonic.blogspot.fr/2007/11/using-storm-and-sqlite-in-multithreaded.html
+    """
+
+    def __init__(self, url, debug):
+        self.url = url
+        self.debug = debug
+        self._local = threading.local()
+
+    def get(self):
+        try:
+            return self._local.store
+        except AttributeError:
+            self._local.store = self.create_store()
+            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 get_storm_store(url, debug=False):
-    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)
+    store_pool = ThreadSafeStorePool(url, debug)
+    return store_pool.get()