Merge remote-tracking branch 'thrisp/sqlalchemy_events_revision'

mitsuhiko · mitsuhiko · commit a5b051031d90 · 2013-05-15T17:27:14.000+01:00
diff --git a/flask_sqlalchemy.py b/flask_sqlalchemy.py
@@ -24,8 +24,7 @@
 from sqlalchemy import orm, event
 from sqlalchemy.orm.exc import UnmappedClassError
 from sqlalchemy.orm.session import Session
-from sqlalchemy.orm.interfaces import MapperExtension, SessionExtension, \
-     EXT_CONTINUE
+from sqlalchemy.event import listen
 from sqlalchemy.interfaces import ConnectionProxy
 from sqlalchemy.engine.url import make_url
 from sqlalchemy.ext.declarative import declarative_base, DeclarativeMeta
@@ -91,7 +90,6 @@ def _include_sqlalchemy(obj):
                 setattr(obj, key, getattr(module, key))
     # Note: obj.Table does not attempt to be a SQLAlchemy Table class.
     obj.Table = _make_table(obj)
-    obj.mapper = signalling_mapper
     obj.relationship = _wrap_with_default_query_class(obj.relationship)
     obj.relation = _wrap_with_default_query_class(obj.relation)
     obj.dynamic_loader = _wrap_with_default_query_class(obj.dynamic_loader)
@@ -155,67 +153,74 @@ def cursor_execute(self, execute, cursor, statement, parameters,
                     _calling_context(self.app_package))))
 
 
-class _SignalTrackingMapperExtension(MapperExtension):
-
-    def after_delete(self, mapper, connection, instance):
-        return self._record(mapper, instance, 'delete')
+class _SignallingSession(Session):
 
-    def after_insert(self, mapper, connection, instance):
-        return self._record(mapper, instance, 'insert')
+    def __init__(self, db, autocommit=False, autoflush=False, **options):
+        self.app = db.get_app()
+        self._model_changes = {}
+        Session.__init__(self, autocommit=autocommit, autoflush=autoflush,
+                         bind=db.engine,
+                         binds=db.get_binds(self.app), **options)
 
-    def after_update(self, mapper, connection, instance):
-        return self._record(mapper, instance, 'update')
+    def get_bind(self, mapper, clause=None):
+        # mapper is None if someone tries to just get a connection
+        if mapper is not None:
+            info = getattr(mapper.mapped_table, 'info', {})
+            bind_key = info.get('bind_key')
+            if bind_key is not None:
+                state = get_state(self.app)
+                return state.db.get_engine(self.app, bind=bind_key)
+        return Session.get_bind(self, mapper, clause)
 
-    def _record(self, mapper, model, operation):
-        s = orm.object_session(model)
-        # Skip the operation tracking when a non signalling session
-        # is used.
-        if isinstance(s, _SignallingSessionExtension):
-            pk = tuple(mapper.primary_key_from_instance(model))
-            s._model_changes[pk] = (model, operation)
-        return EXT_CONTINUE
 
+class _SessionSignalEvents(object):
 
-class _SignallingSessionExtension(SessionExtension):
+    def register(self):
+        listen(Session, 'before_commit', self.session_signal_before_commit)
+        listen(Session, 'after_commit', self.session_signal_after_commit)
+        listen(Session, 'after_rollback', self.session_signal_after_rollback)
 
-    def before_commit(self, session):
+    @staticmethod
+    def session_signal_before_commit(session):
         d = session._model_changes
         if d:
             before_models_committed.send(session.app, changes=d.values())
-        return EXT_CONTINUE
 
-    def after_commit(self, session):
+    @staticmethod
+    def session_signal_after_commit(session):
         d = session._model_changes
         if d:
             models_committed.send(session.app, changes=d.values())
             d.clear()
-        return EXT_CONTINUE
 
-    def after_rollback(self, session):
+    @staticmethod
+    def session_signal_after_rollback(session):
         session._model_changes.clear()
-        return EXT_CONTINUE
 
 
-class _SignallingSession(Session):
+class _MapperSignalEvents(object):
 
-    def __init__(self, db, autocommit=False, autoflush=False, **options):
-        self.app = db.get_app()
-        self._model_changes = {}
-        Session.__init__(self, autocommit=autocommit, autoflush=autoflush,
-                         extension=db.session_extensions,
-                         bind=db.engine,
-                         binds=db.get_binds(self.app), **options)
+    def __init__(self, mapper):
+        self.mapper = mapper
 
