Merge "Add _wrap_db_error() support to SessionTransaction.commit()" into stable/icehouse

2 files changed, 58 insertions, 2 deletions

diff --git a/openstack/common/db/sqlalchemy/session.py b/openstack/common/db/sqlalchemy/session.py
index 2f81b690..ff7e6da0 100644
--- a/openstack/common/db/sqlalchemy/session.py
+++ b/openstack/common/db/sqlalchemy/session.py
@@ -432,9 +432,11 @@ def _wrap_db_error(f):
     def _wrap(self, *args, **kwargs):
         try:
             assert issubclass(
-                self.__class__, sqlalchemy.orm.session.Session
+                self.__class__, (
+                    sqlalchemy.orm.session.Session, SessionTransactionWrapper)
             ), ('_wrap_db_error() can only be applied to methods of '
-                'subclasses of sqlalchemy.orm.session.Session.')
+                'subclasses of sqlalchemy.orm.session.Session or '
+                ' SessionTransactionWrapper')
 
             return f(self, *args, **kwargs)
         except UnicodeEncodeError:
@@ -716,6 +718,25 @@ class Session(sqlalchemy.orm.session.Session):
     def commit(self, *args, **kwargs):
         return super(Session, self).commit(*args, **kwargs)
 
+    def begin(self, **kw):
+        trans = super(Session, self).begin(**kw)
+        trans.__class__ = SessionTransactionWrapper
+        return trans
+
+
+class SessionTransactionWrapper(sqlalchemy.orm.session.SessionTransaction):
+    @property
+    def bind(self):
+        return self.session.bind
+
+    @_wrap_db_error
+    def commit(self, *args, **kwargs):
+        return super(SessionTransactionWrapper, self).commit(*args, **kwargs)
+
+    @_wrap_db_error
+    def rollback(self, *args, **kwargs):
+        return super(SessionTransactionWrapper, self).rollback(*args, **kwargs)
+
 
 def get_maker(engine, autocommit=True, expire_on_commit=False):
     """Return a SQLAlchemy sessionmaker using the given engine."""
diff --git a/tests/unit/db/sqlalchemy/test_sqlalchemy.py b/tests/unit/db/sqlalchemy/test_sqlalchemy.py
index 54840647..0dce8afe 100644
--- a/tests/unit/db/sqlalchemy/test_sqlalchemy.py
+++ b/tests/unit/db/sqlalchemy/test_sqlalchemy.py
@@ -123,6 +123,41 @@ class SessionErrorWrapperTestCase(test_base.DbTestCase):
                         _raise_if_deadlock_error):
                     self.assertRaises(db_exc.DBDeadlock, _session.commit)
 
+    def test_context_commit_wrapper(self):
+        # test a commit that raises at the point of commit within
+        # a context manager "with session.begin():"
+
+        _session = self.sessionmaker()
+
+        trans = _session.begin()
+        # patch SQLAlchemy connection._commit_impl to
+        # raise a deadlock exception during commit only.
+        deadlock_on_commit = mock.Mock(
+            side_effect=sqla_exc.OperationalError(
+                "simulated deadlock", None, None))
+
+        # patch our own _raise_if_deadlock_error filter to treat
+        # this OperationalError as a DBDeadlock.
+        def _raise_if_deadlock_error(operational_error, engine_name):
+            self.assertIsInstance(
+                operational_error, sqla_exc.OperationalError)
+            self.assertTrue(
+                "simulated deadlock" in str(operational_error))
+            raise db_exc.DBDeadlock(operational_error)
+
+        connection = _session.connection()
+
+        with mock.patch.object(
+                connection, '_commit_impl', deadlock_on_commit):
+            def run_as_ctxmanager():
+                with trans:
+                    pass
+
+            with mock.patch.object(
+                    session, "_raise_if_deadlock_error",
+                    _raise_if_deadlock_error):
+                self.assertRaises(db_exc.DBDeadlock, run_as_ctxmanager)
+
     def test_execute_wrapper(self):
         _session = self.sessionmaker()
         with _session.begin():