diff --git a/keystone/common/sql/core.py b/keystone/common/sql/core.py index db3db66f74..5785ac74b6 100644 --- a/keystone/common/sql/core.py +++ b/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 diff --git a/keystone/tests/test_sql_core.py b/keystone/tests/test_sql_core.py index 7c9147faea..013eb86ad2 100644 --- a/keystone/tests/test_sql_core.py +++ b/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)