from sqlalchemy.orm import sessionmaker

from twisted.internet import reactor
from twisted.internet import threads
from twisted.python.threadpool import ThreadPool

POOL_SIZE=5     #how many threads should the db connector pool should have

dbpool = ThreadPool(minthreads=1, maxthreads=POOL_SIZE, name='database')
dbpool.start()
reactor.addSystemEventTrigger('before', 'shutdown', dbpool.stop)

def run_in_db_thread(f):
    """Decorator to run DB queries in Twisted's thread pool"""
    def wrapper(*args, **kwargs):
        return threads.deferToThreadPool(reactor, dbpool,f, *args, **kwargs)
    return wrapper


class WithSession(object):
    '''a with statement for a database session connection'''
    def __init__(self, engine, autocommit=False):
        self.engine = engine
        self.autocommit=autocommit
    
    def __enter__(self):
        self.session = sessionmaker(bind=self.engine)()
        return self.session
    
    def __exit__(self,exc_type, exc_value, traceback):
        if exc_type is None:
            if self.autocommit:
                self.session.commit()
        else:
            self.session.rollback()
        self.session.close()

class DBDefer(object):
    '''a twisted sqlalchemy connector this Decorator adds a session parameter, with a valid session connection'''
    def __init__(self, engine, autocommit=False):
        self.autocommit=autocommit
        self.engine = engine

    def __call__(self, func):
        @run_in_db_thread
        def wrapper(*args, **kwargs):
            with WithSession(self.engine, self.autocommit) as session:
                return func(*args, session=session, **kwargs)
        return wrapper