Handle DB2 disconnect

If the DB2 server was restarted, database connections would fail,
requiring a restart of the Keystone server.

This change adds an on-checkout handler for DB2 similar to the
MySQL handler.

Change-Id: Ie53d9606edac408401961fc6526bb7c99baa836a
Closes-Bug: #1231657

keystone/common/sql/core.py

@@ -210,6 +210,21 @@ def mysql_on_checkout(dbapi_conn, connection_rec, connection_proxy):
             raise
 
 
+def db2_on_checkout(engine, dbapi_conn, connection_rec, connection_proxy):
+    """Ensures that DB2 connections checked out of the pool are alive."""
+
+    cursor = dbapi_conn.cursor()
+    try:
+        cursor.execute('select 1 from (values (1)) AS t1')
+    except Exception as e:
+        is_disconnect = engine.dialect.is_disconnect(e, dbapi_conn, cursor)
+        if is_disconnect:
+            LOG.warn(_('Got database server has gone away: %s'), e)
+            raise DisconnectionError("Database server went away")
+        else:
+            raise
+
+
 # Backends
 class Base(object):
     _engine = None
@@ -248,6 +263,9 @@ def new_engine():
 
             if engine.name == 'mysql':
                 sql.event.listen(engine, 'checkout', mysql_on_checkout)
+            elif engine.name == 'ibm_db_sa':
+                callback = functools.partial(db2_on_checkout, engine)
+                sql.event.listen(engine, 'checkout', callback)
 
             return engine
 

keystone/tests/test_sql_core.py

@@ -205,7 +205,7 @@ def cursor(self):
         return self._cursor
 
 
-class TestCheckoutHandler(tests.TestCase):
+class TestMysqlCheckoutHandler(tests.TestCase):
     def _do_on_checkout(self, failwith=None):
         dbapi_conn = FakeDbapiConn(failwith=failwith)
         connection_rec = None
@@ -238,3 +238,55 @@ def test_error(self):
         self.assertRaises(FakeDbapiConn.OperationalError,
                           self._do_on_checkout,
                           failwith=other_exception)
+
+
+class TestDb2CheckoutHandler(tests.TestCase):
+
+    class FakeEngine(object):
+        class Dialect():
+            DISCONNECT_EXCEPTION = Exception()
+
+            @classmethod
+            def is_disconnect(cls, e, *args):
+                return (e is cls.DISCONNECT_EXCEPTION)
+
+        dialect = Dialect()
+
+    def _do_on_checkout(self, failwith=None):
+        engine = self.FakeEngine()
+        dbapi_conn = FakeDbapiConn(failwith=failwith)
+        connection_rec = None
+        connection_proxy = None
+        sql.db2_on_checkout(engine, dbapi_conn, connection_rec,
+                            connection_proxy)
+
+    def test_checkout_success(self):
+        # If call db2_on_checkout and query doesn't raise anything, then no
+        # problems
+
+        # If this doesn't raise then the test is successful.
+        self._do_on_checkout()
+
+    def test_disconnected(self):
+        # If call db2_on_checkout and query raises exception that engine
+        # dialect says is a disconnect problem, then raises DisconnectionError.
+
+        disconnected_exception = self.FakeEngine.Dialect.DISCONNECT_EXCEPTION
+        self.assertRaises(DisconnectionError,
+                          self._do_on_checkout,
+                          failwith=disconnected_exception)
+
+    def test_error(self):
+        # If call db2_on_checkout and query raises an exception that engine
+        # dialect says is not a disconnect problem, then the original error is
+        # raised.
+
+        # fake engine dialect doesn't look for this exception.
+
+        class OtherException(Exception):
+            pass
+
+        other_exception = OtherException()
+        self.assertRaises(OtherException,
+                          self._do_on_checkout,
+                          failwith=other_exception)