From: Mike Bayer Date: Mon, 14 Sep 2020 12:04:09 +0000 (-0400) Subject: Adapt event exec_once_mutex to asyncio X-Git-Tag: rel_1_4_0b1~98 X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=4d17fe4063adef50c1d529993e0b047f503940e2;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git Adapt event exec_once_mutex to asyncio The pool makes use of a threading.Lock() for the "first_connect" event. if the pool is async make sure this is a greenlet-adapted asyncio lock. Fixes: #5581 Change-Id: If52415839c7ed82135465f1fe93b95d86c305820 --- diff --git a/lib/sqlalchemy/event/attr.py b/lib/sqlalchemy/event/attr.py index 87c6e980f8..abb264f98f 100644 --- a/lib/sqlalchemy/event/attr.py +++ b/lib/sqlalchemy/event/attr.py @@ -41,6 +41,7 @@ from . import registry from .. import exc from .. import util from ..util import threading +from ..util.concurrency import AsyncAdaptedLock class RefCollection(util.MemoizedSlots): @@ -277,6 +278,9 @@ class _EmptyListener(_InstanceLevelDispatch): class _CompoundListener(_InstanceLevelDispatch): __slots__ = "_exec_once_mutex", "_exec_once" + def _set_asyncio(self): + self._exec_once_mutex = AsyncAdaptedLock() + def _memoized_attr__exec_once_mutex(self): return threading.Lock() diff --git a/lib/sqlalchemy/event/base.py b/lib/sqlalchemy/event/base.py index a87c1fe448..c78080738f 100644 --- a/lib/sqlalchemy/event/base.py +++ b/lib/sqlalchemy/event/base.py @@ -241,8 +241,17 @@ class Events(util.with_metaclass(_EventMeta, object)): return target @classmethod - def _listen(cls, event_key, propagate=False, insert=False, named=False): - event_key.base_listen(propagate=propagate, insert=insert, named=named) + def _listen( + cls, + event_key, + propagate=False, + insert=False, + named=False, + asyncio=False, + ): + event_key.base_listen( + propagate=propagate, insert=insert, named=named, asyncio=asyncio + ) @classmethod def _remove(cls, event_key): diff --git a/lib/sqlalchemy/event/registry.py b/lib/sqlalchemy/event/registry.py index 19b9174b71..144dd45dc6 100644 --- a/lib/sqlalchemy/event/registry.py +++ b/lib/sqlalchemy/event/registry.py @@ -244,21 +244,26 @@ class _EventKey(object): return self._key in _key_to_collection def base_listen( - self, propagate=False, insert=False, named=False, retval=None + self, + propagate=False, + insert=False, + named=False, + retval=None, + asyncio=False, ): target, identifier = self.dispatch_target, self.identifier dispatch_collection = getattr(target.dispatch, identifier) + for_modify = dispatch_collection.for_modify(target.dispatch) + if asyncio: + for_modify._set_asyncio() + if insert: - dispatch_collection.for_modify(target.dispatch).insert( - self, propagate - ) + for_modify.insert(self, propagate) else: - dispatch_collection.for_modify(target.dispatch).append( - self, propagate - ) + for_modify.append(self, propagate) @property def _listen_fn(self): diff --git a/lib/sqlalchemy/pool/base.py b/lib/sqlalchemy/pool/base.py index f20b63cf54..87383fef71 100644 --- a/lib/sqlalchemy/pool/base.py +++ b/lib/sqlalchemy/pool/base.py @@ -59,6 +59,8 @@ class Pool(log.Identified): _dialect = _ConnDialect() + _is_asyncio = False + def __init__( self, creator, diff --git a/lib/sqlalchemy/pool/events.py b/lib/sqlalchemy/pool/events.py index 3954f907f4..9443877a91 100644 --- a/lib/sqlalchemy/pool/events.py +++ b/lib/sqlalchemy/pool/events.py @@ -54,6 +54,12 @@ class PoolEvents(event.Events): else: return target + @classmethod + def _listen(cls, event_key, **kw): + target = event_key.dispatch_target + + event_key.base_listen(asyncio=target._is_asyncio) + def connect(self, dbapi_connection, connection_record): """Called at the moment a particular DBAPI connection is first created for a given :class:`_pool.Pool`. diff --git a/lib/sqlalchemy/pool/impl.py b/lib/sqlalchemy/pool/impl.py index e1a9f00db1..ffdd63671a 100644 --- a/lib/sqlalchemy/pool/impl.py +++ b/lib/sqlalchemy/pool/impl.py @@ -218,6 +218,7 @@ class QueuePool(Pool): class AsyncAdaptedQueuePool(QueuePool): + _is_asyncio = True _queue_class = sqla_queue.AsyncAdaptedQueue diff --git a/lib/sqlalchemy/util/_concurrency_py3k.py b/lib/sqlalchemy/util/_concurrency_py3k.py index 3b112ff7db..82125b7713 100644 --- a/lib/sqlalchemy/util/_concurrency_py3k.py +++ b/lib/sqlalchemy/util/_concurrency_py3k.py @@ -96,6 +96,17 @@ try: del context.driver return result + class AsyncAdaptedLock: + def __init__(self): + self.mutex = asyncio.Lock() + + def __enter__(self): + await_fallback(self.mutex.acquire()) + return self + + def __exit__(self, *arg, **kw): + self.mutex.release() + except ImportError: # pragma: no cover greenlet = None diff --git a/lib/sqlalchemy/util/concurrency.py b/lib/sqlalchemy/util/concurrency.py index 4c4ea20d12..e0883aa683 100644 --- a/lib/sqlalchemy/util/concurrency.py +++ b/lib/sqlalchemy/util/concurrency.py @@ -7,6 +7,7 @@ if compat.py3k: from ._concurrency_py3k import await_fallback from ._concurrency_py3k import greenlet from ._concurrency_py3k import greenlet_spawn + from ._concurrency_py3k import AsyncAdaptedLock else: asyncio = None greenlet = None @@ -19,3 +20,6 @@ else: def greenlet_spawn(fn, *args, **kw): raise ValueError("Cannot use this function in py2.") + + def AsyncAdaptedLock(*args, **kw): + raise ValueError("Cannot use this function in py2.") diff --git a/test/ext/asyncio/test_engine_py3k.py b/test/ext/asyncio/test_engine_py3k.py index 59c47c4032..a5d167c2e7 100644 --- a/test/ext/asyncio/test_engine_py3k.py +++ b/test/ext/asyncio/test_engine_py3k.py @@ -50,6 +50,12 @@ class EngineFixture(fixtures.TablesTest): class AsyncEngineTest(EngineFixture): __backend__ = True + @async_test + async def test_init_once_concurrency(self, async_engine): + c1 = async_engine.connect() + c2 = async_engine.connect() + await asyncio.wait([c1, c2]) + @async_test async def test_connect_ctxmanager(self, async_engine): async with async_engine.connect() as conn: