Reworked internal exception raises that emit

a rollback() before re-raising, so that the stack
trace is preserved from sys.exc_info() before entering
the rollback.  This so that the traceback is preserved
when using coroutine frameworks which may have switched
contexts before the rollback function returns.
[ticket:2703]

7 files changed, 89 insertions, 57 deletions

diff --git a/lib/sqlalchemy/engine/base.py b/lib/sqlalchemy/engine/base.py
index e40af6219..b4c9b1e1c 100644
--- a/lib/sqlalchemy/engine/base.py
+++ b/lib/sqlalchemy/engine/base.py
@@ -462,7 +462,6 @@ class Connection(Connectable):
             self.engine.dialect.do_begin(self.connection)
         except Exception, e:
             self._handle_dbapi_exception(e, None, None, None, None)
-            raise
 
     def _rollback_impl(self):
         if self._has_events:
@@ -476,7 +475,6 @@ class Connection(Connectable):
                 self.__transaction = None
             except Exception, e:
                 self._handle_dbapi_exception(e, None, None, None, None)
-                raise
         else:
             self.__transaction = None
 
@@ -491,7 +489,6 @@ class Connection(Connectable):
             self.__transaction = None
         except Exception, e:
             self._handle_dbapi_exception(e, None, None, None, None)
-            raise
 
     def _savepoint_impl(self, name=None):
         if self._has_events:
@@ -693,7 +690,6 @@ class Connection(Connectable):
                                 dialect, self, conn)
         except Exception, e:
             self._handle_dbapi_exception(e, None, None, None, None)
-            raise
 
         ret = ctx._exec_default(default, None)
         if self.should_close_with_result:
@@ -830,7 +826,6 @@ class Connection(Connectable):
             self._handle_dbapi_exception(e,
                         str(statement), parameters,
                         None, None)
-            raise
 
         if context.compiled:
             context.pre_exec()
@@ -877,7 +872,6 @@ class Connection(Connectable):
                                 parameters,
                                 cursor,
                                 context)
-            raise
 
         if self._has_events:
             self.dispatch.after_cursor_execute(self, cursor,
@@ -952,7 +946,6 @@ class Connection(Connectable):
                                 parameters,
                                 cursor,
                                 None)
-            raise
 
     def _safe_close_cursor(self, cursor):
         """Close the given cursor, catching exceptions
@@ -983,22 +976,21 @@ class Connection(Connectable):
                                     cursor,
                                     context):
 
+        exc_info = sys.exc_info()
+
         if not self._is_disconnect:
             self._is_disconnect = isinstance(e, self.dialect.dbapi.Error) and \
                 not self.closed and \
                 self.dialect.is_disconnect(e, self.__connection, cursor)
 
         if self._reentrant_error:
-            # Py3K
-            #raise exc.DBAPIError.instance(statement, parameters, e,
-            #                               self.dialect.dbapi.Error) from e
-            # Py2K
-            raise exc.DBAPIError.instance(statement,
+            util.raise_from_cause(
+                        exc.DBAPIError.instance(statement,
                                             parameters,
                                             e,
-                                            self.dialect.dbapi.Error), \
-                                            None, sys.exc_info()[2]
-            # end Py2K
+                                            self.dialect.dbapi.Error),
+                        exc_info
+                        )
         self._reentrant_error = True
         try:
             # non-DBAPI error - if we already got a context,
@@ -1021,26 +1013,18 @@ class Connection(Connectable):
                     self._safe_close_cursor(cursor)
                 self._autorollback()
 
-            if not should_wrap:
-                return
-
-            # Py3K
-            #raise exc.DBAPIError.instance(
-            #                        statement,
-            #                        parameters,
-            #                        e,
-            #                        self.dialect.dbapi.Error,
-            #                        connection_invalidated=self._is_disconnect) \
-            #                        from e
-            # Py2K
-            raise exc.DBAPIError.instance(
-                                    statement,
-                                    parameters,
-                                    e,
-                                    self.dialect.dbapi.Error,
-                                    connection_invalidated=self._is_disconnect), \
-                                    None, sys.exc_info()[2]
-            # end Py2K
+            if should_wrap:
+                util.raise_from_cause(
+                                    exc.DBAPIError.instance(
+                                        statement,
+                                        parameters,
+                                        e,
+                                        self.dialect.dbapi.Error,
+                                        connection_invalidated=self._is_disconnect),
+                                    exc_info
+                                )
+
+            util.reraise(*exc_info)
 
         finally:
             del self._reentrant_error
@@ -1115,8 +1099,8 @@ class Connection(Connectable):
             trans.commit()
             return ret
         except:
-            trans.rollback()
-            raise
+            with util.safe_reraise():
+                trans.rollback()
 
     def run_callable(self, callable_, *args, **kwargs):
         """Given a callable object or function, execute it, passing
