Ensure DBConnectionError is raised on failed revalidate

The second call to connection.scalar() inside of _connect_ping_listener()
is not subject to the exc_filters system under SQLAlchemy 0.9 and earlier,
in the case that the attempt to revalidate within fails.  This causes
the exception to be propagated outwards as the original
OperationalError (assuming that's what the DBAPI raised, as is typical),
rather than it being wrapped again in DBConnectionError.

SQLAlchemy 1.0 has altered this behavior such that the handle_error()
listener is invoked correctly for revalidiation attempts (it was being
invoked before, but without the correct state), as well as
for initial connect attempts, as part of
https://bitbucket.org/zzzeek/sqlalchemy/issue/3266/.  For 0.9 and earlier,
we here backport this system into oslo/db/sqlalchemy/compat.py,
including that we redefine the connection pool's handler to re-throw
the original dbapi.Error instance directly, then move the handling
of this error in terms of SQLAlchemy wrapping and events into
the Engine and Connection.   The approach here works back to
SQLAlchemy 0.8, and is conditional based on whether or not
SQLAlchemy 1.0 is detected.

Change-Id: I455e957b043318a8500909c66cadde53228b52d6

5 files changed, 379 insertions, 135 deletions

diff --git a/oslo_db/sqlalchemy/compat/__init__.py b/oslo_db/sqlalchemy/compat/__init__.py
index b49d5c4..3dc29ac 100644
--- a/oslo_db/sqlalchemy/compat/__init__.py
+++ b/oslo_db/sqlalchemy/compat/__init__.py
@@ -23,8 +23,6 @@ from oslo_db.sqlalchemy.compat import handle_error as _h_err
 # flake8 won't let me import handle_error directly
 engine_connect = _e_conn.engine_connect
 handle_error = _h_err.handle_error
-handle_connect_context = _h_err.handle_connect_context
 
 __all__ = [
-    'engine_connect', 'handle_error',
-    'handle_connect_context']
+    'engine_connect', 'handle_error']
diff --git a/oslo_db/sqlalchemy/compat/handle_error.py b/oslo_db/sqlalchemy/compat/handle_error.py
index 7e476a0..5169857 100644
--- a/oslo_db/sqlalchemy/compat/handle_error.py
+++ b/oslo_db/sqlalchemy/compat/handle_error.py
@@ -16,10 +16,10 @@ http://docs.sqlalchemy.org/en/rel_0_9/core/events.html.
 
 
 """
-import contextlib
 import sys
 
 import six
+from sqlalchemy.engine import base as engine_base
 from sqlalchemy.engine import Engine
 from sqlalchemy import event
 from sqlalchemy import exc as sqla_exc
@@ -39,117 +39,294 @@ def handle_error(engine, listener):
     if utils.sqla_100:
         event.listen(engine, "handle_error", listener)
         return
-    elif utils.sqla_097:
+
+    assert isinstance(engine, Engine), \
+        "engine argument must be an Engine instance, not a Connection"
+
+    if not utils.sqla_097:
+        _rework_handle_exception_for_events(engine)
+        engine._oslo_handle_error_events.append(listener)
+
+    _rework_connect_and_revalidate_for_events(engine)
+
+    if utils.sqla_097:
         # ctx.engine added per
         # https://bitbucket.org/zzzeek/sqlalchemy/issue/3266/
         def wrap_listener(ctx):
-            ctx.engine = ctx.connection.engine
+            if isinstance(ctx, engine_base.ExceptionContextImpl):
+                ctx.engine = ctx.connection.engine
             return listener(ctx)
         event.listen(engine, "handle_error", wrap_listener)
-        return
-
-    assert isinstance(engine, Engine), \
-        "engine argument must be an Engine instance, not a Connection"
 
-    # use a Connection-wrapper class to wrap _handle_dbapi_exception.
-    if not getattr(engine._connection_cls,
-                   '_oslo_handle_error_wrapper', False):
-        engine._oslo_handle_error_events = []
 
-        class Connection(engine._connection_cls):
-            _oslo_handle_error_wrapper = True
+def _rework_handle_exception_for_events(engine):
+    """Patch the _handle_dbapi_error() system on Connection.
 
-            def _handle_dbapi_exception(self, e, statement, parameters,
-                                        cursor, context):
+    This allows the 0.9.7-style handle_error() event to be available on
+    the Connection object.
 
