Merge "Upgrade exc_filters for 'engine' argument and connect behavior"

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)