From 8513fb588f7f2288cb15dd2bd1e4066b60e139b3 Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Tue, 29 Jul 2014 14:06:43 -0400
Subject: [PATCH] - The exception wrapping system for DBAPI errors can now
 accommodate non-standard DBAPI exceptions, such as the psycopg2
 TransactionRollbackError.  These exceptions will now be raised using the
 closest available subclass in ``sqlalchemy.exc``, in the case of
 TransactionRollbackError, ``sqlalchemy.exc.OperationalError``. fixes #3075

---
 doc/build/changelog/changelog_09.rst    | 11 ++++++++++
 lib/sqlalchemy/exc.py                   |  9 ++++++---
 test/dialect/postgresql/test_dialect.py | 10 +++++++++
 test/engine/test_execute.py             | 27 ++++++++++++++++++++++++-
 4 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/doc/build/changelog/changelog_09.rst b/doc/build/changelog/changelog_09.rst
index c63ed7fbb0..afc51f2400 100644
--- a/doc/build/changelog/changelog_09.rst
+++ b/doc/build/changelog/changelog_09.rst
@@ -13,6 +13,17 @@
 .. changelog::
     :version: 0.9.8
 
+    .. change::
+        :tags: bug, postgresql
+        :versions: 1.0.0
+        :tickets: 3075
+
+        The exception wrapping system for DBAPI errors can now accommodate
+        non-standard DBAPI exceptions, such as the psycopg2
+        TransactionRollbackError.  These exceptions will now be raised
+        using the closest available subclass in ``sqlalchemy.exc``, in the
+        case of TransactionRollbackError, ``sqlalchemy.exc.OperationalError``.
+
     .. change::
         :tags: bug, sql
         :versions: 1.0.0
diff --git a/lib/sqlalchemy/exc.py b/lib/sqlalchemy/exc.py
index 7d333fc01e..a82bae33fd 100644
--- a/lib/sqlalchemy/exc.py
+++ b/lib/sqlalchemy/exc.py
@@ -294,9 +294,12 @@ class DBAPIError(StatementError):
                     statement, params, orig
                 )
 
-            name, glob = orig.__class__.__name__, globals()
-            if name in glob and issubclass(glob[name], DBAPIError):
-                cls = glob[name]
+            glob = globals()
+            for super_ in orig.__class__.__mro__:
+                name = super_.__name__
+                if name in glob and issubclass(glob[name], DBAPIError):
+                    cls = glob[name]
+                    break
 
         return cls(statement, params, orig, connection_invalidated)
 
diff --git a/test/dialect/postgresql/test_dialect.py b/test/dialect/postgresql/test_dialect.py
index a0f9e68959..4ec5c25741 100644
--- a/test/dialect/postgresql/test_dialect.py
+++ b/test/dialect/postgresql/test_dialect.py
@@ -65,6 +65,16 @@ class MiscTest(fixtures.TestBase, AssertsExecutionResults, AssertsCompiledSQL):
         assert testing.db.dialect.dbapi.__version__.\
             startswith(".".join(str(x) for x in v))
 
+    @testing.only_on('postgresql+psycopg2', 'psycopg2-specific feature')
+    def test_psycopg2_non_standard_err(self):
+        from psycopg2.extensions import TransactionRollbackError
+        import psycopg2
+
+        exception = exc.DBAPIError.instance(
+            "some statement", {}, TransactionRollbackError("foo"),
+            psycopg2.Error)
+        assert isinstance(exception, exc.OperationalError)
+
     # currently not passing with pg 9.3 that does not seem to generate
     # any notices here, would rather find a way to mock this
     @testing.only_on('postgresql+psycopg2', 'psycopg2-specific feature')
diff --git a/test/engine/test_execute.py b/test/engine/test_execute.py
index 33e116cdfe..1cd5978217 100644
--- a/test/engine/test_execute.py
+++ b/test/engine/test_execute.py
@@ -20,7 +20,7 @@ from sqlalchemy.engine import result as _result, default
 from sqlalchemy.engine.base import Engine
 from sqlalchemy.testing import fixtures
 from sqlalchemy.testing.mock import Mock, call, patch
-from contextlib import contextmanager
+from contextlib import contextmanager, nested
 
 users, metadata, users_autoinc = None, None, None
 class ExecuteTest(fixtures.TestBase):
@@ -235,6 +235,31 @@ class ExecuteTest(fixtures.TestBase):
             )
             eq_(is_disconnect.call_count, 0)
 
+    def test_exception_wrapping_non_standard_dbapi_error(self):
+        class DBAPIError(Exception):
+            pass
+
+        class OperationalError(DBAPIError):
+            pass
+
+        class NonStandardException(OperationalError):
+            pass
+
+        with nested(
+            patch.object(testing.db.dialect, "dbapi", Mock(Error=DBAPIError)),
+            patch.object(
+                testing.db.dialect, "is_disconnect",
+                lambda *arg: False),
+            patch.object(
+                testing.db.dialect, "do_execute",
+                Mock(side_effect=NonStandardException)),
+        ):
+            with testing.db.connect() as conn:
+                assert_raises(
+                    tsa.exc.OperationalError,
+                    conn.execute, "select 1"
+                )
+
 
     def test_exception_wrapping_non_dbapi_statement(self):
         class MyType(TypeDecorator):
-- 
2.47.3