-                try:
-                    super(Connection, self)._handle_dbapi_exception(
-                        e, statement, parameters, cursor, context)
-                except Exception as reraised_exception:
-                    # all versions:
-                    #   _handle_dbapi_exception reraises all DBAPI errors
-                    # 0.8 and above:
-                    #   reraises all errors unconditionally
-                    pass
-                else:
-                    # 0.7.8:
-                    #   _handle_dbapi_exception does not unconditionally
-                    #   re-raise
-                    reraised_exception = e
-
-                _oslo_handle_error_events = getattr(
-                    self.engine,
-                    '_oslo_handle_error_events',
-                    False)
-
-                newraise = None
-                if _oslo_handle_error_events:
-                    if isinstance(reraised_exception,
-                                  sqla_exc.StatementError):
-                        sqlalchemy_exception = reraised_exception
-                        original_exception = sqlalchemy_exception.orig
-                        self._is_disconnect = is_disconnect = (
-                            isinstance(sqlalchemy_exception,
-                                       sqla_exc.DBAPIError)
-                            and sqlalchemy_exception.connection_invalidated)
-                    else:
-                        sqlalchemy_exception = None
-                        original_exception = reraised_exception
-                        is_disconnect = False
-
-                    # new handle_error event
-                    ctx = ExceptionContextImpl(
-                        original_exception, sqlalchemy_exception,
-                        self.engine, self, cursor, statement,
-                        parameters, context, is_disconnect)
-
-                    for fn in _oslo_handle_error_events:
-                        try:
-                            # handler returns an exception;
-                            # call next handler in a chain
-                            per_fn = fn(ctx)
-                            if per_fn is not None:
-                                ctx.chained_exception = newraise = per_fn
-                        except Exception as _raised:
-                            # handler raises an exception - stop processing
-                            newraise = _raised
-                            break
-
-                    if sqlalchemy_exception and \
-                            self._is_disconnect != ctx.is_disconnect:
-
-                        if not ctx.is_disconnect:
-                            raise NotImplementedError(
-                                "Can't reset 'disconnect' status of exception "
-                                "once it is set with this version of "
-                                "SQLAlchemy")
-
-                        sqlalchemy_exception.connection_invalidated = \
-                            self._is_disconnect = ctx.is_disconnect
-                        if self._is_disconnect:
-                            self._do_disconnect(e)
-
-                if newraise:
-                    six.reraise(type(newraise), newraise, sys.exc_info()[2])
-                else:
-                    six.reraise(type(reraised_exception),
-                                reraised_exception, sys.exc_info()[2])
-
-            def _do_disconnect(self, e):
-                del self._is_disconnect
-                if utils.sqla_094:
-                    dbapi_conn_wrapper = self.connection
-                    self.engine.pool._invalidate(dbapi_conn_wrapper, e)
-                    self.invalidate(e)
+    """
+    engine._oslo_handle_error_events = []
+
+    class Connection(engine._connection_cls):
+        def _handle_dbapi_exception(self, e, statement, parameters,
+                                    cursor, context):
+
+            try:
+                super(Connection, self)._handle_dbapi_exception(
+                    e, statement, parameters, cursor, context)
+            except Exception as reraised_exception:
+                # all versions:
+                #   _handle_dbapi_exception reraises all DBAPI errors
+                # 0.8 and above:
+                #   reraises all errors unconditionally
+                pass
+            else:
+                # 0.7.8:
+                #   _handle_dbapi_exception does not unconditionally
+                #   re-raise
+                reraised_exception = e
+
+            _oslo_handle_error_events = getattr(
+                self.engine,
+                '_oslo_handle_error_events',
+                False)
+
+            newraise = None
+            if _oslo_handle_error_events:
+                if isinstance(reraised_exception,
+                              sqla_exc.StatementError):
+                    sqlalchemy_exception = reraised_exception
+                    original_exception = sqlalchemy_exception.orig
+                    self._is_disconnect = is_disconnect = (
+                        isinstance(sqlalchemy_exception,
+                                   sqla_exc.DBAPIError)
+                        and sqlalchemy_exception.connection_invalidated)
                 else:
-                    dbapi_conn_wrapper = self.connection
-                    self.invalidate(e)
-                    if not hasattr(dbapi_conn_wrapper, '_pool') or \
-                            dbapi_conn_wrapper._pool is self.engine.pool:
-                        self.engine.dispose()
+                    sqlalchemy_exception = None
+                    original_exception = reraised_exception
+                    is_disconnect = False
+
+                # new handle_error event
+                ctx = ExceptionContextImpl(
+                    original_exception, sqlalchemy_exception,
+                    self.engine, self, cursor, statement,
+                    parameters, context, is_disconnect)
+
+                for fn in _oslo_handle_error_events:
+                    try:
+                        # handler returns an exception;
+                        # call next handler in a chain
+                        per_fn = fn(ctx)
+                        if per_fn is not None:
+                            ctx.chained_exception = newraise = per_fn
+                    except Exception as _raised:
+                        # handler raises an exception - stop processing
+                        newraise = _raised
+                        break
+
+                if sqlalchemy_exception and \
+                        self._is_disconnect != ctx.is_disconnect:
+
+                    if not ctx.is_disconnect:
+                        raise NotImplementedError(
+                            "Can't reset 'disconnect' status of exception "
+                            "once it is set with this version of "
+                            "SQLAlchemy")
+
+                    sqlalchemy_exception.connection_invalidated = \
+                        self._is_disconnect = ctx.is_disconnect
+                    if self._is_disconnect:
+                        self._do_disconnect(e)
+
+            if newraise:
+                six.reraise(type(newraise), newraise, sys.exc_info()[2])
+            else:
+                six.reraise(type(reraised_exception),
+                            reraised_exception, sys.exc_info()[2])
+
+        def _do_disconnect(self, e):
+            del self._is_disconnect
+            if utils.sqla_094:
+                dbapi_conn_wrapper = self.connection
+                self.engine.pool._invalidate(dbapi_conn_wrapper, e)
+                self.invalidate(e)
+            else:
+                dbapi_conn_wrapper = self.connection
+                self.invalidate(e)
+                if not hasattr(dbapi_conn_wrapper, '_pool') or \
+                        dbapi_conn_wrapper._pool is self.engine.pool:
+                    self.engine.dispose()
+
+    engine._connection_cls = Connection
+
+
+def _rework_connect_and_revalidate_for_events(engine):
+    """Patch the _revalidate_connection() system on Connection.
+
+    This applies 1.0's _revalidate_connection() approach into an 0.9
+    version of SQLAlchemy, and consists of three steps:
+
+    1. wrap the pool._creator function, which in 0.9 has a local
+    call to sqlalchemy.exc.DBAPIError.instance(), so that this exception is
+    again unwrapped back to the original DBAPI-specific Error, then raise
+    that.  This is essentially the same as if the dbapi.connect() isn't
+    wrapped in the first place, which is how SQLAlchemy 1.0 now functions.
+
+    2. patch the Engine object's raw_connection() method.  In SQLAlchemy 1.0,
+    this is now where the error wrapping occurs when a pool connect attempt
+    is made.  Here, when raw_connection() is called without a hosting
+    Connection, we send exception raises to
+    _handle_dbapi_exception_noconnection(), here copied from SQLAlchemy
+    1.0, which is an alternate version of Connection._handle_dbapi_exception()
+    tailored for an initial connect failure when there is no
+    Connection object being dealt with.  This allows the error handler
+    events to be called.
+
+    3. patch the Connection class to follow 1.0's behavior for
+    _revalidate_connection(); here, the call to engine.raw_connection()
+    will pass the raised error to Connection._handle_dbapi_exception(),
+    again allowing error handler events to be called.
 
