def test_same_name(): v1 = ContextVar('v') v2 = ContextVar('v') v1.set(1) v2.set(2) assert v1.get() == 1 assert v2.get() == 2
def __init__(self, dialect, pool, loop, logging_name=None, echo=None, execution_options=None): self._sa_engine = _SAEngine( dialect, logging_name=logging_name, echo=echo, execution_options=execution_options) self._dialect = dialect self._pool = pool self._loop = loop self._ctx = ContextVar('gino', default=None)
def __init__(self, connection_name: str, fetch_inserted: bool = True, **kwargs) -> None: self.log = logging.getLogger("db_client") self.connection_name = connection_name self.fetch_inserted = fetch_inserted self._current_transaction = ContextVar(self.connection_name, default=self) # Type: dict
async def test_max_count_cancel(loop): t_exec = 0 delay_exec = 0.1 event = Event() executions = 0 leeway = 100 max_count = 5 arg = ContextVar("arg") tasks = [] executing_task = None @aggregate(leeway * 1000, max_count) async def pow(*args: float, power: float = 2) -> List[float]: nonlocal executions, executing_task, t_exec, delay_exec await asyncio.sleep(0.03) t_exec = time.time() executions += 1 executing_task = tasks[arg.get()] event.set() await asyncio.sleep(delay_exec) return [math.pow(num, power) for num in args] async def pho(num: int): arg.set(num) return await pow(float(num)) t = time.time() tasks = [] for i in range(5): tasks.append(loop.create_task(pho(i))) # Execution must have started await event.wait() event.clear() elapsed = t_exec - t assert 0 < elapsed < leeway assert all(not task.done() for task in tasks) assert executions == 1 first_executing_task = executing_task first_executing_task.cancel() # Another task must have tried to execute await event.wait() assert executions == 2 assert first_executing_task.cancelled() assert all( not task.done() for task in tasks if task is not first_executing_task ) # Must have finished await wait(tasks) assert first_executing_task.cancelled() for i, task in enumerate(tasks): if task is first_executing_task: continue assert task.done() assert task.result() == math.pow(i, 2)
class BaseDBAsyncClient: query_class = Query executor_class = BaseExecutor schema_generator = BaseSchemaGenerator capabilities = Capabilities("") def __init__(self, connection_name: str, fetch_inserted: bool = True, **kwargs) -> None: self.log = logging.getLogger("db_client") self.connection_name = connection_name self.fetch_inserted = fetch_inserted self._current_transaction = ContextVar(self.connection_name, default=self) # Type: dict async def create_connection(self, with_db: bool) -> None: raise NotImplementedError() # pragma: nocoverage async def close(self) -> None: raise NotImplementedError() # pragma: nocoverage async def db_create(self) -> None: raise NotImplementedError() # pragma: nocoverage async def db_delete(self) -> None: raise NotImplementedError() # pragma: nocoverage def acquire_connection(self): raise NotImplementedError() # pragma: nocoverage def _in_transaction(self) -> "BaseTransactionWrapper": raise NotImplementedError() # pragma: nocoverage async def execute_insert(self, query: str, values: list) -> Any: raise NotImplementedError() # pragma: nocoverage async def execute_query(self, query: str) -> Sequence[dict]: raise NotImplementedError() # pragma: nocoverage async def execute_script(self, query: str) -> None: raise NotImplementedError() # pragma: nocoverage def get_current_transaction(self) -> "BaseDBAsyncClient": return self._current_transaction.get()
class GinoEngine: """ Connects a :class:`~.dialects.base.Pool` and :class:`~sqlalchemy.engine.interfaces.Dialect` together to provide a source of database connectivity and behavior. A :class:`.GinoEngine` object is instantiated publicly using the :func:`gino.create_engine` function or :func:`db.set_bind() <gino.api.Gino.set_bind>` method. .. seealso:: :doc:`/engine` """ connection_cls = GinoConnection """Customizes the connection class to use, default is :class:`.GinoConnection`.""" def __init__(self, dialect, pool, loop, logging_name=None, echo=None, execution_options=None): self._sa_engine = _SAEngine( dialect, logging_name=logging_name, echo=echo, execution_options=execution_options) self._dialect = dialect self._pool = pool self._loop = loop self._ctx = ContextVar('gino', default=None) @property def dialect(self): """ Read-only property for the :class:`~sqlalchemy.engine.interfaces.Dialect` of this engine. """ return self._dialect @property def raw_pool(self): """ Read-only access to the underlying database connection pool instance. This depends on the actual dialect in use, :class:`~asyncpg.pool.Pool` of asyncpg for example. """ return self._pool.raw_pool def acquire(self, *, timeout=None, reuse=False, lazy=False, reusable=True): """ Acquire a connection from the pool. There are two ways using this method - as an asynchronous context manager:: async with engine.acquire() as conn: # play with the connection which will guarantee the connection is returned to the pool when leaving the ``async with`` block; or as a coroutine:: conn = await engine.acquire() try: # play with the connection finally: await conn.release() where the connection should be manually returned to the pool with :meth:`conn.release() <.GinoConnection.release>`. Within the same context (usually the same :class:`~asyncio.Task`, see also :doc:`/transaction`), a nesting acquire by default re :param timeout: Block up to ``timeout`` seconds until there is one free connection in the pool. Default is ``None`` - block forever until succeeded. This has no effect when ``lazy=True``, and depends on the actual situation when ``reuse=True``. :param reuse: Reuse the latest reusable acquired connection (before it's returned to the pool) in current context if there is one, or borrow a new one if none present. Default is ``False`` for always borrow a new one. This is useful when you are in a nested method call series, wishing to use the same connection without passing it around as parameters. See also: :doc:`/transaction`. A reusing connection is not reusable even if ``reusable=True``. If the reused connection happened to be a lazy one, then the reusing connection is lazy too. :param lazy: Don't acquire the actual underlying connection yet - do it only when needed. Default is ``False`` for always do it immediately. This is useful before entering a code block which may or may not make use of a given connection object. Feeding in a lazy connection will save the borrow-return job if the connection is never used. If setting ``reuse=True`` at the same time, then the reused connection - if any - applies the same laziness. For example, reusing a lazy connection with ``lazy=False`` will cause the reused connection to acquire an underlying connection immediately. :param reusable: Mark this connection as reusable or otherwise. This has no effect if it is a reusing connection. All reusable connections are placed in a stack, any reusing acquire operation will always reuse the top (latest) reusable connection. One reusable connection may be reused by several reusing connections - they all share one same underlying connection. Acquiring a connection with ``reusable=False`` and ``reusing=False`` makes it a cleanly isolated connection which is only referenced once here. :return: A :class:`.GinoConnection` object. """ return _AcquireContext(functools.partial( self._acquire, timeout, reuse, lazy, reusable)) async def _acquire(self, timeout, reuse, lazy, reusable): stack = _ContextualStack(self._ctx) if reuse and stack: dbapi_conn = _ReusingDBAPIConnection(self._dialect.cursor_cls, stack.top) reusable = False else: dbapi_conn = _DBAPIConnection(self._dialect.cursor_cls, self._pool) rv = self.connection_cls(self._dialect, _SAConnection(self._sa_engine, dbapi_conn), stack if reusable else None) dbapi_conn.gino_conn = rv if not lazy: await dbapi_conn.acquire(timeout=timeout) if reusable: stack.push(dbapi_conn) return rv @property def current_connection(self): """ Gets the most recently acquired reusable connection in the context. ``None`` if there is no such connection. :return: :class:`.GinoConnection` """ stack = self._ctx.get() if stack: return stack[-1].gino_conn async def close(self): """ Close the engine, by closing the underlying pool. """ await self._pool.close() async def all(self, clause, *multiparams, **params): """ Acquires a connection with ``reuse=True`` and runs :meth:`~.GinoConnection.all` on it. ``reuse=True`` means you can safely do this without borrowing more than one underlying connection:: async with engine.acquire(): await engine.all('SELECT ...') The same applies for other query methods. """ async with self.acquire(reuse=True) as conn: return await conn.all(clause, *multiparams, **params) async def first(self, clause, *multiparams, **params): """ Runs :meth:`~.GinoConnection.first`, See :meth:`.all`. """ async with self.acquire(reuse=True) as conn: return await conn.first(clause, *multiparams, **params) async def scalar(self, clause, *multiparams, **params): """ Runs :meth:`~.GinoConnection.scalar`, See :meth:`.all`. """ async with self.acquire(reuse=True) as conn: return await conn.scalar(clause, *multiparams, **params) async def status(self, clause, *multiparams, **params): """ Runs :meth:`~.GinoConnection.status`. See also :meth:`.all`. """ async with self.acquire(reuse=True) as conn: return await conn.status(clause, *multiparams, **params) def compile(self, clause, *multiparams, **params): """ A shortcut for :meth:`~gino.dialects.base.AsyncDialectMixin.compile` on the dialect, returns raw SQL string and parameters according to the rules of the dialect. """ return self._dialect.compile(clause, *multiparams, **params) def transaction(self, *args, timeout=None, reuse=True, reusable=True, **kwargs): """ Borrows a new connection and starts a transaction with it. Different to :meth:`.GinoConnection.transaction`, transaction on engine level supports only managed usage:: async with engine.transaction() as tx: # play with transaction here Where the implicitly acquired connection is available as :attr:`tx.connection <gino.transaction.GinoTransaction.connection>`. By default, :meth:`.transaction` acquires connection with ``reuse=True`` and ``reusable=True``, that means it by default tries to create a nested transaction instead of a new transaction on a new connection. You can change the default behavior by setting these two arguments. The other arguments are the same as :meth:`~.GinoConnection.transaction` on connection. .. seealso:: :meth:`.GinoEngine.acquire` :meth:`.GinoConnection.transaction` :class:`~gino.transaction.GinoTransaction` :return: A asynchronous context manager that yields a :class:`~gino.transaction.GinoTransaction` """ return _TransactionContext(self.acquire( timeout=timeout, reuse=reuse, reusable=reusable), (args, kwargs)) def iterate(self, clause, *multiparams, **params): """ Creates a server-side cursor in database for large query results. This requires that there is a reusable connection in the current context, and an active transaction is present. Then its :meth:`.GinoConnection.iterate` is executed and returned. """ connection = self.current_connection if connection is None: raise ValueError( 'No Connection in context, please provide one') return connection.iterate(clause, *multiparams, **params) def update_execution_options(self, **opt): """Update the default execution_options dictionary of this :class:`.GinoEngine`. .. seealso:: :meth:`sqlalchemy.engine.Engine.update_execution_options` :meth:`.GinoConnection.execution_options` """ self._sa_engine.update_execution_options(**opt) async def _run_visitor(self, *args, **kwargs): async with self.acquire(reuse=True) as conn: await getattr(conn, '_run_visitor')(*args, **kwargs)
import asyncio import pytest from aiocontextvars import ContextVar v = ContextVar('v') @pytest.mark.asyncio async def parallel(x): v.set(x) await asyncio.sleep(0.1) assert v.get() == x @pytest.mark.asyncio async def test_parallel(): await asyncio.gather(*map(parallel, range(16))) # noinspection PyUnusedLocal @pytest.mark.asyncio async def test_parallel_with_inherit(): await asyncio.gather(*map(parallel, range(16))) # noinspection PyUnusedLocal @pytest.mark.asyncio async def test_inherit(): v.set('initial') async def sub():
def test_get(): v = ContextVar('v') with pytest.raises(LookupError): v.get() assert v.get(456) == 456 return v
def test_get_with_default(): v = ContextVar('v', default=123) assert v.get() == 123 assert v.get(456) == 456 return v
import asyncio import textwrap import gettext from aiocontextvars import ContextVar import discord from discord.abc import Messageable from discord.ext import commands from meowth.core import checks from meowth.utils import convert_to_bool, make_embed, bold cvar = ContextVar('bot') def ctx_setup(loop): import builtins builtins.__dict__['_'] = use_current_gettext builtins.__dict__['get_ctx'] = cvar.get builtins.__dict__['__cvar__'] = cvar def use_current_gettext(*args, **kwargs): return cvar.get().get_text(*args, **kwargs) class Context(commands.Context): def __init__(self, **kwargs): super().__init__(**kwargs) self.dbi = self.bot.dbi
@property @abstractmethod def user(self) -> User: """Get the current request user""" @property @abstractmethod def roles(self) -> List[str]: """Get current user roles""" @abstractmethod def validate_roles(self, required_roles: List[str] = None): """Check if a user is authenticated""" user_var = ContextVar('user', default=None) class StandardAuthProvider(AuthProvider): def setup(self, user: User) -> None: user_var.set(user) @property def user(self) -> User: if not user_var.get(): raise AuthenticationError("Not authenticated.") return user_var.get() @property def roles(self) -> List[str]:
from abc import ABC, abstractmethod from aiocontextvars import ContextVar from .tenant import Tenant class TenantProvider(ABC): """Tenant service.""" @abstractmethod def setup(self, tenant: Tenant) -> None: "Setup current tenant method to be implemented." @property @abstractmethod def tenant(self) -> Tenant: """Get the current tenant""" tenant_var = ContextVar('tenant', default=None) class StandardTenantProvider(TenantProvider): def setup(self, tenant: Tenant) -> None: tenant_var.set(tenant) @property def tenant(self) -> Tenant: if not tenant_var.get(): raise ValueError('No tenant has been set.') return tenant_var.get()
async def test_max_count_multiple_batches_cancel(loop): t_exec = 0 delay_exec = 0.1 event = Event() executions = 0 leeway = 0.1 max_count = 5 arg = ContextVar('arg') tasks = [] executing_task = None @aggregate(leeway * 1000, max_count) async def pow(*args: float, power: float = 2) -> List[float]: nonlocal executions, executing_task, t_exec, delay_exec t_exec = monotonic() executions += 1 executing_task = tasks[arg.get()] event.set() await asyncio.sleep(delay_exec) return [math.pow(num, power) for num in args] async def pho(num: int): arg.set(num) return await pow(float(num)) t = monotonic() tasks = [] for i in range(9): tasks.append(loop.create_task(pho(i))) # Execution of the first batch must have started await event.wait() event.clear() assert all(not task.done() for task in tasks) assert executions == 1 first_executing_task = executing_task first_executing_task.cancel() # Another task must have tried to execute await event.wait() event.clear() assert executions == 2 assert first_executing_task.cancelled() assert all( not task.done() for task in tasks if task is not first_executing_task ) await wait(tasks[:5]) # First batch must have finished assert first_executing_task.cancelled() for i, task in enumerate(tasks[:5]): if task is first_executing_task: continue assert task.done() assert task.result() == math.pow(i, 2) # Second batch must have started execution await event.wait() elapsed = t_exec - t assert leeway * 0.9 < elapsed < leeway * 2 assert all(not task.done() for task in tasks[5:]) assert executions == 3 # Second batch mast have finished await wait(tasks[5:]) for i, task in enumerate(tasks[5:], start=5): assert task.done() assert task.result() == math.pow(i, 2)
from .base import state_machine from .base import super_check, call_with_super_check from .base.state_machine import StateEntryFailed, StateMachine, event from .base import TransitionFailed from . import events from . import persistence from . import process_comms from . import process_states from . import utils # pylint: disable=too-many-lines __all__ = ['Process', 'ProcessSpec', 'BundleKeys', 'TransitionFailed'] _LOGGER = logging.getLogger(__name__) PROCESS_STACK = ContextVar('process stack', default=[]) class BundleKeys: """ String keys used by the process to save its state in the state bundle. See :func:`save_instance_state` and :func:`load_instance_state`. """ # pylint: disable=too-few-public-methods INPUTS_RAW = 'INPUTS_RAW' INPUTS_PARSED = 'INPUTS_PARSED' OUTPUTS = 'OUTPUTS' class ProcessStateMachineMeta(abc.ABCMeta, state_machine.StateMachineMeta):
Except: - we do it on top of asyncio - we do not pretend implementing an equivalent of Trio; just the philosophy! - no pretensions; if you want more Trio-like, just switch to Trio! """ import asyncio import logging import sys import time from typing import Awaitable, Optional from aiocontextvars import ContextVar from traio.task import NamedFuture, TaskWrapper SCOPE = ContextVar('traio_scope') DEFAULT_LOGGER = logging.getLogger('traio') DEFAULT_LOGGER.setLevel(logging.CRITICAL) PY37 = sys.version_info >= (3, 7) def current_task(loop=None): """ Return current task. Wraps difference between before/after py37 """ if PY37: return asyncio.current_task(loop) return asyncio.Task.current_task(loop) class CancelError(Exception):
try: return "%s.%s.%s" % ( func.im_class.__module__, # type: ignore func.im_class.__name__, # type: ignore func.__name__, ) except Exception: pass func_qualname = ( getattr(func, "__qualname__", None) or getattr(func, "__name__", None) or None ) # type: Optional[str] if not func_qualname: # No idea what it is return None # Methods in Python 3 # Functions # Classes try: return "%s.%s" % (func.__module__, func_qualname) except Exception: pass # Possibly a lambda return func_qualname disable_capture_event = ContextVar("disable_capture_event")
from aiocontextvars import ContextVar from asyncpg import Connection, create_pool from asyncpg.pool import Pool class ConnectionManager(ABC): @abstractmethod async def get(self, pool: str = "") -> Connection: """Get a database connection.""" @abstractmethod async def put(self, connection: Connection, pool: str = "") -> None: """Put a database connection back in its pool.""" connections_var = ContextVar('connections', default=None) class DefaultConnectionManager(ConnectionManager): def __init__(self, settings: List[Dict[str, Any]], default: str = "") -> None: self.settings = settings self.pools: Dict[str, Pool] = {} self.default = default or self.settings[0]['name'] async def get(self, pool: str = "") -> Connection: connection: Connection = connections_var.get(None) if connection is None: if self.pools == {}: await self._setup()