photolib/photoapp/dbutils.py

91 lines
2.7 KiB
Python

import sqlalchemy
import cherrypy
from cherrypy.process import plugins
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.pool import NullPool
from sqlalchemy.orm import sessionmaker
Base = declarative_base()
engine_specific_options = {"sqlite": dict(connect_args={'check_same_thread': False},
poolclass=NullPool,
pool_pre_ping=True),
"mysql": dict(pool_pre_ping=True)}
def get_engine_options(uri):
for engine_prefix, options in engine_specific_options.items():
if uri.startswith(engine_prefix):
return options
return {}
def get_db_engine(uri, debug=False):
engine = sqlalchemy.create_engine(uri, **get_engine_options(uri), echo=debug)
Base.metadata.create_all(engine)
return engine
def get_db_session(uri):
engine = get_db_engine(uri)
session = sessionmaker()
session.configure(bind=engine)
return session
class DbAlias(object):
"""
This provides a shorter alias for the cherrypy.request.db object, which is a database session created bound to the
current request. Since the `db` attribute doesn't exist until a request is received, we cannot simply reference it
with another variable.
And instance of this class acts as an object proxy to the database object in cherrypy.request.db.
"""
def __getattr__(self, attr):
return getattr(cherrypy.request.db, attr)
db = DbAlias()
class SAEnginePlugin(plugins.SimplePlugin):
def __init__(self, bus, dbcon):
plugins.SimplePlugin.__init__(self, bus)
self.sa_engine = dbcon
self.bus.subscribe("bind", self.bind)
def start(self):
Base.metadata.create_all(self.sa_engine)
def bind(self, session):
session.configure(bind=self.sa_engine)
class SATool(cherrypy.Tool):
def __init__(self):
cherrypy.Tool.__init__(self, 'before_request_body',
self.bind_session,
priority=100)
self.session = sqlalchemy.orm.scoped_session(
sqlalchemy.orm.sessionmaker(autoflush=True, autocommit=False))
def _setup(self):
cherrypy.Tool._setup(self)
cherrypy.request.hooks.attach('on_end_resource', self.commit_transaction, priority=80)
def bind_session(self):
cherrypy.engine.publish('bind', self.session)
cherrypy.request.db = self.session
def commit_transaction(self):
cherrypy.request.db = None
try:
self.session.commit()
except Exception:
self.session.rollback()
raise
finally:
self.session.remove()