-    def get_bind(self, mapper, clause=None):
-        # mapper is None if someone tries to just get a connection
-        if mapper is not None:
-            info = getattr(mapper.mapped_table, 'info', {})
-            bind_key = info.get('bind_key')
-            if bind_key is not None:
-                state = get_state(self.app)
-                return state.db.get_engine(self.app, bind=bind_key)
-        return Session.get_bind(self, mapper, clause)
+    def register(self):
+        listen(self.mapper, 'after_delete', self.mapper_signal_after_delete)
+        listen(self.mapper, 'after_insert', self.mapper_signal_after_insert)
+        listen(self.mapper, 'after_update', self.mapper_signal_after_update)
+
+    def mapper_signal_after_delete(self, mapper, connection, target):
+        self._record(mapper, target, 'delete')
 
+    def mapper_signal_after_insert(self, mapper, connection, target):
+        self._record(mapper, target, 'insert')
+
+    def mapper_signal_after_update(self, mapper, connection, target):
+        self._record(mapper, target, 'update')
+
+    @staticmethod
+    def _record(mapper, target, operation):
+        pk = tuple(mapper.primary_key_from_instance(target))
+        orm.object_session(target)._model_changes[pk] = (target, operation)
 
 def get_debug_queries():
     """In debug mode Flask-SQLAlchemy will log all the SQL queries sent
@@ -504,14 +509,6 @@ def get_state(app):
     return app.extensions['sqlalchemy']
 
 
-def signalling_mapper(*args, **kwargs):
-    """Replacement for mapper that injects some extra extensions"""
-    extensions = to_list(kwargs.pop('extension', None), [])
-    extensions.append(_SignalTrackingMapperExtension())
-    kwargs['extension'] = extensions
-    return sqlalchemy.orm.mapper(*args, **kwargs)
-
-
 class _SQLAlchemyState(object):
     """Remembers configuration for the (db, app) tuple."""
 
@@ -619,11 +616,10 @@ class User(db.Model):
         a custom function which will define the SQLAlchemy session's scoping.
     """
 
-    def __init__(self, app=None, use_native_unicode=True,
-                 session_extensions=None, session_options=None):
+    def __init__(self, app=None,
+                 use_native_unicode=True,
+                 session_options=None):
         self.use_native_unicode = use_native_unicode
-        self.session_extensions = to_list(session_extensions, []) + \
-                                  [_SignallingSessionExtension()]
 
         if session_options is None:
             session_options = {}
