Fix missing wrap_db_error for Session.execute() method

We should add wrap_db_error for Session.execute() method, because
Query.update() method will call Session.execute() and it is not inside
Session.flush(), so exceptions from Query.update() wouldn't be wrapped.

Add test for Session.flush() wrapper
Add test for Session.execute() wrapper

Fixes bug 1107890

Change-Id: I96894e502f9f279999b61a1e66469e34a07a022d

2 files changed, 67 insertions, 0 deletions

diff --git a/openstack/common/db/sqlalchemy/session.py b/openstack/common/db/sqlalchemy/session.py
index 700273a..06b198e 100644
--- a/openstack/common/db/sqlalchemy/session.py
+++ b/openstack/common/db/sqlalchemy/session.py
@@ -593,6 +593,10 @@ class Session(sqlalchemy.orm.session.Session):
     def flush(self, *args, **kwargs):
         return super(Session, self).flush(*args, **kwargs)
 
+    @wrap_db_error
+    def execute(self, *args, **kwargs):
+        return super(Session, self).execute(*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 d4894ba..10f7e41 100644
--- a/tests/unit/db/sqlalchemy/test_sqlalchemy.py
+++ b/tests/unit/db/sqlalchemy/test_sqlalchemy.py
@@ -23,8 +23,14 @@ try:
 except ImportError:
     HAS_MYSQLDB = False
 
+from sqlalchemy import Column, MetaData, Table, UniqueConstraint
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy import DateTime, Integer
+
 from openstack.common import context
 from openstack.common import exception
+from openstack.common.db import common as common_db
+from openstack.common.db.sqlalchemy import models
 from openstack.common.db.sqlalchemy import session
 from tests import utils as test_utils
 
@@ -70,3 +76,60 @@ class DbPoolTestCase(test_utils.BaseTestCase):
         self.assertEqual(info['kwargs']['max_idle'], 11)
         self.assertEqual(info['kwargs']['min_size'], 21)
         self.assertEqual(info['kwargs']['max_size'], 42)
+
+
+BASE = declarative_base()
+_TABLE_NAME = '__tmp__test__tmp__'
+
+
+class TmpTable(BASE, models.ModelBase):
+    __tablename__ = _TABLE_NAME
+    id = Column(Integer, primary_key=True)
+    foo = Column(Integer)
+
+
+class SessionErrorWrapperTestCase(test_utils.BaseTestCase):
+    def setUp(self):
+        super(SessionErrorWrapperTestCase, self).setUp()
+        meta = MetaData()
+        meta.bind = session.get_engine()
+        test_table = Table(_TABLE_NAME, meta,
+                           Column('id', Integer, primary_key=True,
+                                  nullable=False),
+                           Column('deleted', Integer, default=0),
+                           Column('deleted_at', DateTime),
+                           Column('updated_at', DateTime),
+                           Column('created_at', DateTime),
+                           Column('foo', Integer),
+                           UniqueConstraint('foo', name='uniq_foo'))
+        test_table.create()
+
+    def tearDown(self):
+        super(SessionErrorWrapperTestCase, self).tearDown()
+        meta = MetaData()
+        meta.bind = session.get_engine()
+        test_table = Table(_TABLE_NAME, meta, autoload=True)
+        test_table.drop()
+
+    def test_flush_wrapper(self):
+        tbl = TmpTable()
+        tbl.update({'foo': 10})
+        tbl.save()
+
+        tbl2 = TmpTable()
+        tbl2.update({'foo': 10})
+        self.assertRaises(common_db.DBDuplicateEntry, tbl2.save)
+
+    def test_execute_wrapper(self):
+        _session = session.get_session()
+        with _session.begin():
+            for i in [10, 20]:
+                tbl = TmpTable()
+                tbl.update({'foo': i})
+                tbl.save(session=_session)
+
+            method = _session.query(TmpTable).\
+                filter_by(foo=10).\
+                update
+            self.assertRaises(common_db.DBDuplicateEntry,
+                              method, {'foo': 20})