-        engine._connection_cls = Connection
-    engine._oslo_handle_error_events.append(listener)
+    """
+
+    _orig_connect = engine.pool._creator
+
+    def connect():
+        try:
+            return _orig_connect()
+        except sqla_exc.DBAPIError as err:
+            original_exception = err.orig
+            raise original_exception
+    engine.pool._creator = connect
+
+    self = engine
+
+    def contextual_connect(close_with_result=False, **kwargs):
+        return self._connection_cls(
+            self,
+            self._wrap_pool_connect(self.pool.connect, None),
+            close_with_result=close_with_result,
+            **kwargs)
+
+    def _wrap_pool_connect(fn, connection):
+        dialect = self.dialect
+        try:
+            return fn()
+        except dialect.dbapi.Error as e:
+            if connection is None:
+                _handle_dbapi_exception_noconnection(
+                    e, dialect, self)
+            else:
+                six.reraise(*sys.exc_info())
+
+    def raw_connection(_connection=None):
+        return self._wrap_pool_connect(
+            self.pool.unique_connection, _connection)
+
+    engine.contextual_connect = contextual_connect
+    engine._wrap_pool_connect = _wrap_pool_connect
+    engine.raw_connection = raw_connection
+
+    class Connection(engine._connection_cls):
+
+        @property
+        def connection(self):
+            "The underlying DB-API connection managed by this Connection."
+            try:
+                return self.__connection
+            except AttributeError:
+                try:
+                    return self._revalidate_connection()
+                except Exception as e:
+                    self._handle_dbapi_exception(e, None, None, None, None)
+
+        def _handle_dbapi_exception(self,
+                                    e,
+                                    statement,
+                                    parameters,
+                                    cursor,
+                                    context):
+            if self.invalidated:
+                # 0.9's _handle_dbapi_exception() can't handle
+                # a Connection that is invalidated already, meaning
+                # its "__connection" attribute is not set.  So if we are
+                # in that case, call our "no connection" invalidator.
+                # this is fine as we are only supporting handle_error listeners
+                # that are applied at the engine level.
+                _handle_dbapi_exception_noconnection(
+                    e, self.dialect, self.engine)
+            else:
+                super(Connection, self)._handle_dbapi_exception(
+                    e, statement, parameters, cursor, context)
+
+        def _revalidate_connection(self):
+            if self._Connection__can_reconnect and self._Connection__invalid:
+                if self._Connection__transaction is not None:
+                    raise sqla_exc.InvalidRequestError(
+                        "Can't reconnect until invalid "
+                        "transaction is rolled back")
+                self._Connection__connection = self.engine.raw_connection(
+                    _connection=self)
+                self._Connection__invalid = False
+                return self._Connection__connection
+            raise sqla_exc.ResourceClosedError("This Connection is closed")
+
+    engine._connection_cls = Connection
+
+
+def _handle_dbapi_exception_noconnection(e, dialect, engine):
+
+    exc_info = sys.exc_info()
+
+    is_disconnect = dialect.is_disconnect(e, None, None)
+
+    should_wrap = isinstance(e, dialect.dbapi.Error)
+
+    if should_wrap:
+        sqlalchemy_exception = sqla_exc.DBAPIError.instance(
+            None,
+            None,
+            e,
+            dialect.dbapi.Error,
+            connection_invalidated=is_disconnect)
+    else:
+        sqlalchemy_exception = None
+
+    newraise = None
+
+    ctx = ExceptionContextImpl(
+        e, sqlalchemy_exception, engine, None, None, None,
+        None, None, is_disconnect)
+
+    if hasattr(engine, '_oslo_handle_error_events'):
+        fns = engine._oslo_handle_error_events
+    else:
+        fns = engine.dispatch.handle_error
+    for fn in fns:
+        try:
+            # handler returns an exception;
+            # call next handler in a chain
+            per_fn = fn(ctx)
+            if per_fn is not None:
+                ctx.chained_exception = newraise = per_fn
+        except Exception as _raised:
+            # handler raises an exception - stop processing
+            newraise = _raised
+            break
+
+    if sqlalchemy_exception and \
+            is_disconnect != ctx.is_disconnect:
+        sqlalchemy_exception.connection_invalidated = \
+            is_disconnect = ctx.is_disconnect
+
+    if newraise:
+        six.reraise(type(newraise), newraise, exc_info[2])
+    elif should_wrap:
+        six.reraise(
+            type(sqlalchemy_exception), sqlalchemy_exception, exc_info[2])
+    else:
+        six.reraise(*exc_info)
 
 
 class ExceptionContextImpl(object):
@@ -266,24 +443,3 @@ class ExceptionContextImpl(object):
     :meth:`.ConnectionEvents.handle_error` handler.
 
     """