@@ -643,6 +639,8 @@ def __init__(self, app=None, use_native_unicode=True,
             self.app = None
 
         _include_sqlalchemy(self)
+        _MapperSignalEvents(self.mapper).register()
+        _SessionSignalEvents().register()
         self.Query = BaseQuery
 
     @property
@@ -662,7 +660,6 @@ def create_scoped_session(self, options=None):
     def make_declarative_base(self):
         """Creates the declarative base."""
         base = declarative_base(cls=Model, name='Model',
-                                mapper=signalling_mapper,
                                 metaclass=_BoundDeclarativeMeta)
         base.query = _QueryProperty(self)
         return base
diff --git a/test_sqlalchemy.py b/test_sqlalchemy.py
@@ -6,7 +6,6 @@
 from datetime import datetime
 import flask
 from flask.ext import sqlalchemy
-from sqlalchemy.orm import sessionmaker
 
 
 def make_todo_model(db):
@@ -71,11 +70,11 @@ def test_query_recording(self):
             queries = sqlalchemy.get_debug_queries()
             self.assertEqual(len(queries), 1)
             query = queries[0]
-            self.assert_('insert into' in query.statement.lower())
+            self.assertTrue('insert into' in query.statement.lower())
             self.assertEqual(query.parameters[0], 'Test 1')
             self.assertEqual(query.parameters[1], 'test')
-            self.assert_('test_sqlalchemy.py' in query.context)
-            self.assert_('test_query_recording' in query.context)
+            self.assertTrue('test_sqlalchemy.py' in query.context)
+            self.assertTrue('test_query_recording' in query.context)
 
     def test_helper_api(self):
         self.assertEqual(self.db.metadata, self.db.Model.metadata)
@@ -132,7 +131,7 @@ def tearDown(self):
     def test_model_signals(self):
         recorded = []
         def committed(sender, changes):
-            self.assert_(isinstance(changes, list))
+            self.assertTrue(isinstance(changes, list))
             recorded.extend(changes)
         with sqlalchemy.models_committed.connected_to(committed,
                                                       sender=self.app):
@@ -179,7 +178,7 @@ def test_basic_pagination(self):
         p = sqlalchemy.Pagination(None, 1, 20, 500, [])
         self.assertEqual(p.page, 1)
         self.assertFalse(p.has_prev)
-        self.assert_(p.has_next)
+        self.assertTrue(p.has_next)
         self.assertEqual(p.total, 500)
         self.assertEqual(p.pages, 25)
         self.assertEqual(p.next_num, 2)
@@ -250,17 +249,17 @@ class Baz(db.Model):
         metadata = db.MetaData()
         metadata.reflect(bind=db.get_engine(app, 'foo'))
         self.assertEqual(len(metadata.tables), 1)
-        self.assert_('foo' in metadata.tables)
+        self.assertTrue('foo' in metadata.tables)
 
         metadata = db.MetaData()
         metadata.reflect(bind=db.get_engine(app, 'bar'))
         self.assertEqual(len(metadata.tables), 1)
-        self.assert_('bar' in metadata.tables)
+        self.assertTrue('bar' in metadata.tables)
 
         metadata = db.MetaData()
         metadata.reflect(bind=db.get_engine(app))
         self.assertEqual(len(metadata.tables), 1)
-        self.assert_('baz' in metadata.tables)
+        self.assertTrue('baz' in metadata.tables)
 
         # do the session have the right binds set?
         self.assertEqual(db.get_binds(app), {
@@ -280,7 +279,7 @@ def test_default_query_class(self):
 
         class Parent(db.Model):
             id = db.Column(db.Integer, primary_key=True)
-            children = db.relationship("Child", backref = db.backref("parents", lazy='dynamic'), lazy='dynamic')
+            children = db.relationship("Child", backref = "parents", lazy='dynamic')
         class Child(db.Model):
             id = db.Column(db.Integer, primary_key=True)
             parent_id = db.Column(db.Integer, db.ForeignKey('parent.id'))
@@ -289,8 +288,8 @@ class Child(db.Model):
         c.parent = p
         self.assertEqual(type(Parent.query), sqlalchemy.BaseQuery)
         self.assertEqual(type(Child.query), sqlalchemy.BaseQuery)
-        self.assert_(isinstance(p.children, sqlalchemy.BaseQuery))
-        self.assert_(isinstance(c.parents, sqlalchemy.BaseQuery))
+        self.assertTrue(isinstance(p.children, sqlalchemy.BaseQuery))
+        #self.assertTrue(isinstance(c.parents, sqlalchemy.BaseQuery))
 
 
 class SQLAlchemyIncludesTestCase(unittest.TestCase):
@@ -419,37 +418,6 @@ class FOOBar(db.Model):
             db.session.add(fb)
             assert fb not in db.session  # because a new scope is generated on each call
 
-class StandardSessionTestCase(unittest.TestCase):
-    def test_insert_update_delete(self):
-        # Ensure _SignalTrackingMapperExtension doesn't croak when
-        # faced with a vanilla SQLAlchemy session.
-        #
-        # Verifies that "AttributeError: 'SessionMaker' object has no attribute '_model_changes'"
-        # is not thrown.
-        app = flask.Flask(__name__)
-        app.config['SQLALCHEMY_ENGINE'] = 'sqlite://'
-        app.config['TESTING'] = True
-        db = sqlalchemy.SQLAlchemy(app)
-        Session = sessionmaker(bind=db.engine)
-
-        class QazWsx(db.Model):
-            id = db.Column(db.Integer, primary_key=True)
-            x = db.Column(db.String, default='')
-
-        db.create_all()
-        session = Session()
-        session.add(QazWsx())
-        session.flush() # issues an INSERT.
-        session.expunge_all()
-        qaz_wsx = session.query(QazWsx).first()
-        assert qaz_wsx.x == ''
-        qaz_wsx.x = 'test'
-        session.flush() # issues an UPDATE.
-        session.expunge_all()
-        qaz_wsx = session.query(QazWsx).first()
-        assert qaz_wsx.x == 'test'
-        session.delete(qaz_wsx) # issues a DELETE.
-        assert session.query(QazWsx).first() is None
 
 
 class CommitOnTeardownTestCase(unittest.TestCase):
@@ -500,7 +468,6 @@ def suite():
     suite.addTest(unittest.makeSuite(CommitOnTeardownTestCase))
     if flask.signals_available:
         suite.addTest(unittest.makeSuite(SignallingTestCase))
-    suite.addTest(unittest.makeSuite(StandardSessionTestCase))
     return suite
 
 
