|
1 from sqlalchemy.orm import sessionmaker |
|
2 |
|
3 from twisted.internet import reactor |
|
4 from twisted.internet import threads |
|
5 from twisted.python.threadpool import ThreadPool |
|
6 |
|
7 POOL_SIZE=5 #how many threads should the db connector pool should have |
|
8 |
|
9 dbpool = ThreadPool(minthreads=1, maxthreads=POOL_SIZE, name='database') |
|
10 dbpool.start() |
|
11 reactor.addSystemEventTrigger('before', 'shutdown', dbpool.stop) |
|
12 |
|
13 def run_in_db_thread(f): |
|
14 """Decorator to run DB queries in Twisted's thread pool""" |
|
15 def wrapper(*args, **kwargs): |
|
16 return threads.deferToThreadPool(reactor, dbpool,f, *args, **kwargs) |
|
17 return wrapper |
|
18 |
|
19 |
|
20 class WithSession(object): |
|
21 '''a with statement for a database session connection''' |
|
22 def __init__(self, engine, autocommit=False): |
|
23 self.engine = engine |
|
24 self.autocommit=autocommit |
|
25 |
|
26 def __enter__(self): |
|
27 self.session = sessionmaker(bind=self.engine)() |
|
28 return self.session |
|
29 |
|
30 def __exit__(self,exc_type, exc_value, traceback): |
|
31 if exc_type is None: |
|
32 if self.autocommit: |
|
33 self.session.commit() |
|
34 else: |
|
35 self.session.rollback() |
|
36 self.session.close() |
|
37 |
|
38 class DBDefer(object): |
|
39 '''a twisted sqlalchemy connector this Decorator adds a session parameter, with a valid session connection''' |
|
40 def __init__(self, engine, autocommit=False): |
|
41 self.autocommit=autocommit |
|
42 self.engine = engine |
|
43 |
|
44 def __call__(self, func): |
|
45 @run_in_db_thread |
|
46 def wrapper(*args, **kwargs): |
|
47 with WithSession(self.engine, self.autocommit) as session: |
|
48 return func(*args, session=session, **kwargs) |
|
49 return wrapper |