cherrypy sqlalchemy plugin

cp_sqlalchemy_plugin.py

from cherrypy.process.plugins import SimplePlugin
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base, declared_attr
from sqlalchemy.orm.scoping import scoped_session
from sqlalchemy.orm.session import sessionmaker
from sqlalchemy.schema import Column
from sqlalchemy.types import Integer
import sqlamp
import cherrypy

# todo-eneldo Documentar
# todo-eneldo capturar exceptions
# todo-eneldo crear test unit


class Base(object):

    @declared_attr
    def __tablename__(cls):
        return cls.__name__.lower()

    id = Column(Integer, primary_key=True)


class SQLAlchemyPlugin(SimplePlugin):

    def __init__(self, engine, echo=False, autoflush=True, autocommit=False):
        SimplePlugin.__init__(self, engine)
        self.uri = cherrypy.config.get("SQLALCHEMY_DATABASE_URI", "sqlite://")
        self.engine = create_engine(self.uri, echo=echo)
        self.session_factory = sessionmaker(autoflush=autoflush, autocommit=autocommit, bind=self.engine)
        self.session = None
        self.Model = declarative_base(cls=Base, bind=self.engine)
        self.TreeModel = declarative_base(cls=Base, bind=self.engine, metaclass=sqlamp.DeclarativeMeta)
        self.subscribe()

    def start(self):
        self.create_all()
        self.Session = scoped_session(self.session_factory)
        self.bus.subscribe("bind-session", self.bind)
        self.bus.subscribe("commit-session", self.commit)

    def stop(self):
        self.bus.unsubscribe("bind-session", self.bind)
        self.bus.unsubscribe("commit-session", self.commit)
        if self.engine:
            self.engine.dispose()
            self.engine = None

    def bind(self):
        self.session = self.Session()
        return self.session

    def commit(self):
        try:
            self.session.commit()
        except:
            self.session.rollback()
        #    raise
        #finally:
        #    self.session.remove()

    def create_all(self):
        self.Model.metadata.create_all()
        self.TreeModel.metadata.create_all()

    def drop_all(self):
        self.Model.metadata.drop_all()
        self.TreeModel.metadata.drop_all()

    @property
    def metadata(self):
        return self.Model.metadata


class SqlalchemyTool(cherrypy.Tool):

    def __init__(self):

        cherrypy.Tool.__init__(self, "on_start_resource",
                               self.bind_session,
                               priority=20)

    def _setup(self):
        cherrypy.Tool._setup(self)
        cherrypy.request.hooks.attach("on_end_resource",
                                      self.commit_transaction,
                                      priority=80)

    def bind_session(self):
        session = cherrypy.engine.publish("bind-session").pop()
        cherrypy.request.db = session

    def commit_transaction(self):
        if not hasattr(cherrypy.request, 'db'):
            return
        cherrypy.request.db = None
        cherrypy.engine.publish('commit-session').pop()

cherrypy.tools.db = SqlalchemyTool()


def sqlalchemySetup(engine, echo=False, autoflush=True, autocommit=False):
    return SQLAlchemyPlugin(engine, echo, autoflush, autocommit)