diff options
Diffstat (limited to 'lib/sqlalchemy/engine')
| -rw-r--r-- | lib/sqlalchemy/engine/base.py | 46 | ||||
| -rw-r--r-- | lib/sqlalchemy/engine/util.py | 79 |
2 files changed, 104 insertions, 21 deletions
diff --git a/lib/sqlalchemy/engine/base.py b/lib/sqlalchemy/engine/base.py index 293dc21b4..663482b1f 100644 --- a/lib/sqlalchemy/engine/base.py +++ b/lib/sqlalchemy/engine/base.py @@ -13,6 +13,7 @@ from .interfaces import Connectable from .interfaces import ExceptionContext from .util import _distill_params from .util import _distill_params_20 +from .util import TransactionalContext from .. import exc from .. import inspection from .. import log @@ -60,6 +61,9 @@ class Connection(Connectable): _is_future = False _sqla_logger_namespace = "sqlalchemy.engine.Connection" + # used by sqlalchemy.engine.util.TransactionalContext + _trans_context_manager = None + def __init__( self, engine, @@ -1683,6 +1687,9 @@ class Connection(Connectable): ): self._invalid_transaction() + elif self._trans_context_manager: + TransactionalContext._trans_ctx_check(self) + if self._is_future and self._transaction is None: self._autobegin() @@ -2182,7 +2189,7 @@ class ExceptionContextImpl(ExceptionContext): self.invalidate_pool_on_disconnect = invalidate_pool_on_disconnect -class Transaction(object): +class Transaction(TransactionalContext): """Represent a database transaction in progress. The :class:`.Transaction` object is procured by @@ -2324,21 +2331,14 @@ class Transaction(object): finally: assert not self.is_active - def __enter__(self): - return self + def _get_subject(self): + return self.connection - def __exit__(self, type_, value, traceback): - if type_ is None and self.is_active: - try: - self.commit() - except: - with util.safe_reraise(): - self.rollback() - else: - if self._deactivated_from_connection: - self.close() - else: - self.rollback() + def _transaction_is_active(self): + return self.is_active + + def _transaction_is_closed(self): + return not self._deactivated_from_connection class MarkerTransaction(Transaction): @@ -2368,6 +2368,10 @@ class MarkerTransaction(Transaction): ) self.connection = connection + + if connection._trans_context_manager: + TransactionalContext._trans_ctx_check(connection) + if connection._nested_transaction is not None: self._transaction = connection._nested_transaction else: @@ -2429,6 +2433,8 @@ class RootTransaction(Transaction): def __init__(self, connection): assert connection._transaction is None + if connection._trans_context_manager: + TransactionalContext._trans_ctx_check(connection) self.connection = connection self._connection_begin_impl() connection._transaction = self @@ -2564,6 +2570,8 @@ class NestedTransaction(Transaction): def __init__(self, connection): assert connection._transaction is not None + if connection._trans_context_manager: + TransactionalContext._trans_ctx_check(connection) self.connection = connection self._savepoint = self.connection._savepoint_impl() self.is_active = True @@ -2935,16 +2943,12 @@ class Engine(Connectable, log.Identified): self.close_with_result = close_with_result def __enter__(self): + self.transaction.__enter__() return self.conn def __exit__(self, type_, value, traceback): try: - if type_ is not None: - if self.transaction.is_active: - self.transaction.rollback() - else: - if self.transaction.is_active: - self.transaction.commit() + self.transaction.__exit__(type_, value, traceback) finally: if not self.close_with_result: self.conn.close() diff --git a/lib/sqlalchemy/engine/util.py b/lib/sqlalchemy/engine/util.py index ede263198..17e3510aa 100644 --- a/lib/sqlalchemy/engine/util.py +++ b/lib/sqlalchemy/engine/util.py @@ -153,3 +153,82 @@ def _distill_params_20(params): return (params,), _no_kw else: raise exc.ArgumentError("mapping or sequence expected for parameters") + + +class TransactionalContext(object): + """Apply Python context manager behavior to transaction objects. + + Performs validation to ensure the subject of the transaction is not + used if the transaction were ended prematurely. + + """ + + _trans_subject = None + + def _transaction_is_active(self): + raise NotImplementedError() + + def _transaction_is_closed(self): + raise NotImplementedError() + + def _get_subject(self): + raise NotImplementedError() + + @classmethod + def _trans_ctx_check(cls, subject): + trans_context = subject._trans_context_manager + if trans_context: + if not trans_context._transaction_is_active(): + raise exc.InvalidRequestError( + "Can't operate on closed transaction inside context " + "manager. Please complete the context manager " + "before emitting further commands." + ) + + def __enter__(self): + subject = self._get_subject() + + # none for outer transaction, may be non-None for nested + # savepoint, legacy nesting cases + trans_context = subject._trans_context_manager + self._outer_trans_ctx = trans_context + + self._trans_subject = subject + subject._trans_context_manager = self + return self + + def __exit__(self, type_, value, traceback): + subject = self._trans_subject + + # simplistically we could assume that + # "subject._trans_context_manager is self". However, any calling + # code that is manipulating __exit__ directly would break this + # assumption. alembic context manager + # is an example of partial use that just calls __exit__ and + # not __enter__ at the moment. it's safe to assume this is being done + # in the wild also + out_of_band_exit = ( + subject is None or subject._trans_context_manager is not self + ) + + if type_ is None and self._transaction_is_active(): + try: + self.commit() + except: + with util.safe_reraise(): + self.rollback() + finally: + if not out_of_band_exit: + subject._trans_context_manager = self._outer_trans_ctx + self._trans_subject = self._outer_trans_ctx = None + else: + try: + if not self._transaction_is_active(): + if not self._transaction_is_closed(): + self.close() + else: + self.rollback() + finally: + if not out_of_band_exit: + subject._trans_context_manager = self._outer_trans_ctx + self._trans_subject = self._outer_trans_ctx = None |