summaryrefslogtreecommitdiff
path: root/lib/sqlalchemy
diff options
context:
space:
mode:
Diffstat (limited to 'lib/sqlalchemy')
-rw-r--r--lib/sqlalchemy/engine/base.py46
-rw-r--r--lib/sqlalchemy/engine/util.py79
-rw-r--r--lib/sqlalchemy/ext/asyncio/base.py4
-rw-r--r--lib/sqlalchemy/ext/asyncio/engine.py27
-rw-r--r--lib/sqlalchemy/ext/asyncio/session.py6
-rw-r--r--lib/sqlalchemy/future/engine.py8
-rw-r--r--lib/sqlalchemy/orm/session.py32
-rw-r--r--lib/sqlalchemy/testing/fixtures.py144
-rw-r--r--lib/sqlalchemy/util/concurrency.py5
9 files changed, 288 insertions, 63 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
diff --git a/lib/sqlalchemy/ext/asyncio/base.py b/lib/sqlalchemy/ext/asyncio/base.py
index fa8c5006e..d11b059fd 100644
--- a/lib/sqlalchemy/ext/asyncio/base.py
+++ b/lib/sqlalchemy/ext/asyncio/base.py
@@ -5,14 +5,14 @@ from . import exc as async_exc
class StartableContext(abc.ABC):
@abc.abstractmethod
- async def start(self) -> "StartableContext":
+ async def start(self, is_ctxmanager=False) -> "StartableContext":
pass
def __await__(self):
return self.start().__await__()
async def __aenter__(self):
- return await self.start()
+ return await self.start(is_ctxmanager=True)
@abc.abstractmethod
async def __aexit__(self, type_, value, traceback):
diff --git a/lib/sqlalchemy/ext/asyncio/engine.py b/lib/sqlalchemy/ext/asyncio/engine.py
index c637b3d90..17ddb614a 100644
--- a/lib/sqlalchemy/ext/asyncio/engine.py
+++ b/lib/sqlalchemy/ext/asyncio/engine.py
@@ -101,7 +101,7 @@ class AsyncConnection(ProxyComparable, StartableContext, AsyncConnectable):
self.sync_engine = async_engine.sync_engine
self.sync_connection = sync_connection
- async def start(self):
+ async def start(self, is_ctxmanager=False):
"""Start this :class:`_asyncio.AsyncConnection` object's context
outside of using a Python ``with:`` block.
@@ -518,19 +518,15 @@ class AsyncEngine(ProxyComparable, AsyncConnectable):
def __init__(self, conn):
self.conn = conn
- async def start(self):
- await self.conn.start()
+ async def start(self, is_ctxmanager=False):
+ await self.conn.start(is_ctxmanager=is_ctxmanager)
self.transaction = self.conn.begin()
await self.transaction.__aenter__()
return self.conn
async def __aexit__(self, type_, value, traceback):
- if type_ is not None:
- await self.transaction.rollback()
- else:
- if self.transaction.is_active:
- await self.transaction.commit()
+ await self.transaction.__aexit__(type_, value, traceback)
await self.conn.close()
def __init__(self, sync_engine: Engine):
@@ -678,7 +674,7 @@ class AsyncTransaction(ProxyComparable, StartableContext):
await greenlet_spawn(self._sync_transaction().commit)
- async def start(self):
+ async def start(self, is_ctxmanager=False):
"""Start this :class:`_asyncio.AsyncTransaction` object's context
outside of using a Python ``with:`` block.
@@ -689,17 +685,14 @@ class AsyncTransaction(ProxyComparable, StartableContext):
if self.nested
else self.connection._sync_connection().begin
)
+ if is_ctxmanager:
+ self.sync_transaction.__enter__()
return self
async def __aexit__(self, type_, value, traceback):
- if type_ is None and self.is_active:
- try:
- await self.commit()
- except:
- with util.safe_reraise():
- await self.rollback()
- else:
- await self.rollback()
+ await greenlet_spawn(
+ self._sync_transaction().__exit__, type_, value, traceback
+ )
def _get_sync_engine_or_connection(async_engine):
diff --git a/lib/sqlalchemy/ext/asyncio/session.py b/lib/sqlalchemy/ext/asyncio/session.py
index d8a5673eb..1b61d6ee3 100644
--- a/lib/sqlalchemy/ext/asyncio/session.py
+++ b/lib/sqlalchemy/ext/asyncio/session.py
@@ -399,15 +399,17 @@ class AsyncSessionTransaction(StartableContext):
await greenlet_spawn(self._sync_transaction().commit)
- async def start(self):
+ async def start(self, is_ctxmanager=False):
self.sync_transaction = await greenlet_spawn(
self.session.sync_session.begin_nested
if self.nested
else self.session.sync_session.begin
)
+ if is_ctxmanager:
+ self.sync_transaction.__enter__()
return self
async def __aexit__(self, type_, value, traceback):
- return await greenlet_spawn(
+ await greenlet_spawn(
self._sync_transaction().__exit__, type_, value, traceback
)
diff --git a/lib/sqlalchemy/future/engine.py b/lib/sqlalchemy/future/engine.py
index cee17a432..ab890ca4f 100644
--- a/lib/sqlalchemy/future/engine.py
+++ b/lib/sqlalchemy/future/engine.py
@@ -359,16 +359,12 @@ class Engine(_LegacyEngine):
def __enter__(self):
self.transaction = self.conn.begin()
+ 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:
self.conn.close()
diff --git a/lib/sqlalchemy/orm/session.py b/lib/sqlalchemy/orm/session.py
index a3ec360d0..cdf3a1585 100644
--- a/lib/sqlalchemy/orm/session.py
+++ b/lib/sqlalchemy/orm/session.py
@@ -31,6 +31,7 @@ from .. import engine
from .. import exc as sa_exc
from .. import sql
from .. import util
+from ..engine.util import TransactionalContext
from ..inspection import inspect
from ..sql import coercions
from ..sql import dml
@@ -475,7 +476,7 @@ class ORMExecuteState(util.MemoizedSlots):
]
-class SessionTransaction(object):
+class SessionTransaction(TransactionalContext):
"""A :class:`.Session`-level transaction.
:class:`.SessionTransaction` is produced from the
@@ -523,6 +524,8 @@ class SessionTransaction(object):
nested=False,
autobegin=False,
):
+ TransactionalContext._trans_ctx_check(session)
+
self.session = session
self._connections = {}
self._parent = parent
@@ -927,21 +930,14 @@ class SessionTransaction(object):
self.session = None
self._connections = None
- def __enter__(self):
- return self
+ def _get_subject(self):
+ return self.session
- def __exit__(self, type_, value, traceback):
- self._assert_active(deactive_ok=True, prepared_ok=True)
- if self.session._transaction is None:
- return
- if type_ is None:
- try:
- self.commit()
- except:
- with util.safe_reraise():
- self.rollback()
- else:
- self.rollback()
+ def _transaction_is_active(self):
+ return self._state is ACTIVE
+
+ def _transaction_is_closed(self):
+ return self._state is CLOSED
class Session(_SessionClassMethods):
@@ -1154,6 +1150,9 @@ class Session(_SessionClassMethods):
_sessions[self.hash_key] = self
+ # used by sqlalchemy.engine.util.TransactionalContext
+ _trans_context_manager = None
+
connection_callable = None
def __enter__(self):
@@ -1252,6 +1251,7 @@ class Session(_SessionClassMethods):
def _autobegin(self):
if not self.autocommit and self._transaction is None:
+
trans = SessionTransaction(self, autobegin=True)
assert self._transaction is trans
return True
@@ -1520,6 +1520,8 @@ class Session(_SessionClassMethods):
)
def _connection_for_bind(self, engine, execution_options=None, **kw):
+ TransactionalContext._trans_ctx_check(self)
+
if self._transaction is not None or self._autobegin():
return self._transaction._connection_for_bind(
engine, execution_options
diff --git a/lib/sqlalchemy/testing/fixtures.py b/lib/sqlalchemy/testing/fixtures.py
index c3eb1b363..581bc8bec 100644
--- a/lib/sqlalchemy/testing/fixtures.py
+++ b/lib/sqlalchemy/testing/fixtures.py
@@ -147,6 +147,150 @@ class TestBase(object):
else:
drop_all_tables_from_metadata(metadata, config.db)
+ @config.fixture(
+ params=[
+ (rollback, second_operation, begin_nested)
+ for rollback in (True, False)
+ for second_operation in ("none", "execute", "begin")
+ for begin_nested in (
+ True,
+ False,
+ )
+ ]
+ )
+ def trans_ctx_manager_fixture(self, request, metadata):
+ rollback, second_operation, begin_nested = request.param
+
+ from sqlalchemy import Table, Column, Integer, func, select
+ from . import eq_
+
+ t = Table("test", metadata, Column("data", Integer))
+ eng = getattr(self, "bind", None) or config.db
+
+ t.create(eng)
+
+ def run_test(subject, trans_on_subject, execute_on_subject):
+ with subject.begin() as trans:
+
+ if begin_nested:
+ if not config.requirements.savepoints.enabled:
+ config.skip_test("savepoints not enabled")
+ if execute_on_subject:
+ nested_trans = subject.begin_nested()
+ else:
+ nested_trans = trans.begin_nested()
+
+ with nested_trans:
+ if execute_on_subject:
+ subject.execute(t.insert(), {"data": 10})
+ else:
+ trans.execute(t.insert(), {"data": 10})
+
+ # for nested trans, we always commit/rollback on the
+ # "nested trans" object itself.
+ # only Session(future=False) will affect savepoint
+ # transaction for session.commit/rollback
+
+ if rollback:
+ nested_trans.rollback()
+ else:
+ nested_trans.commit()
+
+ if second_operation != "none":
+ with assertions.expect_raises_message(
+ sa.exc.InvalidRequestError,
+ "Can't operate on closed transaction "
+ "inside context "
+ "manager. Please complete the context "
+ "manager "
+ "before emitting further commands.",
+ ):
+ if second_operation == "execute":
+ if execute_on_subject:
+ subject.execute(
+ t.insert(), {"data": 12}
+ )
+ else:
+ trans.execute(t.insert(), {"data": 12})
+ elif second_operation == "begin":
+ if execute_on_subject:
+ subject.begin_nested()
+ else:
+ trans.begin_nested()
+
+ # outside the nested trans block, but still inside the
+ # transaction block, we can run SQL, and it will be
+ # committed
+ if execute_on_subject:
+ subject.execute(t.insert(), {"data": 14})
+ else:
+ trans.execute(t.insert(), {"data": 14})
+
+ else:
+ if execute_on_subject:
+ subject.execute(t.insert(), {"data": 10})
+ else:
+ trans.execute(t.insert(), {"data": 10})
+
+ if trans_on_subject:
+ if rollback:
+ subject.rollback()
+ else:
+ subject.commit()
+ else:
+ if rollback:
+ trans.rollback()
+ else:
+ trans.commit()
+
+ if second_operation != "none":
+ with assertions.expect_raises_message(
+ sa.exc.InvalidRequestError,
+ "Can't operate on closed transaction inside "
+ "context "
+ "manager. Please complete the context manager "
+ "before emitting further commands.",
+ ):
+ if second_operation == "execute":
+ if execute_on_subject:
+ subject.execute(t.insert(), {"data": 12})
+ else:
+ trans.execute(t.insert(), {"data": 12})
+ elif second_operation == "begin":
+ if hasattr(trans, "begin"):
+ trans.begin()
+ else:
+ subject.begin()
+ elif second_operation == "begin_nested":
+ if execute_on_subject:
+ subject.begin_nested()
+ else:
+ trans.begin_nested()
+
+ expected_committed = 0
+ if begin_nested:
+ # begin_nested variant, we inserted a row after the nested
+ # block
+ expected_committed += 1
+ if not rollback:
+ # not rollback variant, our row inserted in the target
+ # block itself would be committed
+ expected_committed += 1
+
+ if execute_on_subject:
+ eq_(
+ subject.scalar(select(func.count()).select_from(t)),
+ expected_committed,
+ )
+ else:
+ with subject.connect() as conn:
+ eq_(
+ conn.scalar(select(func.count()).select_from(t)),
+ expected_committed,
+ )
+
+ return run_test
+
_connection_fixture_connection = None
diff --git a/lib/sqlalchemy/util/concurrency.py b/lib/sqlalchemy/util/concurrency.py
index e26f305d9..60db9cfff 100644
--- a/lib/sqlalchemy/util/concurrency.py
+++ b/lib/sqlalchemy/util/concurrency.py
@@ -24,6 +24,11 @@ if not have_greenlet:
asyncio = None # noqa F811
def _not_implemented():
+ # this conditional is to prevent pylance from considering
+ # greenlet_spawn() etc as "no return" and dimming out code below it
+ if have_greenlet:
+ return None
+
if not compat.py3k:
raise ValueError("Cannot use this function in py2.")
else:
View Lorry Controller Status