-
-
-@contextlib.contextmanager
-def handle_connect_context(handler, engine):
-    """Wrap connect() routines with a "handle error" context."""
-    try:
-        yield
-    except Exception as e:
-        if utils.sqla_100:
-            raise
-
-        if isinstance(e, sqla_exc.StatementError):
-            s_exc, orig = e, e.orig
-        else:
-            s_exc, orig = None, e
-
-        ctx = ExceptionContextImpl(
-            orig, s_exc, engine, None, None,
-            None, None, None, False
-        )
-        handler(ctx)
diff --git a/oslo_db/sqlalchemy/exc_filters.py b/oslo_db/sqlalchemy/exc_filters.py
index aa5b55a..5df1eb7 100644
--- a/oslo_db/sqlalchemy/exc_filters.py
+++ b/oslo_db/sqlalchemy/exc_filters.py
@@ -346,14 +346,10 @@ def register_engine(engine):
 
 
 def handle_connect_error(engine):
-    """Handle connect error.
+    """Connect to the engine, including handle_error handlers.
 
-    Provide a special context that will allow on-connect errors
-    to be treated within the filtering context.
-
-    This routine is dependent on SQLAlchemy version, as version 1.0.0
-    provides this functionality natively.
+    The compat library now builds this into the engine.connect()
+    system as per SQLAlchemy 1.0's behavior.
 
     """
-    with compat.handle_connect_context(handler, engine):
-        return engine.connect()
+    return engine.connect()
diff --git a/oslo_db/sqlalchemy/session.py b/oslo_db/sqlalchemy/session.py
index 24bf31d..ef253b9 100644
--- a/oslo_db/sqlalchemy/session.py
+++ b/oslo_db/sqlalchemy/session.py
@@ -576,7 +576,7 @@ def _test_connection(engine, max_retries, retry_interval):
     de_ref = None
     for attempt in attempts:
         try:
-            return exc_filters.handle_connect_error(engine)
+            return engine.connect()
         except exception.DBConnectionError as de:
             msg = _LW('SQL connection failed. %s attempts left.')
             LOG.warning(msg, max_retries - attempt)
diff --git a/oslo_db/tests/sqlalchemy/test_exc_filters.py b/oslo_db/tests/sqlalchemy/test_exc_filters.py
index b2845a2..edab9d6 100644
--- a/oslo_db/tests/sqlalchemy/test_exc_filters.py
+++ b/oslo_db/tests/sqlalchemy/test_exc_filters.py
@@ -94,13 +94,16 @@ class TestsExceptionFilter(_SQLAExceptionMatcher, oslo_test_base.BaseTestCase):
         self.engine.connect().close()  # initialize
 
     @contextlib.contextmanager
-    def _dbapi_fixture(self, dialect_name):
+    def _dbapi_fixture(self, dialect_name, is_disconnect=False):
         engine = self.engine
         with test_utils.nested(
             mock.patch.object(engine.dialect.dbapi,
                               "Error",
                               self.Error),
             mock.patch.object(engine.dialect, "name", dialect_name),
+            mock.patch.object(engine.dialect,
+                              "is_disconnect",
+                              lambda *args: is_disconnect)
         ):
             yield
 
@@ -846,3 +849,94 @@ class TestDBConnectRetry(TestsExceptionFilter):
             self.OperationalError("blah blah -39981 blah blah"),
             2, 3
         )
+
+
+class TestDBConnectPingWrapping(TestsExceptionFilter):
+
+    def setUp(self):
+        super(TestDBConnectPingWrapping, self).setUp()
+        compat.engine_connect(self.engine, session._connect_ping_listener)
+
+    @contextlib.contextmanager
+    def _fixture(
+            self, dialect_name, exception, good_conn_count,
+            is_disconnect=True):
+        engine = self.engine
+
+        # empty out the connection pool
+        engine.dispose()
+
+        connect_fn = engine.dialect.connect
+        real_do_execute = engine.dialect.do_execute
+
+        counter = itertools.count(1)
+
+        def cant_execute(*arg, **kw):
+            value = next(counter)
+            if value > good_conn_count:
+                raise exception
+            else:
+                return real_do_execute(*arg, **kw)
+
+        def cant_connect(*arg, **kw):
+            value = next(counter)
+            if value > good_conn_count:
+                raise exception
+            else:
+                return connect_fn(*arg, **kw)
+
+        with self._dbapi_fixture(dialect_name, is_disconnect=is_disconnect):
+            with mock.patch.object(engine.dialect, "connect", cant_connect):
+                with mock.patch.object(
+                        engine.dialect, "do_execute", cant_execute):
+                    yield
+
+    def _test_ping_listener_disconnected(
+            self, dialect_name, exc_obj, is_disconnect=True):
+        with self._fixture(dialect_name, exc_obj, 3, is_disconnect):
+            conn = self.engine.connect()
+            self.assertEqual(conn.scalar(sqla.select([1])), 1)
+            conn.close()
+
+        with self._fixture(dialect_name, exc_obj, 1, is_disconnect):
+            self.assertRaises(
+                exception.DBConnectionError,
+                self.engine.connect
+            )
+            self.assertRaises(
+                exception.DBConnectionError,
+                self.engine.connect
+            )
+            self.assertRaises(
+                exception.DBConnectionError,
+                self.engine.connect
+            )
+
+        with self._fixture(dialect_name, exc_obj, 1, is_disconnect):
+            self.assertRaises(
+                exception.DBConnectionError,
+                self.engine.contextual_connect
+            )
+            self.assertRaises(
+                exception.DBConnectionError,
+                self.engine.contextual_connect
+            )
+            self.assertRaises(
+                exception.DBConnectionError,
+                self.engine.contextual_connect
+            )
+
+    def test_mysql_w_disconnect_flag(self):
+        for code in [2002, 2003, 2002]:
+            self._test_ping_listener_disconnected(
+                "mysql",
+                self.OperationalError('%d MySQL server has gone away' % code)
+            )
+
+    def test_mysql_wo_disconnect_flag(self):
+        for code in [2002, 2003]:
+            self._test_ping_listener_disconnected(
+                "mysql",
+                self.OperationalError('%d MySQL server has gone away' % code),
+                is_disconnect=False
+            )