@@ -1222,8 +1206,8 @@ class Transaction(object):
             try:
                 self.commit()
             except:
-                self.rollback()
-                raise
+                with util.safe_reraise():
+                    self.rollback()
         else:
             self.rollback()
 
@@ -1548,8 +1532,8 @@ class Engine(Connectable, log.Identified):
         try:
             trans = conn.begin()
         except:
-            conn.close()
-            raise
+            with util.safe_reraise():
+                conn.close()
         return Engine._trans_ctx(conn, trans, close_with_result)
 
     def transaction(self, callable_, *args, **kwargs):
diff --git a/lib/sqlalchemy/engine/default.py b/lib/sqlalchemy/engine/default.py
index abb9f0fc3..daa9fe085 100644
--- a/lib/sqlalchemy/engine/default.py
+++ b/lib/sqlalchemy/engine/default.py
@@ -737,7 +737,6 @@ class DefaultExecutionContext(interfaces.ExecutionContext):
             except Exception, e:
                 self.root_connection._handle_dbapi_exception(
                                 e, None, None, None, self)
-                raise
         else:
             inputsizes = {}
             for key in self.compiled.bind_names.values():
@@ -756,7 +755,6 @@ class DefaultExecutionContext(interfaces.ExecutionContext):
             except Exception, e:
                 self.root_connection._handle_dbapi_exception(
                                 e, None, None, None, self)
-                raise
 
     def _exec_default(self, default, type_):
         if default.is_sequence:
diff --git a/lib/sqlalchemy/engine/result.py b/lib/sqlalchemy/engine/result.py
index 1c148e1f0..88930081e 100644
--- a/lib/sqlalchemy/engine/result.py
+++ b/lib/sqlalchemy/engine/result.py
@@ -443,7 +443,6 @@ class ResultProxy(object):
         except Exception, e:
             self.connection._handle_dbapi_exception(
                               e, None, None, self.cursor, self.context)
-            raise
 
     @property
     def lastrowid(self):
@@ -467,7 +466,6 @@ class ResultProxy(object):
             self.connection._handle_dbapi_exception(
                                  e, None, None,
                                  self._saved_cursor, self.context)
-            raise
 
     @property
     def returns_rows(self):
@@ -752,7 +750,6 @@ class ResultProxy(object):
             self.connection._handle_dbapi_exception(
                                     e, None, None,
                                     self.cursor, self.context)
-            raise
 
     def fetchmany(self, size=None):
         """Fetch many rows, just like DB-API
@@ -772,7 +769,6 @@ class ResultProxy(object):
             self.connection._handle_dbapi_exception(
                                     e, None, None,
                                     self.cursor, self.context)
-            raise
 
     def fetchone(self):
         """Fetch one row, just like DB-API ``cursor.fetchone()``.
@@ -792,7 +788,6 @@ class ResultProxy(object):
             self.connection._handle_dbapi_exception(
                                     e, None, None,
                                     self.cursor, self.context)
