500 lines
14 KiB
Python
500 lines
14 KiB
Python
|
# ext/asyncio/session.py
|
||
|
|
# Copyright (C) 2020-2021 the SQLAlchemy authors and contributors
|
||
|
|
# <see AUTHORS file>
|
||
|
|
#
|
||
|
|
# This module is part of SQLAlchemy and is released under
|
||
|
|
# the MIT License: http://www.opensource.org/licenses/mit-license.php
|
||
|
|
from . import engine
|
||
|
|
from . import result as _result
|
||
|
|
from .base import ReversibleProxy
|
||
|
|
from .base import StartableContext
|
||
|
|
from ... import util
|
||
|
|
from ...orm import object_session
|
||
|
|
from ...orm import Session
|
||
|
|
from ...orm import state as _instance_state
|
||
|
|
from ...util.concurrency import greenlet_spawn
|
||
|
|
|
||
|
|
|
||
|
|
@util.create_proxy_methods(
|
||
|
|
Session,
|
||
|
|
":class:`_orm.Session`",
|
||
|
|
":class:`_asyncio.AsyncSession`",
|
||
|
|
classmethods=["object_session", "identity_key"],
|
||
|
|
methods=[
|
||
|
|
"__contains__",
|
||
|
|
"__iter__",
|
||
|
|
"add",
|
||
|
|
"add_all",
|
||
|
|
"expire",
|
||
|
|
"expire_all",
|
||
|
|
"expunge",
|
||
|
|
"expunge_all",
|
||
|
|
"get_bind",
|
||
|
|
"is_modified",
|
||
|
|
"in_transaction",
|
||
|
|
"in_nested_transaction",
|
||
|
|
],
|
||
|
|
attributes=[
|
||
|
|
"dirty",
|
||
|
|
"deleted",
|
||
|
|
"new",
|
||
|
|
"identity_map",
|
||
|
|
"is_active",
|
||
|
|
"autoflush",
|
||
|
|
"no_autoflush",
|
||
|
|
"info",
|
||
|
|
],
|
||
|
|
)
|
||
|
|
class AsyncSession(ReversibleProxy):
|
||
|
|
"""Asyncio version of :class:`_orm.Session`.
|
||
|
|
|
||
|
|
|
||
|
|
.. versionadded:: 1.4
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
__slots__ = (
|
||
|
|
"binds",
|
||
|
|
"bind",
|
||
|
|
"sync_session",
|
||
|
|
"_proxied",
|
||
|
|
"_slots_dispatch",
|
||
|
|
)
|
||
|
|
|
||
|
|
dispatch = None
|
||
|
|
|
||
|
|
def __init__(self, bind=None, binds=None, **kw):
|
||
|
|
kw["future"] = True
|
||
|
|
if bind:
|
||
|
|
self.bind = bind
|
||
|
|
bind = engine._get_sync_engine_or_connection(bind)
|
||
|
|
|
||
|
|
if binds:
|
||
|
|
self.binds = binds
|
||
|
|
binds = {
|
||
|
|
key: engine._get_sync_engine_or_connection(b)
|
||
|
|
for key, b in binds.items()
|
||
|
|
}
|
||
|
|
|
||
|
|
self.sync_session = self._proxied = self._assign_proxied(
|
||
|
|
Session(bind=bind, binds=binds, **kw)
|
||
|
|
)
|
||
|
|
|
||
|
|
async def refresh(
|
||
|
|
self, instance, attribute_names=None, with_for_update=None
|
||
|
|
):
|
||
|
|
"""Expire and refresh the attributes on the given instance.
|
||
|
|
|
||
|
|
A query will be issued to the database and all attributes will be
|
||
|
|
refreshed with their current database value.
|
||
|
|
|
||
|
|
This is the async version of the :meth:`_orm.Session.refresh` method.
|
||
|
|
See that method for a complete description of all options.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
return await greenlet_spawn(
|
||
|
|
self.sync_session.refresh,
|
||
|
|
instance,
|
||
|
|
attribute_names=attribute_names,
|
||
|
|
with_for_update=with_for_update,
|
||
|
|
)
|
||
|
|
|
||
|
|
async def run_sync(self, fn, *arg, **kw):
|
||
|
|
"""Invoke the given sync callable passing sync self as the first
|
||
|
|
argument.
|
||
|
|
|
||
|
|
This method maintains the asyncio event loop all the way through
|
||
|
|
to the database connection by running the given callable in a
|
||
|
|
specially instrumented greenlet.
|
||
|
|
|
||
|
|
E.g.::
|
||
|
|
|
||
|
|
with AsyncSession(async_engine) as session:
|
||
|
|
await session.run_sync(some_business_method)
|
||
|
|
|
||
|
|
.. note::
|
||
|
|
|
||
|
|
The provided callable is invoked inline within the asyncio event
|
||
|
|
loop, and will block on traditional IO calls. IO within this
|
||
|
|
callable should only call into SQLAlchemy's asyncio database
|
||
|
|
APIs which will be properly adapted to the greenlet context.
|
||
|
|
|
||
|
|
.. seealso::
|
||
|
|
|
||
|
|
:ref:`session_run_sync`
|
||
|
|
"""
|
||
|
|
|
||
|
|
return await greenlet_spawn(fn, self.sync_session, *arg, **kw)
|
||
|
|
|
||
|
|
async def execute(
|
||
|
|
self,
|
||
|
|
statement,
|
||
|
|
params=None,
|
||
|
|
execution_options=util.EMPTY_DICT,
|
||
|
|
bind_arguments=None,
|
||
|
|
**kw
|
||
|
|
):
|
||
|
|
"""Execute a statement and return a buffered
|
||
|
|
:class:`_engine.Result` object."""
|
||
|
|
|
||
|
|
execution_options = execution_options.union({"prebuffer_rows": True})
|
||
|
|
|
||
|
|
return await greenlet_spawn(
|
||
|
|
self.sync_session.execute,
|
||
|
|
statement,
|
||
|
|
params=params,
|
||
|
|
execution_options=execution_options,
|
||
|
|
bind_arguments=bind_arguments,
|
||
|
|
**kw
|
||
|
|
)
|
||
|
|
|
||
|
|
async def scalar(
|
||
|
|
self,
|
||
|
|
statement,
|
||
|
|
params=None,
|
||
|
|
execution_options=util.EMPTY_DICT,
|
||
|
|
bind_arguments=None,
|
||
|
|
**kw
|
||
|
|
):
|
||
|
|
"""Execute a statement and return a scalar result."""
|
||
|
|
|
||
|
|
result = await self.execute(
|
||
|
|
statement,
|
||
|
|
params=params,
|
||
|
|
execution_options=execution_options,
|
||
|
|
bind_arguments=bind_arguments,
|
||
|
|
**kw
|
||
|
|
)
|
||
|
|
return result.scalar()
|
||
|
|
|
||
|
|
async def get(
|
||
|
|
self,
|
||
|
|
entity,
|
||
|
|
ident,
|
||
|
|
options=None,
|
||
|
|
populate_existing=False,
|
||
|
|
with_for_update=None,
|
||
|
|
identity_token=None,
|
||
|
|
):
|
||
|
|
"""Return an instance based on the given primary key identifier,
|
||
|
|
or ``None`` if not found.
|
||
|
|
|
||
|
|
|
||
|
|
"""
|
||
|
|
return await greenlet_spawn(
|
||
|
|
self.sync_session.get,
|
||
|
|
entity,
|
||
|
|
ident,
|
||
|
|
options=options,
|
||
|
|
populate_existing=populate_existing,
|
||
|
|
with_for_update=with_for_update,
|
||
|
|
identity_token=identity_token,
|
||
|
|
)
|
||
|
|
|
||
|
|
async def stream(
|
||
|
|
self,
|
||
|
|
statement,
|
||
|
|
params=None,
|
||
|
|
execution_options=util.EMPTY_DICT,
|
||
|
|
bind_arguments=None,
|
||
|
|
**kw
|
||
|
|
):
|
||
|
|
"""Execute a statement and return a streaming
|
||
|
|
:class:`_asyncio.AsyncResult` object."""
|
||
|
|
|
||
|
|
execution_options = execution_options.union({"stream_results": True})
|
||
|
|
|
||
|
|
result = await greenlet_spawn(
|
||
|
|
self.sync_session.execute,
|
||
|
|
statement,
|
||
|
|
params=params,
|
||
|
|
execution_options=execution_options,
|
||
|
|
bind_arguments=bind_arguments,
|
||
|
|
**kw
|
||
|
|
)
|
||
|
|
return _result.AsyncResult(result)
|
||
|
|
|
||
|
|
async def delete(self, instance):
|
||
|
|
"""Mark an instance as deleted.
|
||
|
|
|
||
|
|
The database delete operation occurs upon ``flush()``.
|
||
|
|
|
||
|
|
As this operation may need to cascade along unloaded relationships,
|
||
|
|
it is awaitable to allow for those queries to take place.
|
||
|
|
|
||
|
|
|
||
|
|
"""
|
||
|
|
return await greenlet_spawn(self.sync_session.delete, instance)
|
||
|
|
|
||
|
|
async def merge(self, instance, load=True):
|
||
|
|
"""Copy the state of a given instance into a corresponding instance
|
||
|
|
within this :class:`_asyncio.AsyncSession`.
|
||
|
|
|
||
|
|
"""
|
||
|
|
return await greenlet_spawn(
|
||
|
|
self.sync_session.merge, instance, load=load
|
||
|
|
)
|
||
|
|
|
||
|
|
async def flush(self, objects=None):
|
||
|
|
"""Flush all the object changes to the database.
|
||
|
|
|
||
|
|
.. seealso::
|
||
|
|
|
||
|
|
:meth:`_orm.Session.flush`
|
||
|
|
|
||
|
|
"""
|
||
|
|
await greenlet_spawn(self.sync_session.flush, objects=objects)
|
||
|
|
|
||
|
|
def get_transaction(self):
|
||
|
|
"""Return the current root transaction in progress, if any.
|
||
|
|
|
||
|
|
:return: an :class:`_asyncio.AsyncSessionTransaction` object, or
|
||
|
|
``None``.
|
||
|
|
|
||
|
|
.. versionadded:: 1.4.18
|
||
|
|
|
||
|
|
"""
|
||
|
|
trans = self.sync_session.get_transaction()
|
||
|
|
if trans is not None:
|
||
|
|
return AsyncSessionTransaction._retrieve_proxy_for_target(trans)
|
||
|
|
else:
|
||
|
|
return None
|
||
|
|
|
||
|
|
def get_nested_transaction(self):
|
||
|
|
"""Return the current nested transaction in progress, if any.
|
||
|
|
|
||
|
|
:return: an :class:`_asyncio.AsyncSessionTransaction` object, or
|
||
|
|
``None``.
|
||
|
|
|
||
|
|
.. versionadded:: 1.4.18
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
trans = self.sync_session.get_nested_transaction()
|
||
|
|
if trans is not None:
|
||
|
|
return AsyncSessionTransaction._retrieve_proxy_for_target(trans)
|
||
|
|
else:
|
||
|
|
return None
|
||
|
|
|
||
|
|
async def connection(self):
|
||
|
|
r"""Return a :class:`_asyncio.AsyncConnection` object corresponding to
|
||
|
|
this :class:`.Session` object's transactional state.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
sync_connection = await greenlet_spawn(self.sync_session.connection)
|
||
|
|
return engine.AsyncConnection._retrieve_proxy_for_target(
|
||
|
|
sync_connection
|
||
|
|
)
|
||
|
|
|
||
|
|
def begin(self, **kw):
|
||
|
|
"""Return an :class:`_asyncio.AsyncSessionTransaction` object.
|
||
|
|
|
||
|
|
The underlying :class:`_orm.Session` will perform the
|
||
|
|
"begin" action when the :class:`_asyncio.AsyncSessionTransaction`
|
||
|
|
object is entered::
|
||
|
|
|
||
|
|
async with async_session.begin():
|
||
|
|
# .. ORM transaction is begun
|
||
|
|
|
||
|
|
Note that database IO will not normally occur when the session-level
|
||
|
|
transaction is begun, as database transactions begin on an
|
||
|
|
on-demand basis. However, the begin block is async to accommodate
|
||
|
|
for a :meth:`_orm.SessionEvents.after_transaction_create`
|
||
|
|
event hook that may perform IO.
|
||
|
|
|
||
|
|
For a general description of ORM begin, see
|
||
|
|
:meth:`_orm.Session.begin`.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
return AsyncSessionTransaction(self)
|
||
|
|
|
||
|
|
def begin_nested(self, **kw):
|
||
|
|
"""Return an :class:`_asyncio.AsyncSessionTransaction` object
|
||
|
|
which will begin a "nested" transaction, e.g. SAVEPOINT.
|
||
|
|
|
||
|
|
Behavior is the same as that of :meth:`_asyncio.AsyncSession.begin`.
|
||
|
|
|
||
|
|
For a general description of ORM begin nested, see
|
||
|
|
:meth:`_orm.Session.begin_nested`.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
return AsyncSessionTransaction(self, nested=True)
|
||
|
|
|
||
|
|
async def rollback(self):
|
||
|
|
"""Rollback the current transaction in progress."""
|
||
|
|
return await greenlet_spawn(self.sync_session.rollback)
|
||
|
|
|
||
|
|
async def commit(self):
|
||
|
|
"""Commit the current transaction in progress."""
|
||
|
|
return await greenlet_spawn(self.sync_session.commit)
|
||
|
|
|
||
|
|
async def close(self):
|
||
|
|
"""Close out the transactional resources and ORM objects used by this
|
||
|
|
:class:`_asyncio.AsyncSession`.
|
||
|
|
|
||
|
|
This expunges all ORM objects associated with this
|
||
|
|
:class:`_asyncio.AsyncSession`, ends any transaction in progress and
|
||
|
|
:term:`releases` any :class:`_asyncio.AsyncConnection` objects which
|
||
|
|
this :class:`_asyncio.AsyncSession` itself has checked out from
|
||
|
|
associated :class:`_asyncio.AsyncEngine` objects. The operation then
|
||
|
|
leaves the :class:`_asyncio.AsyncSession` in a state which it may be
|
||
|
|
used again.
|
||
|
|
|
||
|
|
.. tip::
|
||
|
|
|
||
|
|
The :meth:`_asyncio.AsyncSession.close` method **does not prevent
|
||
|
|
the Session from being used again**. The
|
||
|
|
:class:`_asyncio.AsyncSession` itself does not actually have a
|
||
|
|
distinct "closed" state; it merely means the
|
||
|
|
:class:`_asyncio.AsyncSession` will release all database
|
||
|
|
connections and ORM objects.
|
||
|
|
|
||
|
|
|
||
|
|
.. seealso::
|
||
|
|
|
||
|
|
:ref:`session_closing` - detail on the semantics of
|
||
|
|
:meth:`_asyncio.AsyncSession.close`
|
||
|
|
|
||
|
|
"""
|
||
|
|
return await greenlet_spawn(self.sync_session.close)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
async def close_all(self):
|
||
|
|
"""Close all :class:`_asyncio.AsyncSession` sessions."""
|
||
|
|
return await greenlet_spawn(self.sync_session.close_all)
|
||
|
|
|
||
|
|
async def __aenter__(self):
|
||
|
|
return self
|
||
|
|
|
||
|
|
async def __aexit__(self, type_, value, traceback):
|
||
|
|
await self.close()
|
||
|
|
|
||
|
|
def _maker_context_manager(self):
|
||
|
|
# no @contextlib.asynccontextmanager until python3.7, gr
|
||
|
|
return _AsyncSessionContextManager(self)
|
||
|
|
|
||
|
|
|
||
|
|
class _AsyncSessionContextManager:
|
||
|
|
def __init__(self, async_session):
|
||
|
|
self.async_session = async_session
|
||
|
|
|
||
|
|
async def __aenter__(self):
|
||
|
|
self.trans = self.async_session.begin()
|
||
|
|
await self.trans.__aenter__()
|
||
|
|
return self.async_session
|
||
|
|
|
||
|
|
async def __aexit__(self, type_, value, traceback):
|
||
|
|
await self.trans.__aexit__(type_, value, traceback)
|
||
|
|
await self.async_session.__aexit__(type_, value, traceback)
|
||
|
|
|
||
|
|
|
||
|
|
class AsyncSessionTransaction(ReversibleProxy, StartableContext):
|
||
|
|
"""A wrapper for the ORM :class:`_orm.SessionTransaction` object.
|
||
|
|
|
||
|
|
This object is provided so that a transaction-holding object
|
||
|
|
for the :meth:`_asyncio.AsyncSession.begin` may be returned.
|
||
|
|
|
||
|
|
The object supports both explicit calls to
|
||
|
|
:meth:`_asyncio.AsyncSessionTransaction.commit` and
|
||
|
|
:meth:`_asyncio.AsyncSessionTransaction.rollback`, as well as use as an
|
||
|
|
async context manager.
|
||
|
|
|
||
|
|
|
||
|
|
.. versionadded:: 1.4
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
__slots__ = ("session", "sync_transaction", "nested")
|
||
|
|
|
||
|
|
def __init__(self, session, nested=False):
|
||
|
|
self.session = session
|
||
|
|
self.nested = nested
|
||
|
|
self.sync_transaction = None
|
||
|
|
|
||
|
|
@property
|
||
|
|
def is_active(self):
|
||
|
|
return (
|
||
|
|
self._sync_transaction() is not None
|
||
|
|
and self._sync_transaction().is_active
|
||
|
|
)
|
||
|
|
|
||
|
|
def _sync_transaction(self):
|
||
|
|
if not self.sync_transaction:
|
||
|
|
self._raise_for_not_started()
|
||
|
|
return self.sync_transaction
|
||
|
|
|
||
|
|
async def rollback(self):
|
||
|
|
"""Roll back this :class:`_asyncio.AsyncTransaction`."""
|
||
|
|
await greenlet_spawn(self._sync_transaction().rollback)
|
||
|
|
|
||
|
|
async def commit(self):
|
||
|
|
"""Commit this :class:`_asyncio.AsyncTransaction`."""
|
||
|
|
|
||
|
|
await greenlet_spawn(self._sync_transaction().commit)
|
||
|
|
|
||
|
|
async def start(self, is_ctxmanager=False):
|
||
|
|
self.sync_transaction = self._assign_proxied(
|
||
|
|
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):
|
||
|
|
await greenlet_spawn(
|
||
|
|
self._sync_transaction().__exit__, type_, value, traceback
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def async_object_session(instance):
|
||
|
|
"""Return the :class:`_asyncio.AsyncSession` to which the given instance
|
||
|
|
belongs.
|
||
|
|
|
||
|
|
This function makes use of the sync-API function
|
||
|
|
:class:`_orm.object_session` to retrieve the :class:`_orm.Session` which
|
||
|
|
refers to the given instance, and from there links it to the original
|
||
|
|
:class:`_asyncio.AsyncSession`.
|
||
|
|
|
||
|
|
If the :class:`_asyncio.AsyncSession` has been garbage collected, the
|
||
|
|
return value is ``None``.
|
||
|
|
|
||
|
|
This functionality is also available from the
|
||
|
|
:attr:`_orm.InstanceState.async_session` accessor.
|
||
|
|
|
||
|
|
:param instance: an ORM mapped instance
|
||
|
|
:return: an :class:`_asyncio.AsyncSession` object, or ``None``.
|
||
|
|
|
||
|
|
.. versionadded:: 1.4.18
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
session = object_session(instance)
|
||
|
|
if session is not None:
|
||
|
|
return async_session(session)
|
||
|
|
else:
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def async_session(session):
|
||
|
|
"""Return the :class:`_asyncio.AsyncSession` which is proxying the given
|
||
|
|
:class:`_orm.Session` object, if any.
|
||
|
|
|
||
|
|
:param session: a :class:`_orm.Session` instance.
|
||
|
|
:return: a :class:`_asyncio.AsyncSession` instance, or ``None``.
|
||
|
|
|
||
|
|
.. versionadded:: 1.4.18
|
||
|
|
|
||
|
|
"""
|
||
|
|
return AsyncSession._retrieve_proxy_for_target(session, regenerate=False)
|
||
|
|
|
||
|
|
|
||
|
|
_instance_state._async_provider = async_session
|