Merge pull request pallets-eco#109 from paxan/bugfix/record-ops-only-for-signalling-sessions

mitsuhiko · mitsuhiko · commit 0d2965c04dee · 2013-01-26T17:24:33.000-08:00
Record operations only for signalling sessions
diff --git a/flask_sqlalchemy.py b/flask_sqlalchemy.py
@@ -166,8 +166,12 @@ def after_update(self, mapper, connection, instance):
         return self._record(mapper, instance, 'update')
 
     def _record(self, mapper, model, operation):
-        pk = tuple(mapper.primary_key_from_instance(model))
-        orm.object_session(model)._model_changes[pk] = (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
 
 
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):
@@ -414,6 +415,38 @@ 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):
 
@@ -463,6 +496,7 @@ def suite():
     suite.addTest(unittest.makeSuite(CommitOnTeardownTestCase))
     if flask.signals_available:
         suite.addTest(unittest.makeSuite(SignallingTestCase))
+    suite.addTest(unittest.makeSuite(StandardSessionTestCase))
     return suite
 
 
