Record operations only for signalling sessions (reprise)

martijnvermaat · martijnvermaat · commit b698ff05f706 · 2013-06-15T18:03:20.000+02:00
This is a rebase of commit cf40cae (GitHub pull request pallets-eco#109 by @paxan) on current master. The original changes were (I assume accidentally) undone in commit eb208b5 (GitHub pull request pallets-eco#123 by @thrisp). The original commit message is copied below. In some scenarios model classes derived from SQLAlchemy(...).Model may be used with basic SQLAlchemy session instances which aren't extended with signalling functionality. Before this fix, when you try using a vanilla session with such models, you would get exception like this: File ".../flask_sqlalchemy.py", line 163, in after_insert return self._record(mapper, instance, 'insert') File ".../flask_sqlalchemy.py", line 172, in _record s._model_changes[pk] = (model, operation) AttributeError: 'SessionMaker' object has no attribute '_model_changes' With this fix, this problem is resolved.
diff --git a/flask_sqlalchemy.py b/flask_sqlalchemy.py
@@ -219,8 +219,11 @@ def mapper_signal_after_update(self, mapper, connection, target):
 
     @staticmethod
     def _record(mapper, target, operation):
-        pk = tuple(mapper.primary_key_from_instance(target))
-        orm.object_session(target)._model_changes[pk] = (target, operation)
+        s = orm.object_session(target)
+        if isinstance(s, _SignallingSession):
+            pk = tuple(mapper.primary_key_from_instance(target))
+            s._model_changes[pk] = (target, operation)
+
 
 def get_debug_queries():
     """In debug mode Flask-SQLAlchemy will log all the SQL queries sent
diff --git a/test_sqlalchemy.py b/test_sqlalchemy.py
@@ -6,6 +6,7 @@
 from datetime import datetime
 import flask
 from flask.ext import sqlalchemy
+from sqlalchemy.orm import sessionmaker
 
 
 def make_todo_model(db):
@@ -454,6 +455,39 @@ def test_roll_back_on_failure(self):
         self.assertEqual(self.client.get('/').data, '')
 
 
+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
+
+
 def suite():
     suite = unittest.TestSuite()
     suite.addTest(unittest.makeSuite(BasicAppTestCase))
@@ -468,6 +502,7 @@ def suite():
     suite.addTest(unittest.makeSuite(CommitOnTeardownTestCase))
     if flask.signals_available:
         suite.addTest(unittest.makeSuite(SignallingTestCase))
+    suite.addTest(unittest.makeSuite(StandardSessionTestCase))
     return suite
 
 
