iro/controller/database.py
branchdevel
changeset 96 883b25b0b760
parent 93 ee10253258ad
child 102 85ac86f7ffc4
equal deleted inserted replaced
94:4349f10d77a6 96:883b25b0b760
     1 from sqlalchemy import create_engine
     1 from sqlalchemy import create_engine, pool
       
     2 
     2 from sqlalchemy.orm import sessionmaker
     3 from sqlalchemy.orm import sessionmaker
       
     4 
       
     5 
       
     6 from twisted.internet import threads
       
     7 
     3 from ..model import Base
     8 from ..model import Base
     4 
     9 
     5 engine = create_engine('sqlite:///:memory:', echo=True)
    10 engine = create_engine('sqlite:///:memory:', echo=True)
     6 
    11 
     7 def createDatabase():
    12 def createDatabase():
     8     Base.metadata.create_all(engine)
    13     Base.metadata.create_all(engine)
     9 
    14 
    10 Session = sessionmaker(bind=engine)
    15 Session = sessionmaker(bind=engine)
    11 
    16 
       
    17 def toThread(f):
       
    18     def wrapper(*args, **kwargs):
       
    19         return threads.deferToThread(f, *args, **kwargs)
       
    20     return wrapper
       
    21 
       
    22 
    12 class WithSession():
    23 class WithSession():
    13     def __init__(self,autocommit=False):
    24     def __init__(self,autocommit=False):
    14         self.autocommit=autocommit
    25         self.autocommit=autocommit
    15 
    26     
    16     def __enter__(self):
    27     def __enter__(self):
    17         self.conn = engine.connect()
    28         self.conn = engine.connect()
    18         self.session = Session(bind=self.conn)
    29         self.session = Session(bind=self.conn)
    19         return self.session
    30         return self.session
    20     
    31     
    25         else:
    36         else:
    26             self.session.rollback()
    37             self.session.rollback()
    27         self.session.close()
    38         self.session.close()
    28         self.conn.close()
    39         self.conn.close()
    29 
    40 
       
    41 class DBDefer(object):
       
    42     def __init__(self, dsn, poolclass = pool.SingletonThreadPool, *args, **kargs):
       
    43         self.engine = create_engine(dsn, poolclass=poolclass, *args, **kargs)
       
    44 
       
    45     def __call__(self, func):
       
    46         @toThread
       
    47         def wrapper(*args, **kwargs):
       
    48             session = sessionmaker(bind=self.engine)()
       
    49             try:
       
    50                 print func,args, kwargs
       
    51                 return func(*args, session=session, **kwargs)
       
    52             except:
       
    53                 session.rollback()
       
    54                 raise
       
    55             finally:
       
    56                 session.close()
       
    57         return wrapper