-            raise
 
     def first(self):
         """Fetch the first row and then close the result set unconditionally.
@@ -809,7 +804,6 @@ class ResultProxy(object):
             self.connection._handle_dbapi_exception(
                                     e, None, None,
                                     self.cursor, self.context)
-            raise
 
         try:
             if row is not None:
diff --git a/lib/sqlalchemy/orm/session.py b/lib/sqlalchemy/orm/session.py
index 71e617e36..91e4f9736 100644
--- a/lib/sqlalchemy/orm/session.py
+++ b/lib/sqlalchemy/orm/session.py
@@ -341,8 +341,8 @@ class SessionTransaction(object):
                 for t in set(self._connections.values()):
                     t[1].prepare()
             except:
-                self.rollback()
-                raise
+                with util.safe_reraise():
+                    self.rollback()
 
         self._state = PREPARED
 
@@ -441,8 +441,8 @@ class SessionTransaction(object):
             try:
                 self.commit()
             except:
-                self.rollback()
-                raise
+                with util.safe_reraise():
+                    self.rollback()
         else:
             self.rollback()
 
@@ -1928,8 +1928,8 @@ class Session(_SessionClassMethods):
             transaction.commit()
 
         except:
-            transaction.rollback(_capture_exception=True)
-            raise
+            with util.safe_reraise():
+                transaction.rollback(_capture_exception=True)
 
     def is_modified(self, instance, include_collections=True,
                             passive=True):
diff --git a/lib/sqlalchemy/util/__init__.py b/lib/sqlalchemy/util/__init__.py
index 57bbdca85..3fa06c793 100644
--- a/lib/sqlalchemy/util/__init__.py
+++ b/lib/sqlalchemy/util/__init__.py
@@ -6,7 +6,8 @@
 
 from .compat import callable, cmp, reduce,  \
     threading, py3k, py3k_warning, jython, pypy, cpython, win32, set_types, \
-    pickle, dottedgetter, parse_qsl, namedtuple, next, WeakSet
+    pickle, dottedgetter, parse_qsl, namedtuple, next, WeakSet, reraise, \
+    raise_from_cause
 
 from ._collections import KeyedTuple, ImmutableContainer, immutabledict, \
     Properties, OrderedProperties, ImmutableProperties, OrderedDict, \
@@ -26,7 +27,7 @@ from .langhelpers import iterate_attributes, class_hierarchy, \
     duck_type_collection, assert_arg_type, symbol, dictlike_iteritems,\
     classproperty, set_creation_order, warn_exception, warn, NoneType,\
     constructor_copy, methods_equivalent, chop_traceback, asint,\
-    generic_repr, counter, PluginLoader, hybridmethod
+    generic_repr, counter, PluginLoader, hybridmethod, safe_reraise
 
 from .deprecations import warn_deprecated, warn_pending_deprecation, \
     deprecated, pending_deprecation
diff --git a/lib/sqlalchemy/util/compat.py b/lib/sqlalchemy/util/compat.py
index 2a0f06f8e..0f0a2f674 100644
--- a/lib/sqlalchemy/util/compat.py
+++ b/lib/sqlalchemy/util/compat.py
@@ -140,3 +140,29 @@ else:
     def b(s):
         return s
 
+
+if py3k:
+    def reraise(tp, value, tb=None, cause=None):
+        if cause is not None:
+            value.__cause__ = cause
+        if value.__traceback__ is not tb:
+            raise value.with_traceback(tb)
+        raise value
+
+    def raise_from_cause(exception, exc_info):
+        exc_type, exc_value, exc_tb = exc_info
+        reraise(type(exception), exception, tb=exc_tb, cause=exc_value)
+else:
+    exec("""def reraise(tp, value, tb=None, cause=None):
+        raise tp, value, tb
+    """)
+
+    def raise_from_cause(exception, exc_info):
+        # not as nice as that of Py3K, but at least preserves
+        # the code line where the issue occurred
+        exc_type, exc_value, exc_tb = exc_info
+        reraise(type(exception), exception, tb=exc_tb)
+
+
+
+
diff --git a/lib/sqlalchemy/util/langhelpers.py b/lib/sqlalchemy/util/langhelpers.py
index e3aed24d8..bba8ad734 100644
--- a/lib/sqlalchemy/util/langhelpers.py
+++ b/lib/sqlalchemy/util/langhelpers.py
@@ -20,6 +20,7 @@ from .compat import set_types, threading, \
 from functools import update_wrapper
 from .. import exc
 import hashlib
+from . import compat
 
 def md5_hex(x):
     # Py3K
@@ -28,6 +29,34 @@ def md5_hex(x):
     m.update(x)
     return m.hexdigest()
 
+class safe_reraise(object):
+    """Reraise an exception after invoking some
+    handler code.
+
+    Stores the existing exception info before
+    invoking so that it is maintained across a potential
+    coroutine context switch.
+
+    e.g.::
+
+        try:
+            sess.commit()
+        except:
+            with safe_reraise():
+                sess.rollback()
+
+    """
+
+    def __enter__(self):
+        self._exc_info = sys.exc_info()
+
+    def __exit__(self, type_, value, traceback):
+        if type_ is None:
+            exc_type, exc_value, exc_tb = self._exc_info
+            compat.reraise(exc_type, exc_value, exc_tb)
+        else:
+            compat.reraise(type_, value, traceback)
+
 def decode_slice(slc):
     """decode a slice object as sent to __getitem__.