Upgrade exc_filters for 'engine' argument and connect behavior

This patch applies upgrades to the sqlalchemy/exc_filters.py and
sqlalchemy/compat/handle_error.py compatibility layers to accommodate
new changes in SQLAlchemy 1.0.  SQLA 1.0 will now route errors that
occur upon connect through the handle_error() event, just like any other,
so that when 1.0 is present we no longer need to use
exc_filters.handle_connect_error; the method becomes a passthrough
as far as running the event handler.  Additionally, SQLAlchemy 1.0
has added the "engine" parameter to ExceptionContext, specifically
to suit the case when the initial connect has failed and there is
no Connection object; the compatibility layer here now emulates
this behavior for SQLAlchemy versions prior to 1.0.

Change-Id: I61714f3c32625a621eaba501d20346519b8b12c7

4 files changed, 61 insertions, 25 deletions

diff --git a/oslo/db/sqlalchemy/compat/__init__.py b/oslo/db/sqlalchemy/compat/__init__.py
index aaaf200..2ffe293 100644
--- a/oslo/db/sqlalchemy/compat/__init__.py
+++ b/oslo/db/sqlalchemy/compat/__init__.py
@@ -23,6 +23,8 @@ 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
-ExceptionContextImpl = _h_err.ExceptionContextImpl
+handle_connect_context = _h_err.handle_connect_context
 
-__all__ = ['engine_connect', 'handle_error', 'ExceptionContextImpl']
+__all__ = [
+    'engine_connect', 'handle_error',
+    'handle_connect_context']
diff --git a/oslo/db/sqlalchemy/compat/handle_error.py b/oslo/db/sqlalchemy/compat/handle_error.py
index 6f0debe..1537480 100644
--- a/oslo/db/sqlalchemy/compat/handle_error.py
+++ b/oslo/db/sqlalchemy/compat/handle_error.py
@@ -16,6 +16,7 @@ http://docs.sqlalchemy.org/en/rel_0_9/core/events.html.
 
 
 """
+import contextlib
 import sys
 
 import six
@@ -35,10 +36,17 @@ def handle_error(engine, listener):
     in order to support safe re-raise of the exception.
 
     """
-
-    if utils.sqla_097:
+    if utils.sqla_100:
         event.listen(engine, "handle_error", listener)
         return
+    elif utils.sqla_097:
+        # ctx.engine added per
+        # https://bitbucket.org/zzzeek/sqlalchemy/issue/3266/
+        def wrap_listener(ctx):
+            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"
@@ -92,7 +100,7 @@ def handle_error(engine, listener):
                     # new handle_error event
                     ctx = ExceptionContextImpl(
                         original_exception, sqlalchemy_exception,
-                        self, cursor, statement,
+                        self.engine, self, cursor, statement,
                         parameters, context, is_disconnect)
 
                     for fn in _oslo_handle_error_events:
@@ -153,8 +161,9 @@ class ExceptionContextImpl(object):
     """
 
     def __init__(self, exception, sqlalchemy_exception,
-                 connection, cursor, statement, parameters,
+                 engine, connection, cursor, statement, parameters,
                  context, is_disconnect):
+        self.engine = engine
         self.connection = connection
         self.sqlalchemy_exception = sqlalchemy_exception
         self.original_exception = exception
@@ -166,7 +175,17 @@ class ExceptionContextImpl(object):
     connection = None
     """The :class:`.Connection` in use during the exception.
 
-    This member is always present.
+    This member is present, except in the case of a failure when
+    first connecting.
+
+
+    """
+
+    engine = None
+    """The :class:`.Engine` in use during the exception.
+
+    This member should always be present, even in the case of a failure
+    when first connecting.
 
     """
 
@@ -247,3 +266,24 @@ 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/compat/utils.py b/oslo/db/sqlalchemy/compat/utils.py
index 7ba83c3..fa6c3e7 100644
--- a/oslo/db/sqlalchemy/compat/utils.py
+++ b/oslo/db/sqlalchemy/compat/utils.py
@@ -19,6 +19,7 @@ _SQLA_VERSION = tuple(
     for num in sqlalchemy.__version__.split(".")
 )
 
+sqla_100 = _SQLA_VERSION >= (1, 0, 0)
 sqla_097 = _SQLA_VERSION >= (0, 9, 7)
 sqla_094 = _SQLA_VERSION >= (0, 9, 4)
 sqla_090 = _SQLA_VERSION >= (0, 9, 0)
diff --git a/oslo/db/sqlalchemy/exc_filters.py b/oslo/db/sqlalchemy/exc_filters.py
index b476482..f997237 100644
--- a/oslo/db/sqlalchemy/exc_filters.py
+++ b/oslo/db/sqlalchemy/exc_filters.py
@@ -314,13 +314,13 @@ def handler(context):
     more specific exception class are attempted first.
 
     """
-    def _dialect_registries(connection):
-        if connection.dialect.name in _registry:
-            yield _registry[connection.dialect.name]
+    def _dialect_registries(engine):
+        if engine.dialect.name in _registry:
+            yield _registry[engine.dialect.name]
         if '*' in _registry:
             yield _registry['*']
 
-    for per_dialect in _dialect_registries(context.connection):
+    for per_dialect in _dialect_registries(context.engine):
         for exc in (
                 context.sqlalchemy_exception,
                 context.original_exception):
@@ -334,7 +334,7 @@ def handler(context):
                                 fn(
                                     exc,
                                     match,
-                                    context.connection.dialect.name,
+                                    context.engine.dialect.name,
                                     context.is_disconnect)
                             except exception.DBConnectionError:
                                 context.is_disconnect = True
@@ -349,18 +349,11 @@ def handle_connect_error(engine):
     """Handle connect error.
 
     Provide a special context that will allow on-connect errors
-    to be raised within the filtering context.
+    to be treated within the filtering context.
+
+    This routine is dependent on SQLAlchemy version, as version 1.0.0
+    provides this functionality natively.
+
     """
-    try:
+    with compat.handle_connect_context(handler, engine):
         return engine.connect()
-    except Exception as e:
-        if isinstance(e, sqla_exc.StatementError):
-            s_exc, orig = e, e.orig
-        else:
-            s_exc, orig = None, e
-
-        ctx = compat.ExceptionContextImpl(
-            orig, s_exc, engine, None,
-            None, None, None, False
-        )
-        handler(ctx)